burn_backend/backend/ops/
collective.rs

1use crate::{
2    Backend, TensorMetadata,
3    distributed::ReduceOperation,
4    tensor::{BoolTensor, Device, FloatTensor, IntTensor},
5};
6use alloc::vec::Vec;
7use burn_std::{Bytes, DType, Slice, TensorData};
8use core::{marker::PhantomData, ops::Range};
9
10/// Collectives over a group of devices, called once with every member's tensor.
11///
12/// Member `i` is position `i` of the tensors given, each on its own member's device and every one
13/// of one shape. A chunk is a member's share of a dimension, as even as possible with the first
14/// members holding one more; a dimension split into chunks holds at least one element per member,
15/// since a backend never receives an empty slice. The defaults emulate each collective with
16/// transfers, slices, concatenations and additions, so every backend runs them; a backend with a
17/// native collective between its devices overrides them.
18pub trait CollectiveOps<B: Backend> {
19    /// Every member receives `op` over every member's tensor.
20    fn float_all_reduce(tensors: Vec<FloatTensor<B>>, op: ReduceOperation) -> Vec<FloatTensor<B>> {
21        EmulatedCollectives::<B>::float_all_reduce(tensors, op)
22    }
23
24    /// Every member receives every member's tensor, concatenated along `dim` in member order.
25    fn float_all_gather(tensors: Vec<FloatTensor<B>>, dim: usize) -> Vec<FloatTensor<B>> {
26        EmulatedCollectives::<B>::float_all_gather(tensors, dim)
27    }
28
29    /// Member `i` receives chunk `i` along `dim` of `op` over every member's tensor.
30    fn float_reduce_scatter(
31        tensors: Vec<FloatTensor<B>>,
32        dim: usize,
33        op: ReduceOperation,
34    ) -> Vec<FloatTensor<B>> {
35        EmulatedCollectives::<B>::float_reduce_scatter(tensors, dim, op)
36    }
37
38    /// Member `i` receives chunk `i` along `split_dim` of every member's tensor, concatenated along
39    /// `concat_dim` in member order.
40    fn float_all_to_all(
41        tensors: Vec<FloatTensor<B>>,
42        split_dim: usize,
43        concat_dim: usize,
44    ) -> Vec<FloatTensor<B>> {
45        EmulatedCollectives::<B>::float_all_to_all(tensors, split_dim, concat_dim)
46    }
47
48    /// Each of `devices` receives `tensor`.
49    fn float_broadcast(tensor: FloatTensor<B>, devices: &[Device<B>]) -> Vec<FloatTensor<B>> {
50        EmulatedCollectives::<B>::float_broadcast(tensor, devices)
51    }
52
53    /// Every member receives `op` over every member's tensor.
54    fn int_all_reduce(tensors: Vec<IntTensor<B>>, op: ReduceOperation) -> Vec<IntTensor<B>> {
55        EmulatedCollectives::<B>::int_all_reduce(tensors, op)
56    }
57
58    /// Every member receives every member's tensor, concatenated along `dim` in member order.
59    fn int_all_gather(tensors: Vec<IntTensor<B>>, dim: usize) -> Vec<IntTensor<B>> {
60        EmulatedCollectives::<B>::int_all_gather(tensors, dim)
61    }
62
63    /// Member `i` receives chunk `i` along `dim` of `op` over every member's tensor.
64    fn int_reduce_scatter(
65        tensors: Vec<IntTensor<B>>,
66        dim: usize,
67        op: ReduceOperation,
68    ) -> Vec<IntTensor<B>> {
69        EmulatedCollectives::<B>::int_reduce_scatter(tensors, dim, op)
70    }
71
72    /// Member `i` receives chunk `i` along `split_dim` of every member's tensor, concatenated along
73    /// `concat_dim` in member order.
74    fn int_all_to_all(
75        tensors: Vec<IntTensor<B>>,
76        split_dim: usize,
77        concat_dim: usize,
78    ) -> Vec<IntTensor<B>> {
79        EmulatedCollectives::<B>::int_all_to_all(tensors, split_dim, concat_dim)
80    }
81
82    /// Each of `devices` receives `tensor`.
83    fn int_broadcast(tensor: IntTensor<B>, devices: &[Device<B>]) -> Vec<IntTensor<B>> {
84        EmulatedCollectives::<B>::int_broadcast(tensor, devices)
85    }
86
87    /// Every member receives every member's tensor, concatenated along `dim` in member order.
88    fn bool_all_gather(tensors: Vec<BoolTensor<B>>, dim: usize) -> Vec<BoolTensor<B>> {
89        EmulatedCollectives::<B>::bool_all_gather(tensors, dim)
90    }
91
92    /// Member `i` receives chunk `i` along `split_dim` of every member's tensor, concatenated along
93    /// `concat_dim` in member order.
94    fn bool_all_to_all(
95        tensors: Vec<BoolTensor<B>>,
96        split_dim: usize,
97        concat_dim: usize,
98    ) -> Vec<BoolTensor<B>> {
99        EmulatedCollectives::<B>::bool_all_to_all(tensors, split_dim, concat_dim)
100    }
101
102    /// Each of `devices` receives `tensor`.
103    fn bool_broadcast(tensor: BoolTensor<B>, devices: &[Device<B>]) -> Vec<BoolTensor<B>> {
104        EmulatedCollectives::<B>::bool_broadcast(tensor, devices)
105    }
106}
107
108/// The collectives of every tensor kind emulated with ordinary tensor ops, transfers between
109/// devices, slices, concatenations and additions. Its methods are [`CollectiveOps`]' defaults, by a
110/// name an override can call when it falls back.
111pub struct EmulatedCollectives<B> {
112    _backend: PhantomData<B>,
113}
114
115impl<B: Backend> EmulatedCollectives<B> {
116    /// See [`CollectiveOps::float_all_reduce`].
117    pub fn float_all_reduce(
118        tensors: Vec<FloatTensor<B>>,
119        op: ReduceOperation,
120    ) -> Vec<FloatTensor<B>> {
121        Self::all_reduce(tensors, B::float_to_device, |parts| {
122            Self::float_reduced(parts, op)
123        })
124    }
125
126    /// See [`CollectiveOps::float_all_gather`].
127    pub fn float_all_gather(tensors: Vec<FloatTensor<B>>, dim: usize) -> Vec<FloatTensor<B>> {
128        Self::all_gather(tensors, dim, B::float_to_device, B::float_cat)
129    }
130
131    /// See [`CollectiveOps::float_reduce_scatter`].
132    pub fn float_reduce_scatter(
133        tensors: Vec<FloatTensor<B>>,
134        dim: usize,
135        op: ReduceOperation,
136    ) -> Vec<FloatTensor<B>> {
137        Self::reduce_scatter(tensors, dim, B::float_to_device, B::float_slice, |parts| {
138            Self::float_reduced(parts, op)
139        })
140    }
141
142    /// See [`CollectiveOps::float_all_to_all`].
143    pub fn float_all_to_all(
144        tensors: Vec<FloatTensor<B>>,
145        split_dim: usize,
146        concat_dim: usize,
147    ) -> Vec<FloatTensor<B>> {
148        Self::all_to_all(
149            tensors,
150            split_dim,
151            concat_dim,
152            B::float_to_device,
153            B::float_slice,
154            B::float_cat,
155        )
156    }
157
158    /// See [`CollectiveOps::float_broadcast`].
159    pub fn float_broadcast(tensor: FloatTensor<B>, devices: &[Device<B>]) -> Vec<FloatTensor<B>> {
160        Self::broadcast(tensor, devices, B::float_to_device)
161    }
162
163    /// See [`CollectiveOps::int_all_reduce`].
164    pub fn int_all_reduce(tensors: Vec<IntTensor<B>>, op: ReduceOperation) -> Vec<IntTensor<B>> {
165        Self::all_reduce(tensors, B::int_to_device, |parts| {
166            Self::int_reduced(parts, op)
167        })
168    }
169
170    /// See [`CollectiveOps::int_all_gather`].
171    pub fn int_all_gather(tensors: Vec<IntTensor<B>>, dim: usize) -> Vec<IntTensor<B>> {
172        Self::all_gather(tensors, dim, B::int_to_device, B::int_cat)
173    }
174
175    /// See [`CollectiveOps::int_reduce_scatter`].
176    pub fn int_reduce_scatter(
177        tensors: Vec<IntTensor<B>>,
178        dim: usize,
179        op: ReduceOperation,
180    ) -> Vec<IntTensor<B>> {
181        Self::reduce_scatter(tensors, dim, B::int_to_device, B::int_slice, |parts| {
182            Self::int_reduced(parts, op)
183        })
184    }
185
186    /// See [`CollectiveOps::int_all_to_all`].
187    pub fn int_all_to_all(
188        tensors: Vec<IntTensor<B>>,
189        split_dim: usize,
190        concat_dim: usize,
191    ) -> Vec<IntTensor<B>> {
192        Self::all_to_all(
193            tensors,
194            split_dim,
195            concat_dim,
196            B::int_to_device,
197            B::int_slice,
198            B::int_cat,
199        )
200    }
201
202    /// See [`CollectiveOps::int_broadcast`].
203    pub fn int_broadcast(tensor: IntTensor<B>, devices: &[Device<B>]) -> Vec<IntTensor<B>> {
204        Self::broadcast(tensor, devices, B::int_to_device)
205    }
206
207    /// See [`CollectiveOps::bool_all_gather`].
208    pub fn bool_all_gather(tensors: Vec<BoolTensor<B>>, dim: usize) -> Vec<BoolTensor<B>> {
209        Self::all_gather(tensors, dim, B::bool_to_device, B::bool_cat)
210    }
211
212    /// See [`CollectiveOps::bool_all_to_all`].
213    pub fn bool_all_to_all(
214        tensors: Vec<BoolTensor<B>>,
215        split_dim: usize,
216        concat_dim: usize,
217    ) -> Vec<BoolTensor<B>> {
218        Self::all_to_all(
219            tensors,
220            split_dim,
221            concat_dim,
222            B::bool_to_device,
223            B::bool_slice,
224            B::bool_cat,
225        )
226    }
227
228    /// See [`CollectiveOps::bool_broadcast`].
229    pub fn bool_broadcast(tensor: BoolTensor<B>, devices: &[Device<B>]) -> Vec<BoolTensor<B>> {
230        Self::broadcast(tensor, devices, B::bool_to_device)
231    }
232
233    /// Reduces once on the first member and broadcasts the result, so each tensor crosses a link
234    /// once each way rather than to every member.
235    fn all_reduce<T, D, R>(tensors: Vec<T>, to_device: D, reduce: R) -> Vec<T>
236    where
237        T: TensorMetadata<Device = Device<B>>,
238        D: Fn(T, &Device<B>) -> T,
239        R: Fn(Vec<T>) -> T,
240    {
241        let devices = Self::devices(&tensors);
242        let Some(root) = devices.first() else {
243            return Vec::new();
244        };
245        let reduced = reduce(Self::moved_to(&tensors, root, &to_device));
246        Self::broadcast(reduced, &devices, to_device)
247    }
248
249    fn all_gather<T, D, C>(tensors: Vec<T>, dim: usize, to_device: D, cat: C) -> Vec<T>
250    where
251        T: TensorMetadata<Device = Device<B>>,
252        D: Fn(T, &Device<B>) -> T,
253        C: Fn(Vec<T>, usize) -> T,
254    {
255        Self::devices(&tensors)
256            .iter()
257            .map(|device| cat(Self::moved_to(&tensors, device, &to_device), dim))
258            .collect()
259    }
260
261    /// Only each member's chunk moves to it.
262    fn reduce_scatter<T, D, S, R>(
263        tensors: Vec<T>,
264        dim: usize,
265        to_device: D,
266        slice: S,
267        reduce: R,
268    ) -> Vec<T>
269    where
270        T: TensorMetadata<Device = Device<B>>,
271        D: Fn(T, &Device<B>) -> T,
272        S: Fn(T, &[Slice]) -> T,
273        R: Fn(Vec<T>) -> T,
274    {
275        let split = DimSplit::new(&tensors, dim);
276        Self::devices(&tensors)
277            .iter()
278            .enumerate()
279            .map(|(member, device)| {
280                let chunks =
281                    Self::member_chunks(&tensors, &split, member, device, &to_device, &slice);
282                reduce(chunks)
283            })
284            .collect()
285    }
286
287    /// Only the chunks that change member move.
288    fn all_to_all<T, D, S, C>(
289        tensors: Vec<T>,
290        split_dim: usize,
291        concat_dim: usize,
292        to_device: D,
293        slice: S,
294        cat: C,
295    ) -> Vec<T>
296    where
297        T: TensorMetadata<Device = Device<B>>,
298        D: Fn(T, &Device<B>) -> T,
299        S: Fn(T, &[Slice]) -> T,
300        C: Fn(Vec<T>, usize) -> T,
301    {
302        let split = DimSplit::new(&tensors, split_dim);
303        Self::devices(&tensors)
304            .iter()
305            .enumerate()
306            .map(|(member, device)| {
307                let chunks =
308                    Self::member_chunks(&tensors, &split, member, device, &to_device, &slice);
309                cat(chunks, concat_dim)
310            })
311            .collect()
312    }
313
314    fn broadcast<T, D>(tensor: T, devices: &[Device<B>], to_device: D) -> Vec<T>
315    where
316        T: Clone,
317        D: Fn(T, &Device<B>) -> T,
318    {
319        devices
320            .iter()
321            .map(|device| to_device(tensor.clone(), device))
322            .collect()
323    }
324
325    fn devices<T>(tensors: &[T]) -> Vec<Device<B>>
326    where
327        T: TensorMetadata<Device = Device<B>>,
328    {
329        tensors.iter().map(TensorMetadata::device).collect()
330    }
331
332    fn moved_to<T, D>(tensors: &[T], device: &Device<B>, to_device: &D) -> Vec<T>
333    where
334        T: Clone,
335        D: Fn(T, &Device<B>) -> T,
336    {
337        tensors
338            .iter()
339            .map(|tensor| to_device(tensor.clone(), device))
340            .collect()
341    }
342
343    /// `member`'s chunk of every member's tensor, moved to `device`.
344    fn member_chunks<T, D, S>(
345        tensors: &[T],
346        split: &DimSplit,
347        member: usize,
348        device: &Device<B>,
349        to_device: &D,
350        slice: &S,
351    ) -> Vec<T>
352    where
353        T: Clone,
354        D: Fn(T, &Device<B>) -> T,
355        S: Fn(T, &[Slice]) -> T,
356    {
357        let slices = split.slices(member);
358        tensors
359            .iter()
360            .map(|tensor| to_device(slice(tensor.clone(), &slices), device))
361            .collect()
362    }
363
364    /// A mean divides each part first, as NCCL averages floats, so a sum past the type's range
365    /// still averages.
366    fn float_reduced(parts: Vec<FloatTensor<B>>, op: ReduceOperation) -> FloatTensor<B> {
367        let members = parts.len();
368        let parts = parts.into_iter().map(|part| match op {
369            ReduceOperation::Sum => part,
370            ReduceOperation::Mean => B::float_div_scalar(part, members.into()),
371        });
372        parts.reduce(B::float_add).expect("a group has a member")
373    }
374
375    /// A mean divides the sum, so it truncates once, as the backend's integer division does.
376    fn int_reduced(parts: Vec<IntTensor<B>>, op: ReduceOperation) -> IntTensor<B> {
377        let members = parts.len();
378        let sum = parts
379            .into_iter()
380            .reduce(B::int_add)
381            .expect("a group has a member");
382        match op {
383            ReduceOperation::Sum => sum,
384            ReduceOperation::Mean => B::int_div_scalar(sum, members.into()),
385        }
386    }
387}
388
389/// Dimension `dim`, of length `len`, cut into one chunk per member, a member being one of the
390/// tensors a collective is given: as even as possible, with the first members holding one more.
391///
392/// A caller that splits a tensor itself must cut it the same way, so this is public.
393#[derive(Clone, Copy, Debug)]
394pub struct DimSplit {
395    /// The dimension the chunks cut along.
396    pub dim: usize,
397    /// How many elements that dimension holds.
398    pub len: usize,
399    /// How many chunks it is cut into, one per member.
400    pub members: usize,
401}
402
403impl DimSplit {
404    fn new<T: TensorMetadata>(tensors: &[T], dim: usize) -> Self {
405        let split = Self {
406            dim,
407            len: tensors.first().map_or(0, |tensor| tensor.shape()[dim]),
408            members: tensors.len(),
409        };
410        assert!(
411            split.len >= split.members,
412            "dimension {dim} holds {} elements, fewer than the {} members it splits over",
413            split.len,
414            split.members,
415        );
416        split
417    }
418
419    /// The indices along `dim` that `member` holds.
420    pub fn range(&self, member: usize) -> Range<usize> {
421        let base = self.len / self.members;
422        let longer = self.len % self.members;
423        let start = member * base + member.min(longer);
424        start..start + base + usize::from(member < longer)
425    }
426
427    /// The slices that select `member`'s chunk.
428    pub fn slices(&self, member: usize) -> Vec<Slice> {
429        let range = self.range(member);
430        (0..=self.dim)
431            .map(|axis| match axis == self.dim {
432                true => Slice::new(range.start as isize, Some(range.end as isize), 1),
433                false => Slice::full(),
434            })
435            .collect()
436    }
437
438    /// Each member's chunk of `data`, in member order, cut on the host one at a time.
439    ///
440    /// # Panics
441    ///
442    /// When `data` is quantized, since its packed values and scales do not cut along a dim, or
443    /// when its `dim` does not hold `len` elements.
444    pub fn data_chunks(self, data: &TensorData) -> impl Iterator<Item = TensorData> + '_ {
445        assert!(
446            !matches!(data.dtype(), DType::QFloat(_)),
447            "A quantized tensor is not cut along a dim"
448        );
449        let shape = data.shape();
450        assert!(
451            shape.as_slice().get(self.dim) == Some(&self.len),
452            "dimension {} of a tensor of shape {shape:?} does not hold {} elements",
453            self.dim,
454            self.len,
455        );
456        (0..self.members).map(move |member| {
457            let mut chunk_shape = shape.clone();
458            chunk_shape[self.dim] = self.range(member).len();
459            TensorData::from_bytes(self.chunk_bytes(data, member), chunk_shape, data.dtype())
460        })
461    }
462
463    /// `member`'s bytes of `data`. A chunk in one piece is a window into `data` where its storage
464    /// has windows, so a file-backed tensor reads only that range.
465    fn chunk_bytes(&self, data: &TensorData, member: usize) -> Bytes {
466        let pieces = self.byte_ranges(data, member);
467        if pieces.len() == 1
468            && let Some(piece) = pieces.clone().next()
469            && let Ok(window) = data.bytes().view(piece.start, piece.end)
470        {
471            return window;
472        }
473        concatenated(data.as_bytes(), pieces)
474    }
475
476    /// Where `member`'s chunk sits in `data`: one byte range for each index of the dims before
477    /// `dim`.
478    fn byte_ranges(
479        &self,
480        data: &TensorData,
481        member: usize,
482    ) -> impl ExactSizeIterator<Item = Range<usize>> + Clone + use<> {
483        let shape = data.shape().as_slice();
484        let blocks: usize = shape[..self.dim].iter().product();
485        let row_size = shape[self.dim + 1..].iter().product::<usize>() * data.dtype().size();
486        let rows = self.range(member);
487        let len = self.len;
488        (0..blocks).map(move |block| {
489            let start = (block * len + rows.start) * row_size;
490            start..start + rows.len() * row_size
491        })
492    }
493}
494
495/// The `ranges` of `bytes` end to end, in a u128 buffer: it is aligned to `MAX_ALIGN`, so the
496/// result is not copied again to reach it.
497fn concatenated(bytes: &[u8], ranges: impl Iterator<Item = Range<usize>> + Clone) -> Bytes {
498    let len: usize = ranges.clone().map(|range| range.len()).sum();
499    let mut concatenated =
500        Bytes::from_elems(Vec::<u128>::with_capacity(len.div_ceil(size_of::<u128>())));
501    for range in ranges {
502        concatenated.extend_from_byte_slice(&bytes[range]);
503    }
504    concatenated
505}
506
507#[cfg(test)]
508mod tests {
509    use super::*;
510    use alloc::vec;
511    #[cfg(feature = "std")]
512    use burn_std::Reader;
513
514    #[test]
515    fn a_split_gives_the_first_members_one_more() {
516        let split = DimSplit {
517            dim: 0,
518            len: 7,
519            members: 3,
520        };
521        let ranges: Vec<_> = (0..3).map(|member| split.range(member)).collect();
522
523        assert_eq!(ranges, [0..3, 3..5, 5..7]);
524    }
525
526    #[test]
527    fn data_chunks_cut_every_block_along_an_inner_dim() {
528        let data = TensorData::new((0..12).collect::<Vec<i32>>(), [2, 3, 2]);
529        let split = DimSplit {
530            dim: 1,
531            len: 3,
532            members: 2,
533        };
534
535        let chunks: Vec<_> = split.data_chunks(&data).collect();
536
537        assert_eq!(
538            chunks,
539            [
540                TensorData::new(vec![0, 1, 2, 3, 6, 7, 8, 9], [2, 2, 2]),
541                TensorData::new(vec![4, 5, 10, 11], [2, 1, 2]),
542            ]
543        );
544    }
545
546    #[cfg(feature = "std")]
547    #[test]
548    fn data_chunks_along_the_first_dim_read_only_their_range_of_a_file() {
549        let whole = TensorData::new((0..12).collect::<Vec<i32>>(), [4, 3]);
550        let path = std::env::temp_dir().join(format!("burn-data-chunks-{}", std::process::id()));
551        std::fs::write(&path, whole.as_bytes()).unwrap();
552        let bytes = Bytes::from_file(&path, 48, 0);
553        let data = TensorData::from_bytes(bytes, [4, 3], DType::I32);
554        let split = DimSplit {
555            dim: 0,
556            len: 4,
557            members: 2,
558        };
559
560        let chunks: Vec<_> = split.data_chunks(&data).collect();
561
562        assert_eq!(
563            chunks,
564            [
565                TensorData::new(vec![0, 1, 2, 3, 4, 5], [2, 3]),
566                TensorData::new(vec![6, 7, 8, 9, 10, 11], [2, 3]),
567            ]
568        );
569        assert!(data.bytes().read(Reader::new().no_copy()).is_err());
570        std::fs::remove_file(path).unwrap();
571    }
572
573    #[test]
574    #[should_panic(expected = "dimension 2 of a tensor of shape")]
575    fn data_chunks_refuse_a_dim_the_tensor_lacks() {
576        let data = TensorData::new((0..12).collect::<Vec<i32>>(), [2, 6]);
577        let split = DimSplit {
578            dim: 2,
579            len: 4,
580            members: 2,
581        };
582
583        let _ = split.data_chunks(&data);
584    }
585}
burn_backend/backend/ops/collective.rs

burn_backend/backend/ops/
collective.rs
