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

use ruda_tensor::Backend;

use crate::{PeerId, local::tensor_map::CollectiveTensorMap};

#[cfg(feature = "tracing")]
use crate::local::tensor_map::get_common_shape;

/// Sums the tensors on one device and returns the result
#[cfg_attr(feature = "tracing", tracing::instrument(
    level="trace",
    skip(tensors),
    fields(shape = ?get_common_shape::<B>(&tensors).unwrap())
))]
pub(crate) fn reduce_sum_centralized<B: Backend>(
    mut tensors: CollectiveTensorMap<B>,
    central: &PeerId,
) -> B::FloatTensorPrimitive {
    let mut central_tensor = tensors
        .remove(central)
        .expect("Source device id is in the map");
    let central_device = B::float_device(&central_tensor);

    for (_, tensor) in tensors {
        let rhs = B::float_to_device(tensor.clone(), &central_device);
        central_tensor = B::float_add(central_tensor, rhs);
    }

    central_tensor
}