burn_backend/backend/distributed/
ops.rs

1use alloc::vec::Vec;
2
3use crate::{
4    Backend, DeviceId, TensorMetadata,
5    distributed::CollectiveTensor,
6    tensor::{Device, FloatTensor},
7};
8
9use crate::distributed::{DistributedConfig, DistributedParams, ReduceOperation};
10
11#[cfg(feature = "std")]
12use crate::distributed::{
13    close_distributed_sync_server, get_distributed_sync_client, start_distributed_sync_server,
14};
15
16/// Mutable reference to a float tensor.
17#[derive(Clone)]
18pub struct TensorRef<B: Backend>(pub *mut FloatTensor<B>);
19unsafe impl<B> Sync for TensorRef<B> where B: Backend {}
20unsafe impl<B> Send for TensorRef<B> where B: Backend {}
21
22// TODO : The following functions should be moved in the `burn-autodiff` crate. The difficulty is in not discriminating between
23// the dispatch backend and its inner dispatched backend when calling the communication server API. This implementation makes
24// it easy by implementing `DistributedOps` for `DispatchBackend`.
25// The functions in question :
26// * `start_communication_server`
27// * `close_communication_server`
28// * `register_sync_parameters`
29// * `submit_sync_collective`
30// * `submit_gradient_sync`
31
32/// Operations on communication tensors.
33pub trait DistributedOps<B: Backend> {
34    /// Start the communication server used to orchestrate tensor syncing between devices.
35    ///
36    /// # Arguments
37    ///
38    /// * `devices` - The devices to orchestrate.
39    fn start_communication_server(devices: &[B::Device], config: DistributedConfig) {
40        #[cfg(feature = "std")]
41        start_distributed_sync_server::<B>(devices, config);
42        #[cfg(not(feature = "std"))]
43        let _ = (devices, config);
44    }
45
46    /// Close the communication server used to orchestrate syncing between devices.
47    ///
48    /// # Arguments
49    ///
50    /// * `device` - A device on the backend.
51    fn close_communication_server(_device: &B::Device) {
52        #[cfg(feature = "std")]
53        close_distributed_sync_server::<B>();
54    }
55
56    /// Register the parameters that will require gradient synchronization for the upcoming backward pass.
57    ///
58    /// This must be called before the backward pass on each device so the gradient sync server
59    /// can coordinate collective operations across all devices.
60    ///
61    /// # Arguments
62    ///
63    /// * `device` - The device calling the initialization.
64    /// * `distributed_params` - A list of [`DistributedParams`] of the tensors to sync.
65    fn register_sync_parameters(_device: &B::Device, distributed_params: Vec<DistributedParams>) {
66        #[cfg(feature = "std")]
67        if let Some(sync_client) = get_distributed_sync_client::<B>() {
68            sync_client.register_sync_parameters(distributed_params);
69        };
70        #[cfg(not(feature = "std"))]
71        let _ = distributed_params;
72    }
73
74    /// Tell the gradient sync server that this device has submitted all its sync operations and is ready to be synchronized.
75    ///
76    /// # Arguments
77    ///
78    /// * `device` - The device on which to sync.
79    fn submit_sync_collective(device: &B::Device) {
80        #[cfg(feature = "std")]
81        if let Some(sync_client) = get_distributed_sync_client::<B>() {
82            sync_client.submit_sync_collective(device.clone());
83        };
84        #[cfg(not(feature = "std"))]
85        let _ = device;
86    }
87
88    /// Submit a gradient tensor for synchronization across all devices.
89    ///
90    /// The gradient is sent to the gradient sync server, which will launch the all-reduce
91    /// operation once all devices have submitted their gradient for this parameter.
92    ///
93    /// # Arguments
94    ///
95    /// * `tensor` - The tensor to synchronize.
96    /// * `distributed_params` - The [`DistributedParams`] for the parameter.
97    fn submit_gradient_sync(tensor: TensorRef<B>, distributed_params: DistributedParams) {
98        #[cfg(feature = "std")]
99        if let Some(sync_client) = get_distributed_sync_client::<B>() {
100            sync_client.submit_gradient_sync(tensor, distributed_params);
101        };
102        #[cfg(not(feature = "std"))]
103        let _ = (tensor, distributed_params);
104    }
105
106    /// all_reduce operation.
107    ///
108    /// # Arguments
109    ///
110    /// * `tensors` - The tensors on which to perform all_reduce.
111    /// * `op` - The [`ReduceOperation`].
112    ///
113    /// # Returns
114    ///
115    /// The corresponding [CollectiveTensor].
116    fn all_reduce(
117        _tensor: FloatTensor<B>,
118        _op: ReduceOperation,
119        _device_ids: Vec<DeviceId>,
120    ) -> CollectiveTensor<B> {
121        unimplemented!()
122    }
123
124    /// Sync the collective operations.
125    ///
126    /// # Arguments
127    ///
128    /// * `device` - The device to sync.
129    fn sync_collective(_device: &B::Device) {
130        unimplemented!()
131    }
132
133    /// Get the device of the tensor reference.
134    ///
135    /// # Arguments
136    ///
137    /// * `tensor` - The tensor reference.
138    ///
139    /// # Returns
140    ///
141    /// The device of the tensor.
142    ///
143    /// # Safety
144    ///
145    /// Ensure that the tensors are not accessed/modified when calling.
146    unsafe fn comm_device(tensor: &TensorRef<B>) -> Device<B> {
147        unsafe { &(*tensor.0) }.device()
148    }
149
150    /// Returns a clone of the float tensor from the tensor reference.
151    ///
152    /// # Arguments
153    ///
154    /// * `tensor` - The tensor.
155    ///
156    /// # Returns
157    ///
158    /// A float tensor containing a copy of the data of the given tensor.
159    ///
160    /// # Safety
161    ///
162    /// Ensure that the tensors are not accessed/modified when calling.
163    unsafe fn float_from_ref(tensor: &TensorRef<B>) -> FloatTensor<B> {
164        unsafe { (*tensor.0).clone() }
165    }
166}
burn_backend/backend/distributed/ops.rs

burn_backend/backend/distributed/
ops.rs
