nuvec_core/
distributed.rs

1//! Distributed groups and collective operations, as MLX's `mx.distributed`.
2//!
3//! A [`Group`] is a set of processes with a rank and a size. Collectives (`all_sum`,
4//! `all_max`, `all_min`, `all_gather`, `send`, `recv`) are lazy array operations over a
5//! group: they run when their result is evaluated, compose with `grad`, `vmap` and
6//! `compile` as MLX defines, and run on the CPU over the arrays' shared memory (on a GPU
7//! stream they wait for the GPU work producing their inputs). Without a configured group,
8//! the default group has size 1 and every collective is the identity, so single-process
9//! code runs unchanged.
10//!
11//! The transport is a [`Communicator`]; `nuvec-distributed` provides MLX's TCP ring and
12//! installs it as the default group.
13
14use std::panic::Location;
15use std::sync::{Arc, Mutex, PoisonError, RwLock};
16
17use crate::array::Data;
18use crate::buffer::Buffer;
19use crate::kernels::Operand;
20use crate::primitive::Primitive;
21use crate::{Array, DType, Error, IntoShape, Result};
22
23/// The reduction of [`Communicator::all_reduce`].
24#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
25#[non_exhaustive]
26pub enum AllReduce {
27    /// Element-wise sum (booleans: or).
28    Sum,
29    /// Element-wise maximum (booleans: or).
30    Max,
31    /// Element-wise minimum (booleans: and).
32    Min,
33}
34
35/// The transport of a [`Group`]: every rank of the group calls the same operations in the
36/// same order. Buffers are row-contiguous elements; a failure names the peer.
37pub trait Communicator: Send + Sync {
38    /// This process's rank, `0..size`.
39    fn rank(&self) -> usize;
40    /// The number of processes.
41    fn size(&self) -> usize;
42    /// Reduces `data` (elements of `dtype`) with every rank's `data`, in place; every rank
43    /// gets the same bits.
44    fn all_reduce(&self, op: AllReduce, dtype: DType, data: &mut [u8]) -> Result<()>;
45    /// Every rank's `input` (equal lengths), concatenated in rank order into `output`.
46    fn all_gather(&self, input: &[u8], output: &mut [u8]) -> Result<()>;
47    /// Sends `data` to rank `dst`.
48    fn send(&self, dst: usize, data: &[u8]) -> Result<()>;
49    /// Receives exactly `data.len()` bytes from rank `src`.
50    fn recv(&self, src: usize, data: &mut [u8]) -> Result<()>;
51    /// A subgroup of the ranks with the same `color`, ordered by `key` (MLX's `split`).
52    fn split(&self, color: i32, key: i32) -> Result<Arc<dyn Communicator>> {
53        let _ = (color, key);
54        Err(Error::InvalidArgument(
55            "distributed: this backend does not split groups (as MLX's ring)".into(),
56        ))
57    }
58}
59
60/// A collective operation's kind, with its parameters ([`Primitive::Collective`]).
61#[derive(Debug, Clone, PartialEq, Eq, Hash)]
62#[non_exhaustive]
63pub enum CollectiveOp {
64    /// [`AllReduce`] over the group.
65    Reduce(AllReduce),
66    /// Every rank's array, concatenated along the first axis in rank order.
67    Gather,
68    /// The input sent to rank `dst`; the output is the input.
69    Send {
70        /// The receiving rank.
71        dst: usize,
72    },
73    /// An array of `shape` and `dtype` received from rank `src`.
74    Recv {
75        /// The sending rank.
76        src: usize,
77        /// The received shape.
78        shape: Vec<usize>,
79        /// The received dtype.
80        dtype: DType,
81    },
82}
83
84impl CollectiveOp {
85    /// The operation's name (MLX's).
86    pub fn name(&self) -> &'static str {
87        match self {
88            Self::Reduce(AllReduce::Sum) => "all_sum",
89            Self::Reduce(AllReduce::Max) => "all_max",
90            Self::Reduce(AllReduce::Min) => "all_min",
91            Self::Gather => "all_gather",
92            Self::Send { .. } => "send",
93            Self::Recv { .. } => "recv",
94        }
95    }
96}
97
98/// A group's registered transport and the lock that keeps its operations in program
99/// order (one collective on the wire at a time).
100struct Registered {
101    comm: Arc<dyn Communicator>,
102    order: Mutex<()>,
103}
104
105/// Groups by id (`id - 1`); id 0 is the size-1 group.
106static GROUPS: Mutex<Vec<Arc<Registered>>> = Mutex::new(Vec::new());
107static DEFAULT: RwLock<Option<Group>> = RwLock::new(None);
108
109fn registered(id: usize) -> Result<Arc<Registered>> {
110    let groups = GROUPS.lock().unwrap_or_else(PoisonError::into_inner);
111    id.checked_sub(1)
112        .and_then(|i| groups.get(i).cloned())
113        .ok_or_else(|| Error::InvalidArgument(format!("distributed: no group {id}")))
114}
115
116/// A group of processes (MLX's `Group`): cheap to clone.
117#[derive(Clone)]
118pub struct Group {
119    id: usize,
120    rank: usize,
121    size: usize,
122}
123
124impl std::fmt::Debug for Group {
125    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
126        write!(f, "Group(rank {} of {})", self.rank, self.size)
127    }
128}
129
130impl Group {
131    /// The group of this process alone: every collective is the identity.
132    pub fn single() -> Self {
133        Self {
134            id: 0,
135            rank: 0,
136            size: 1,
137        }
138    }
139
140    /// A group over `comm` (a backend's transport).
141    pub fn new(comm: Arc<dyn Communicator>) -> Self {
142        let (rank, size) = (comm.rank(), comm.size());
143        if size <= 1 {
144            return Self::single();
145        }
146        let mut groups = GROUPS.lock().unwrap_or_else(PoisonError::into_inner);
147        groups.push(Arc::new(Registered {
148            comm,
149            order: Mutex::new(()),
150        }));
151        Self {
152            id: groups.len(),
153            rank,
154            size,
155        }
156    }
157
158    /// This process's rank in the group.
159    pub fn rank(&self) -> usize {
160        self.rank
161    }
162
163    /// The number of processes in the group.
164    pub fn size(&self) -> usize {
165        self.size
166    }
167
168    /// The ranks with the same `color` as a new group, ordered by `key` (then by rank), as
169    /// MLX's `Group.split`. Backends may not support it (MLX's ring does not).
170    pub fn split(&self, color: i32, key: i32) -> Result<Group> {
171        if self.size == 1 {
172            return Ok(self.clone());
173        }
174        Ok(Group::new(registered(self.id)?.comm.split(color, key)?))
175    }
176}
177
178/// Makes `group` the default group of the collectives (what `nuvec-distributed`'s `init`
179/// does).
180pub fn set_default_group(group: Group) {
181    *DEFAULT.write().unwrap_or_else(PoisonError::into_inner) = Some(group);
182}
183
184/// The default group: the one set by [`set_default_group`], else this process alone.
185pub fn default_group() -> Group {
186    DEFAULT
187        .read()
188        .unwrap_or_else(PoisonError::into_inner)
189        .clone()
190        .unwrap_or_else(Group::single)
191}
192
193fn group_or_default(group: Option<&Group>) -> Group {
194    group.cloned().unwrap_or_else(default_group)
195}
196
197/// The size of the group with this id (1 for id 0).
198pub(crate) fn group_size(id: usize) -> Result<usize> {
199    match id {
200        0 => Ok(1),
201        id => Ok(registered(id)?.comm.size()),
202    }
203}
204
205/// The rank of this process in the group with this id.
206pub(crate) fn group_rank(id: usize) -> Result<usize> {
207    match id {
208        0 => Ok(0),
209        id => Ok(registered(id)?.comm.rank()),
210    }
211}
212
213fn record(op: CollectiveOp, group: &Group, inputs: Vec<Array>) -> Result<Array> {
214    let mut out = Array::from_op(
215        Primitive::Collective {
216            op,
217            group: group.id,
218        },
219        inputs,
220    )?;
221    Ok(out.swap_remove(0))
222}
223
224fn reduce(x: &Array, op: AllReduce, group: Option<&Group>) -> Result<Array> {
225    let group = group_or_default(group);
226    if group.size == 1 {
227        return Ok(x.clone());
228    }
229    record(CollectiveOp::Reduce(op), &group, vec![x.clone()])
230}
231
232fn gather(x: &Array, group: Option<&Group>) -> Result<Array> {
233    let group = group_or_default(group);
234    if group.size == 1 {
235        return Ok(x.clone());
236    }
237    record(CollectiveOp::Gather, &group, vec![x.clone()])
238}
239
240fn peer(op: &str, group: &Group, peer: usize) -> Result<()> {
241    if group.size == 1 || peer >= group.size || peer == group.rank {
242        return Err(Error::InvalidArgument(format!(
243            "{op}: rank {peer} is not another rank of a group of {}",
244            group.size
245        )));
246    }
247    Ok(())
248}
249
250fn send_op(x: &Array, dst: usize, group: Option<&Group>) -> Result<Array> {
251    let group = group_or_default(group);
252    peer("send", &group, dst)?;
253    record(CollectiveOp::Send { dst }, &group, vec![x.clone()])
254}
255
256fn recv_op(shape: Vec<usize>, dtype: DType, src: usize, group: Option<&Group>) -> Result<Array> {
257    let group = group_or_default(group);
258    peer("recv", &group, src)?;
259    record(CollectiveOp::Recv { src, shape, dtype }, &group, Vec::new())
260}
261
262/// The header a `send` puts before its data, so a mismatched `recv` is an error rather
263/// than reinterpreted bytes: a tag, the dtype's name length and name, the rank and the
264/// dimensions, as little-endian `u64`s.
265fn header(dtype: DType, shape: &[usize]) -> Vec<u8> {
266    let name = dtype.name().as_bytes();
267    let mut h = Vec::with_capacity(32 + name.len() + 8 * shape.len());
268    h.extend_from_slice(b"nuvecsnd");
269    h.extend_from_slice(&(name.len() as u64).to_le_bytes());
270    h.extend_from_slice(name);
271    h.extend_from_slice(&(shape.len() as u64).to_le_bytes());
272    for &d in shape {
273        h.extend_from_slice(&(d as u64).to_le_bytes());
274    }
275    h
276}
277
278/// Evaluates a collective on the CPU: `operands` are the inputs' layouts and data.
279pub(crate) fn eval_cpu(
280    op: &CollectiveOp,
281    group: usize,
282    operands: &[Operand<'_>],
283) -> Result<Vec<Data>> {
284    let entry = registered(group)?;
285    let _order = entry.order.lock().unwrap_or_else(PoisonError::into_inner);
286    let comm = &entry.comm;
287    let data = match op {
288        CollectiveOp::Reduce(reduce) => {
289            let (meta, data) = operands[0];
290            let mut out = crate::kernels::contiguous(meta, data)?;
291            comm.all_reduce(*reduce, meta.dtype, out.bytes_mut())?;
292            out
293        }
294        CollectiveOp::Gather => {
295            let (meta, data) = operands[0];
296            let input = crate::kernels::contiguous(meta, data)?;
297            let mut out = Buffer::try_alloc(input.bytes().len() * comm.size())?;
298            comm.all_gather(input.bytes(), out.bytes_mut())?;
299            out
300        }
301        CollectiveOp::Send { dst } => {
302            let (meta, data) = operands[0];
303            let input = crate::kernels::contiguous(meta, data)?;
304            let h = header(meta.dtype, &meta.shape);
305            comm.send(*dst, &(h.len() as u64).to_le_bytes())?;
306            comm.send(*dst, &h)?;
307            comm.send(*dst, &(input.bytes().len() as u64).to_le_bytes())?;
308            comm.send(*dst, input.bytes())?;
309            input
310        }
311        CollectiveOp::Recv { src, shape, dtype } => {
312            let want = header(*dtype, shape);
313            let mut len = [0u8; 8];
314            comm.recv(*src, &mut len)?;
315            let len = u64::from_le_bytes(len) as usize;
316            if len > (1 << 20) {
317                return Err(Error::InvalidArgument(format!(
318                    "recv: rank {src} sent no array header"
319                )));
320            }
321            let mut got = vec![0u8; len];
322            comm.recv(*src, &mut got)?;
323            let mut bytes = [0u8; 8];
324            comm.recv(*src, &mut bytes)?;
325            let bytes = u64::from_le_bytes(bytes) as usize;
326            if got != want {
327                // The payload is read and dropped, so the connection stays aligned.
328                let mut sink = vec![0u8; bytes.min(1 << 20)];
329                let mut left = bytes;
330                while left > 0 {
331                    let n = left.min(sink.len());
332                    comm.recv(*src, &mut sink[..n])?;
333                    left -= n;
334                }
335                return Err(Error::InvalidArgument(format!(
336                    "recv: rank {src} sent another dtype or shape than {} {shape:?}",
337                    dtype.name()
338                )));
339            }
340            let n: usize = shape.iter().product();
341            let mut out = Buffer::try_alloc(n * dtype.size())?;
342            comm.recv(*src, out.bytes_mut())?;
343            out
344        }
345    };
346    Ok(vec![Data::new(data)])
347}
348
349/// The element-wise sum of `x` over the ranks of `group` (the default group if `None`), as
350/// MLX's `all_sum`: every rank gets the same bits.
351///
352/// ```
353/// use nuvec_core::{array, distributed};
354///
355/// // Without a configured group, the group is this process alone.
356/// let x = array![1.0, 2.0];
357/// assert_eq!(distributed::all_sum(&x, None).to_vec::<f32>().unwrap(), [1.0, 2.0]);
358/// ```
359#[track_caller]
360pub fn all_sum(x: &Array, group: Option<&Group>) -> Array {
361    Array::sticky("all_sum", Location::caller(), &[x], || {
362        reduce(x, AllReduce::Sum, group)
363    })
364}
365
366/// The element-wise maximum of `x` over the ranks of `group` (MLX's `all_max`).
367#[track_caller]
368pub fn all_max(x: &Array, group: Option<&Group>) -> Array {
369    Array::sticky("all_max", Location::caller(), &[x], || {
370        reduce(x, AllReduce::Max, group)
371    })
372}
373
374/// The element-wise minimum of `x` over the ranks of `group` (MLX's `all_min`).
375#[track_caller]
376pub fn all_min(x: &Array, group: Option<&Group>) -> Array {
377    Array::sticky("all_min", Location::caller(), &[x], || {
378        reduce(x, AllReduce::Min, group)
379    })
380}
381
382/// Every rank's `x` concatenated along the first axis in rank order (a scalar gathers to a
383/// vector), as MLX's `all_gather`.
384#[track_caller]
385pub fn all_gather(x: &Array, group: Option<&Group>) -> Array {
386    Array::sticky("all_gather", Location::caller(), &[x], || gather(x, group))
387}
388
389/// Sends `x` to rank `dst` of `group` when the result (which is `x`) is evaluated, as
390/// MLX's `send`. MLX's ring backend sends to its neighbours only.
391#[track_caller]
392pub fn send(x: &Array, dst: usize, group: Option<&Group>) -> Array {
393    Array::sticky("send", Location::caller(), &[x], || send_op(x, dst, group))
394}
395
396/// An array of `shape` and `dtype` received from rank `src` of `group`, as MLX's `recv`. A
397/// sender's other dtype or shape is an error when it is evaluated.
398#[track_caller]
399pub fn recv(shape: impl IntoShape, dtype: DType, src: usize, group: Option<&Group>) -> Array {
400    let spec = shape.shape_spec();
401    Array::constructor("recv", Location::caller(), dtype, &spec, |shape| {
402        recv_op(shape.to_vec(), dtype, src, group)
403    })
404}
405
406/// [`recv`] of `x`'s shape and dtype (MLX's `recv_like`).
407#[track_caller]
408pub fn recv_like(x: &Array, src: usize, group: Option<&Group>) -> Array {
409    let at = Location::caller();
410    Array::sticky("recv_like", at, &[x], || {
411        recv_op(x.dims().to_vec(), x.dtype(), src, group)
412    })
413}
414
415/// `x` unchanged, whose gradient is summed over `group` (MLX's `sum_gradients`): the input
416/// of a layer whose output is sharded across the group, so each rank's gradient covers
417/// every rank's slice of the output.
418#[track_caller]
419pub fn sum_gradients(x: &Array, group: Option<&Group>) -> Array {
420    let group = group_or_default(group);
421    if group.size == 1 {
422        return x.clone();
423    }
424    Array::sticky("sum_gradients", Location::caller(), &[x], || {
425        let f = crate::custom_function_flat(|xs| Ok(xs.to_vec())).with_vjp(
426            move |_primals, _outputs, cotangents| {
427                cotangents
428                    .iter()
429                    .map(|c| reduce(c, AllReduce::Sum, Some(&group)))
430                    .collect()
431            },
432        );
433        Ok(f.call(std::slice::from_ref(x))?.remove(0))
434    })
435}
436
437impl crate::tree::Static for Group {
438    fn hash_static(&self, mut h: &mut dyn std::hash::Hasher) {
439        std::hash::Hash::hash(&(self.id, self.rank, self.size), &mut h);
440    }
441}
nuvec_core/distributed.rs

nuvec_core/
distributed.rs
