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
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129

mod tests {
    use std::sync::mpsc::SyncSender;

    use ruda_core::rand::get_seeded_rng;
    use ruda_tensor::api::{Tensor, TensorData, TensorPrimitive, Tolerance, backend::Backend};

    use serial_test::serial;

    use crate::PeerId;

    #[cfg(not(any(
        feature = "test-cuda",
        feature = "test-wgpu",
        feature = "test-metal",
        feature = "test-vulkan"
    )))]
    pub type TestBackend = ruda_tensor_host::Host;

    #[cfg(feature = "test-cuda")]
    pub type TestBackend = ruda_tensor_device::cuda::Cuda<f32>;

    #[cfg(feature = "test-wgpu")]
    pub type TestBackend = ruda_tensor_wgpu::Wgpu<f32>;

    #[cfg(feature = "test-metal")]
    pub type TestBackend = ruda_tensor_wgpu::Wgpu<f32>;

    #[cfg(feature = "test-vulkan")]
    pub type TestBackend = ruda_tensor_wgpu::Wgpu<f32>;

    use crate::{BroadcastStrategy, CollectiveConfig, broadcast, register, reset_collective};

    pub fn run_peer<B: Backend>(
        id: PeerId,
        config: CollectiveConfig,
        input: Option<TensorData>,
        output: SyncSender<Tensor<B, 1>>,
    ) {
        let device = B::Device::default();

        register::<B>(id, device.clone(), config).unwrap();

        let tensor = input.map(|data| B::float_from_data(data, &device));
        let tensor = broadcast::<B>(id, tensor).unwrap();
        let tensor = Tensor::<B, 1>::from_primitive(TensorPrimitive::Float(tensor));

        output.send(tensor).unwrap();
    }

    fn generate_random_input(shape: Vec<usize>) -> TensorData {
        TensorData::random::<f32, _, _>(
            shape.clone(),
            ruda_tensor::api::Distribution::Default,
            &mut get_seeded_rng(),
        )
    }

    fn test_broadcast<B: Backend>(
        device_count: usize,
        strategy: BroadcastStrategy,
        tensor_size: usize,
    ) {
        reset_collective::<TestBackend>();

        let (send, recv) = std::sync::mpsc::sync_channel(32);

        let shape = vec![tensor_size];

        let input = generate_random_input(shape);

        let config = CollectiveConfig::default()
            .with_num_devices(device_count)
            .with_local_broadcast_strategy(strategy);

        for id in 0..device_count {
            // The peer #0 is the root: it sends the tensor
            let input = if id == 0 { Some(input.clone()) } else { None };

            std::thread::spawn({
                let config = config.clone();
                let send = send.clone();
                move || run_peer::<B>(id.into(), config, input, send)
            });
        }

        // Expect all peers to receive the input tensor
        let tol: Tolerance<f32> = Tolerance::balanced();
        for _ in 0..device_count {
            let tensor = recv.recv().unwrap().to_data();
            input.assert_approx_eq(&tensor, tol);
        }
    }

    #[test]
    #[serial]
    pub fn test_broadcast_centralized_sum() {
        test_broadcast::<TestBackend>(4, BroadcastStrategy::Centralized, 4);
    }

    #[test]
    #[serial]
    pub fn test_broadcast_centralized_mean() {
        test_broadcast::<TestBackend>(4, BroadcastStrategy::Centralized, 4);
    }

    #[test]
    #[serial]
    pub fn test_broadcast_binary_tree_sum() {
        test_broadcast::<TestBackend>(4, BroadcastStrategy::Tree(2), 4);
    }

    #[test]
    #[serial]
    pub fn test_broadcast_binary_tree_mean() {
        test_broadcast::<TestBackend>(4, BroadcastStrategy::Tree(2), 4);
    }

    #[test]
    #[serial]
    pub fn test_broadcast_5_tree_sum() {
        test_broadcast::<TestBackend>(4, BroadcastStrategy::Tree(5), 4);
    }

    #[test]
    #[serial]
    pub fn test_broadcast_5_tree_mean() {
        test_broadcast::<TestBackend>(4, BroadcastStrategy::Tree(5), 4);
    }
}