ruCCL 0.21.1

Ruda collective communication algorithms and orchestration.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88

use crate::{PeerId, local::tensor_map::CollectiveTensorMap};
use ruda_tensor::{Backend, DeviceOps};
use std::collections::HashMap;

/// Performs an all-reduce on the provided tensors in a b-tree structure with `arity`.
/// Similar to [reduce_sum_tree](reduce_sum_tree), but this function broadcasts the result with
/// the same tree algorithm.
/// The returned tensors are on the same devices as the corresponding inputs
#[cfg_attr(
    feature = "tracing",
    tracing::instrument(level = "trace", skip(tensors))
)]
pub(crate) fn all_reduce_sum_tree<B: Backend>(
    tensors: CollectiveTensorMap<B>,
    arity: u32,
) -> CollectiveTensorMap<B> {
    let mut input = tensors.into_iter().collect::<Vec<_>>();

    // Sort to put devices of the same type together
    input.sort_by(|a, b| {
        let dev_a = B::float_device(&a.1);
        let dev_b = B::float_device(&b.1);
        dev_a.id().cmp(&dev_b.id())
    });
    // Recursive all-reduce
    let out = all_reduce_sum_tree_inner::<B>(input, arity);

    let mut tensors = HashMap::new();
    for (id, tensor) in out {
        tensors.insert(id, tensor);
    }
    tensors
}

/// Recursive function that sums `tensors` and redistributes the result to the host devices
#[cfg_attr(
    feature = "tracing",
    tracing::instrument(level = "trace", skip(tensors))
)]
fn all_reduce_sum_tree_inner<B: Backend>(
    mut tensors: Vec<(PeerId, B::FloatTensorPrimitive)>,
    arity: u32,
) -> Vec<(PeerId, B::FloatTensorPrimitive)> {
    let mut parent_tensors = vec![];
    let mut children_groups = vec![];

    // Phase 1: Sum tensors in groups of `arity` + 1
    while !tensors.is_empty() {
        // Maps ids to devices for each child of this parent
        let mut children = vec![];
        let (parent, mut parent_tensor) = tensors.remove(0);
        let parent_device = B::float_device(&parent_tensor);

        for _ in 0..arity {
            if tensors.is_empty() {
                break;
            }
            let (child, mut child_tensor) = tensors.remove(0);
            let child_device = B::float_device(&child_tensor);
            children.push((child, child_device));
            child_tensor = B::float_to_device(child_tensor, &parent_device);
            parent_tensor = B::float_add(parent_tensor, child_tensor);
        }

        parent_tensors.push((parent, parent_tensor));
        children_groups.push(children);
    }

    if parent_tensors.len() > 1 {
        // Parents are not yet at the root, do the upper part of the tree
        parent_tensors = all_reduce_sum_tree_inner::<B>(parent_tensors, arity);
    }

    // Phase 2: Redistribute result from each parent to the respective devices
    for (parent, parent_tensor) in parent_tensors {
        let children = children_groups.remove(0);
        for (child, child_device) in children {
            // replace child tensors with result
            tensors.push((
                child,
                B::float_to_device(parent_tensor.clone(), &child_device),
            ));
        }
        tensors.push((parent, parent_tensor));
    }

    tensors
}