burn_backend/backend/distributed/
ops.rs

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

burn_backend/backend/distributed/
ops.rs
