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

use ruda_tensor::Backend;

use crate::local::tensor_map::{CollectiveTensorMap, get_peer_devices};
use crate::local::{broadcast_centralized, reduce_sum_centralized};

/// Perform an all-reduce operation by reducing all tensors on one device, and broadcasting the
/// result to all other devices
///
/// Internally, this is just a call to `reduce` followed by a `broadcast`
#[cfg_attr(
    feature = "tracing",
    tracing::instrument(level = "trace", skip(tensors))
)]
pub(crate) fn all_reduce_sum_centralized<B: Backend>(
    tensors: CollectiveTensorMap<B>,
) -> CollectiveTensorMap<B> {
    // Get corresponding devices for each peer
    let peer_devices = get_peer_devices::<B>(&tensors);
    let central_device = *tensors.keys().next().unwrap();

    // Reduce to central device
    let central_tensor = reduce_sum_centralized::<B>(tensors, &central_device);

    // Broadcast result to all
    broadcast_centralized::<B>(peer_devices, central_device, central_tensor)
}