collective config refactoring to add broadcast and reduce

Author: Cielbird

crates/burn-collective/multinode-tests/src/bin/node.rs

@@ -74,7 +74,7 @@ fn launch_threads<B: Backend, const D: usize>(
 
         // Put all the parameters in the config
         let config = CollectiveConfig::default()
-            .with_all_reduce_kind(test_input.all_reduce_kind)
+            .with_all_reduce_kind(test_input.all_reduce_op)
             .with_num_devices(test_input.device_count)
             .with_device_id(id.into())
             .with_node_id(test_input.node_id)

crates/burn-collective/multinode-tests/src/bin/test_launcher.rs

@@ -13,7 +13,7 @@ use burn::{
     prelude::Backend,
     tensor::{Shape, Tensor, TensorData},
 };
-use burn_collective::{AllReduceStrategy, ReduceKind};
+use burn_collective::{AllReduceStrategy, ReduceOperation};
 use burn_collective_multinode_tests::shared::NodeTestData;
 use burn_common::rand::{SeedableRng, StdRng};
 
@@ -39,7 +39,7 @@ async fn main() {
     let nodes = launch_nodes(
         topology,
         tensor_shape,
-        ReduceKind::Mean,
+        ReduceOperation::Mean,
         AllReduceStrategy::Ring,
         AllReduceStrategy::Tree(2),
     );
@@ -131,13 +131,13 @@ fn launch_orchestrator(test_files_dir: &str) -> Child {
 fn launch_nodes(
     topology: Vec<u32>,
     tensor_shape: Shape,
-    reduce_kind: ReduceKind,
+    reduce_op: ReduceOperation,
     global_strategy: AllReduceStrategy,
     local_strategy: AllReduceStrategy,
 ) -> Vec<(String, Child)> {
     let total_device_count = topology.iter().sum();
     let (inputs, expected) =
-        generate_random_input(tensor_shape, reduce_kind, total_device_count, 42);
+        generate_random_input(tensor_shape, reduce_op, total_device_count, 42);
 
     // URL for the global orchestrator on port 3000
     let global_url = "ws://localhost:3000";
@@ -156,7 +156,7 @@ fn launch_nodes(
             expected.clone(),
             node_count,
             global_address.clone(),
-            reduce_kind,
+            reduce_op,
             global_strategy,
             local_strategy,
         );
@@ -190,7 +190,7 @@ fn write_node_input(
     expected: TensorData,
     node_count: u32,
     global_address: Address,
-    reduce_kind: ReduceKind,
+    reduce_op: ReduceOperation,
     global_strategy: AllReduceStrategy,
     local_strategy: AllReduceStrategy,
 ) -> String {
@@ -213,7 +213,7 @@ fn write_node_input(
         global_address,
         node_address,
         data_service_port,
-        all_reduce_kind: reduce_kind,
+        all_reduce_op: reduce_op,
         global_strategy,
         local_strategy,
         inputs,
@@ -231,7 +231,7 @@ fn write_node_input(
 /// Generates random input tensors and expected output based on the provided shape and reduce kind.
 fn generate_random_input(
     shape: Shape,
-    reduce_kind: ReduceKind,
+    reduce_kind: ReduceOperation,
     input_count: u32,
     seed: u64,
 ) -> (Vec<TensorData>, TensorData) {
@@ -256,7 +256,7 @@ fn generate_random_input(
         expected_tensor = expected_tensor.add(input_tensor);
     }
 
-    if reduce_kind == ReduceKind::Mean {
+    if reduce_kind == ReduceOperation::Mean {
         expected_tensor = expected_tensor.div_scalar(input_count);
     }
 

crates/burn-collective/multinode-tests/src/shared.rs

@@ -1,5 +1,5 @@
 use burn::tensor::TensorData;
-use burn_collective::{AllReduceStrategy, NodeId, ReduceKind};
+use burn_collective::{AllReduceStrategy, NodeId, ReduceOperation};
 use burn_communication::Address;
 use serde::{Deserialize, Serialize};
 
@@ -17,8 +17,8 @@ pub struct NodeTestData {
     pub node_address: Address,
     /// Node's data service port, for initializing the p2p tensor data service
     pub data_service_port: u16,
-    /// What kind of aggregation
-    pub all_reduce_kind: ReduceKind,
+    /// What kind of all-reduce
+    pub all_reduce_op: ReduceOperation,
     /// Node's data service port, for initializing the p2p tensor data service
     pub global_strategy: AllReduceStrategy,
     /// What kind of aggregation

crates/burn-collective/src/api.rs

@@ -1,7 +1,7 @@
-use burn_tensor::{Tensor, backend::Backend};
+use burn_tensor::{backend::Backend, Tensor};
 
 use crate::{
-    CollectiveConfig, global::shared::GlobalCollectiveError, local_server::get_collective_client,
+    global::shared::GlobalCollectiveError, local_server::get_collective_client, CollectiveConfig, DeviceId, ReduceOperation
 };
 
 /// Errors from collective operations
@@ -33,32 +33,94 @@ pub enum CollectiveError {
 /// Registers a device. `num_devices` must be the same for every register,
 /// and `device_id` must be unique.
 ///
+/// * `id` - The peer id of the caller
+/// 
 /// With auto-diff backends, make sure to use the inner backend.
-pub fn register<B: Backend>(config: &CollectiveConfig) -> Result<(), CollectiveError> {
+pub fn register<B: Backend>(
+    id: DeviceId,
+    config: CollectiveConfig,
+) -> Result<(), CollectiveError> {
     let mut client = get_collective_client::<B>();
-    client.register(config)
+    client.register(id, config)
 }
 
 /// Calls for an all-reduce operation with the given parameters, and returns the result.
 /// The `params` must be the same as the parameters passed by the other nodes.
+///
+/// * `id` - The peer id of the caller
+/// * `tensor` - The input tensor to reduce with the peers' tensors
+/// * `config` - Config of the collective operation, must be coherent with the other calls
 pub fn all_reduce<B: Backend, const D: usize>(
+    id: DeviceId,
     tensor: Tensor<B, D>,
-    config: &CollectiveConfig,
+    op: ReduceOperation,
 ) -> Result<Tensor<B, D>, CollectiveError> {
     let client = get_collective_client::<B>();
     let device = tensor.device();
     let tensor = tensor.into_primitive().tensor();
-    let primitive = client.all_reduce(tensor, config)?;
+    let primitive = client.all_reduce(id, tensor, op)?;
     let tensor =
         Tensor::from_primitive(burn_tensor::TensorPrimitive::Float(primitive)).to_device(&device);
 
     Ok(tensor)
 }
 
+/// Broadcasts, or recives a broadcasted tensor.
+///
+/// * `id` - The peer id of the caller
+/// * `tensor` - If defined, this tensor will be broadcasted. Otherwise, this call will receive
+///     the broadcasted tensor.
+/// * `root` - The peer that will broadcast the tensor.
+/// * `config` - Config of the collective operation, must be coherent with the other calls
+///
+/// Returns the broadcasted tensor.
+pub fn broadcast<B: Backend, const D: usize>(
+    id: DeviceId,
+    tensor: Option<Tensor<B, D>>,
+    _device: B::Device, // TODO `register` should return a client, and collective ops should be done on the client.
+    root: DeviceId,
+) -> Result<Tensor<B, D>, CollectiveError> {
+    let client = get_collective_client::<B>();
+    let tensor = tensor.map(|tensor| {
+        tensor.device();
+        tensor.into_primitive().tensor()
+    });
+    let primitive = client.broadcast(id, tensor, root)?;
+    let tensor =
+        Tensor::from_primitive(burn_tensor::TensorPrimitive::Float(primitive));
+
+    Ok(tensor)
+}
+
+/// Reduces a tensor onto one device.
+///
+/// * `id` - The peer id of the caller
+/// * `tensor` - The tensor to send as input
+/// * `root` - The ID of the peer that will receive the result.
+/// * `config` - Config of the collective operation, must be coherent with the other calls
+///
+/// Returns Ok(None) if the root tensor is not the caller. Otherwise, returns the reduced tensor.
+pub fn reduce<B: Backend, const D: usize>(
+    id: DeviceId,
+    tensor: Tensor<B, D>,
+    op: ReduceOperation,
+    root: DeviceId,
+) -> Result<Option<Tensor<B, D>>, CollectiveError> {
+    let client = get_collective_client::<B>();
+    let device = tensor.device();
+    let tensor = tensor.into_primitive().tensor();
+    let primitive = client.reduce(id, tensor, op, root)?;
+    let tensor = primitive.map(|primitive| {
+        Tensor::from_primitive(burn_tensor::TensorPrimitive::Float(primitive)).to_device(&device)
+    });
+
+    Ok(tensor)
+}
+
 /// Closes the collective session, unregistering the device
-pub fn finish_collective<B: Backend>(config: &CollectiveConfig) -> Result<(), CollectiveError> {
+pub fn finish_collective<B: Backend>(id: DeviceId) -> Result<(), CollectiveError> {
     let client = get_collective_client::<B>();
-    client.finish(config.device_id)
+    client.finish(id)
 }
 
 /// Resets the local collective server. All registered callers and ongoing operations are forgotten

crates/burn-collective/src/client.rs

@@ -0,0 +1,117 @@
+use std::sync::mpsc::SyncSender;
+
+use burn_tensor::backend::Backend;
+
+use crate::{
+    CollectiveConfig, CollectiveError, DeviceId, ReduceOperation,
+    local_server::{
+        AllReduceResult, BroadcastResult, FinishResult, Message, ReduceResult, RegisterResult,
+    },
+};
+
+#[derive(Clone)]
+pub(crate) struct LocalCollectiveClient<B: Backend> {
+    pub channel: SyncSender<Message<B>>,
+}
+
+impl<B: Backend> LocalCollectiveClient<B> {
+    pub(crate) fn reset(&self) {
+        self.channel.send(Message::Reset).unwrap();
+    }
+
+    pub(crate) fn register(&mut self, id: DeviceId, config: CollectiveConfig) -> RegisterResult {
+        if config.is_valid() {
+            return Err(CollectiveError::InvalidConfig);
+        }
+
+        let (callback, rec) = std::sync::mpsc::sync_channel::<RegisterResult>(1);
+
+        self.channel
+            .send(Message::Register {
+                device_id: id,
+                config,
+                callback,
+            })
+            .unwrap();
+
+        rec.recv()
+            .unwrap_or(Err(CollectiveError::LocalServerMissing))
+    }
+
+    pub(crate) fn all_reduce(
+        &self,
+        id: DeviceId,
+        tensor: B::FloatTensorPrimitive,
+        op: ReduceOperation,
+    ) -> AllReduceResult<B::FloatTensorPrimitive> {
+        let (callback, rec) =
+            std::sync::mpsc::sync_channel::<AllReduceResult<B::FloatTensorPrimitive>>(1);
+        let msg = Message::AllReduce {
+            device_id: id,
+            tensor,
+            op,
+            callback,
+        };
+
+        self.channel.send(msg).unwrap();
+
+        // returns a tensor primitive that may or may not be on the correct device,
+        // depending on the strategy used.
+        rec.recv()
+            .unwrap_or(Err(CollectiveError::LocalServerMissing))
+    }
+
+    pub(crate) fn reduce(
+        &self,
+        id: DeviceId,
+        tensor: B::FloatTensorPrimitive,
+        op: ReduceOperation,
+        root: DeviceId,
+    ) -> ReduceResult<B::FloatTensorPrimitive> {
+        let (callback, rec) =
+            std::sync::mpsc::sync_channel::<ReduceResult<B::FloatTensorPrimitive>>(1);
+        let msg = Message::Reduce {
+            device_id: id,
+            tensor,
+            op,
+            root,
+            callback,
+        };
+
+        self.channel.send(msg).unwrap();
+
+        // returns a tensor or none depending on if this device is the root
+        rec.recv()
+            .unwrap_or(Err(CollectiveError::LocalServerMissing))
+    }
+
+    pub(crate) fn broadcast(
+        &self,
+        id: DeviceId,
+        tensor: Option<B::FloatTensorPrimitive>,
+        root: DeviceId,
+    ) -> BroadcastResult<B::FloatTensorPrimitive> {
+        let (callback, rec) =
+            std::sync::mpsc::sync_channel::<BroadcastResult<B::FloatTensorPrimitive>>(1);
+        let msg = Message::Broadcast {
+            device_id: id,
+            tensor,
+            root,
+            callback,
+        };
+
+        self.channel.send(msg).unwrap();
+
+        // returns a tensor or none depending on if this device is the root
+        rec.recv()
+            .unwrap_or(Err(CollectiveError::LocalServerMissing))
+    }
+
+    pub(crate) fn finish(&self, id: DeviceId) -> FinishResult {
+        let (callback, rec) = std::sync::mpsc::sync_channel::<FinishResult>(1);
+        self.channel.send(Message::Finish { id, callback }).unwrap();
+
+        rec.recv()
+            .unwrap_or(Err(CollectiveError::LocalServerMissing))
+    }
+}