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

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

/// Performs a reduce on the provided tensors in a b-tree structure with `arity`.
#[cfg_attr(
    feature = "tracing",
    tracing::instrument(level = "trace", skip(tensors))
)]
pub(crate) fn reduce_sum_tree<B: Backend>(
    mut tensors: CollectiveTensorMap<B>,
    root: &PeerId,
    arity: u32,
) -> B::FloatTensorPrimitive {
    // Convert hash map to vector of key-value pairs because order matters
    let mut input = vec![];
    let root_tensor = tensors.remove(root).unwrap();
    for (_, tensor) in tensors.drain() {
        input.push(tensor);
    }

    // Sort to put devices of the same type together
    input.sort_by(|a, b| {
        let dev_a = B::float_device(a);
        let dev_b = B::float_device(b);
        dev_a.id().cmp(&dev_b.id())
    });

    // put the root first
    input.insert(0, root_tensor);

    reduce_sum_tree_inner::<B>(input, arity)
}

/// Recursive function that sums `tensors`
///
/// Traverses `tensors` and reduces in a post-order traversal. The first tensor in the list is
/// chosen as the root
#[cfg_attr(
    feature = "tracing",
    tracing::instrument(level = "trace", skip(tensors))
)]
fn reduce_sum_tree_inner<B: Backend>(
    mut tensors: Vec<B::FloatTensorPrimitive>,
    arity: u32,
) -> B::FloatTensorPrimitive {
    let mut parents = vec![];
    let mut children_groups = vec![];

    // Sum tensors in groups of `arity` + 1
    while !tensors.is_empty() {
        let mut children = vec![];
        let 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_tensor = tensors.remove(0);
            children.push(B::float_device(&child_tensor));
            let rhs = B::float_to_device(child_tensor, &parent_device);
            parent_tensor = B::float_add(parent_tensor, rhs);
        }

        parents.push(parent_tensor);
        children_groups.push(children);
    }

    if parents.len() > 1 {
        // Parents are not yet at the root, do the upper part of the tree
        reduce_sum_tree_inner::<B>(parents, arity)
    } else {
        // Root of tree
        parents.remove(0)
    }
}