silvic_core/
collective.rs

1//! The collectives distributed training runs on, in the shape of LightGBM's `Network` (allreduce, allgather) but
2//! over a transport the caller supplies.
3//!
4//! Every rank calls the same collectives in the same order with buffers of the same length; a rank that diverges
5//! deadlocks or fails the job, as in MPI. In quantized training every histogram value is an integer below 2^53, so
6//! `allreduce_sum_f64` is exact in any reduction order and a model trained on N ranks equals one trained on one.
7
8use crate::error::Result;
9
10pub trait Collective: Send {
11    fn rank(&self) -> usize;
12
13    fn world_size(&self) -> usize;
14
15    /// Elementwise sum across ranks; every rank ends with the sum.
16    fn allreduce_sum_f64(&mut self, data: &mut [f64]) -> Result<()>;
17
18    fn allreduce_sum_i64(&mut self, data: &mut [i64]) -> Result<()>;
19
20    /// Elementwise maximum across ranks.
21    fn allreduce_max_f64(&mut self, data: &mut [f64]) -> Result<()>;
22
23    /// Every rank's bytes, indexed by rank, on every rank.
24    fn allgather_bytes(&mut self, local: &[u8]) -> Result<Vec<Vec<u8>>>;
25}
26
27/// A world of one rank: every collective returns its input.
28#[derive(Debug, Default, Clone, Copy)]
29pub struct LocalCollective;
30
31impl Collective for LocalCollective {
32    fn rank(&self) -> usize {
33        0
34    }
35
36    fn world_size(&self) -> usize {
37        1
38    }
39
40    fn allreduce_sum_f64(&mut self, _data: &mut [f64]) -> Result<()> {
41        Ok(())
42    }
43
44    fn allreduce_sum_i64(&mut self, _data: &mut [i64]) -> Result<()> {
45        Ok(())
46    }
47
48    fn allreduce_max_f64(&mut self, _data: &mut [f64]) -> Result<()> {
49        Ok(())
50    }
51
52    fn allgather_bytes(&mut self, local: &[u8]) -> Result<Vec<Vec<u8>>> {
53        Ok(vec![local.to_vec()])
54    }
55}
silvic_core/collective.rs

silvic_core/
collective.rs
