pub struct ProcessGroup { /* private fields */ }Expand description
A handle to the collective-communication world: a rank, a size, and the transport that connects them.
Wraps an Arc<dyn Transport> so the same group can be cloned into
the generator loop, the weight loader, and (later) per-layer
collective ops without re-establishing connections.
Implementations§
Source§impl ProcessGroup
impl ProcessGroup
pub fn new(transport: Arc<dyn Transport>) -> Self
pub fn rank(&self) -> u32
pub fn world_size(&self) -> u32
Sourcepub fn is_leader(&self) -> bool
pub fn is_leader(&self) -> bool
true on the lowest rank — the canonical place to print, sample,
or own the final logits in a pipeline.
pub fn transport(&self) -> &Arc<dyn Transport>
pub fn barrier(&self) -> Result<(), CollectiveError>
Sourcepub fn send_f32(
&self,
to: u32,
tag: u32,
data: &[f32],
) -> Result<(), CollectiveError>
pub fn send_f32( &self, to: u32, tag: u32, data: &[f32], ) -> Result<(), CollectiveError>
Send a hidden-state (or any) tensor to to. tag must be below
TAG_RESERVED_BASE (those tags are reserved for collectives).
Sourcepub fn recv_f32(&self, from: u32, tag: u32) -> Result<Vec<f32>, CollectiveError>
pub fn recv_f32(&self, from: u32, tag: u32) -> Result<Vec<f32>, CollectiveError>
Receive a tensor from from. Length is taken from the wire.
Sourcepub fn all_reduce(
&self,
data: &mut [f32],
op: ReduceKind,
) -> Result<(), CollectiveError>
pub fn all_reduce( &self, data: &mut [f32], op: ReduceKind, ) -> Result<(), CollectiveError>
In-place all-reduce: on return every rank holds
op({each rank's input}). All ranks must pass the same length.
Bandwidth-optimal ring all-reduce (reduce-scatter then
all-gather): every rank moves only ~2·(n-1)/n · len floats in and
out, independent of n, versus the old gather-to-root that funneled
O(n)·len through rank 0. Deadlock-free over [NetTransport]
because each rank’s background reader thread always drains its
socket, so a send never blocks on a peer that is itself sending.
Sourcepub fn all_reduce_typed(
&self,
data: &mut [u8],
dtype: DType,
op: ReduceKind,
) -> Result<(), CollectiveError>
pub fn all_reduce_typed( &self, data: &mut [u8], dtype: DType, op: ReduceKind, ) -> Result<(), CollectiveError>
In-place all-reduce over any float dtype (F16/BF16/F32/F64),
operating on the raw little-endian bytes. The wire carries the native
dtype — an fp16 reduction moves half the bytes of the f32 path, the
direct bandwidth win of mixed precision — while the reduction itself
accumulates in f64 so range isn’t lost mid-sum. Same bandwidth-optimal
ring as Self::all_reduce. data.len() must be a multiple of the
dtype’s element size, and every rank must pass the same dtype + length.
Sourcepub fn spawn_all_reduce(
self: &Arc<Self>,
data: Vec<f32>,
op: ReduceKind,
) -> JoinHandle<Vec<f32>>
pub fn spawn_all_reduce( self: &Arc<Self>, data: Vec<f32>, op: ReduceKind, ) -> JoinHandle<Vec<f32>>
Non-blocking all-reduce: takes ownership of data, runs the
collective on a background thread, and returns a handle whose
join() yields the reduced vector. This is the building block for
overlapping gradient communication with backward compute — fire it
per bucket as gradients are produced, keep computing, then join.
Sourcepub fn all_gather(&self, local: &[f32]) -> Result<Vec<f32>, CollectiveError>
pub fn all_gather(&self, local: &[f32]) -> Result<Vec<f32>, CollectiveError>
All-gather: every rank contributes local (same length on every
rank) and receives the concatenation in rank order, length
world_size * local.len().
Sourcepub fn broadcast(
&self,
root: u32,
data: &mut [f32],
) -> Result<(), CollectiveError>
pub fn broadcast( &self, root: u32, data: &mut [f32], ) -> Result<(), CollectiveError>
Broadcast data from root to every rank. On non-root ranks
data is overwritten; its length must already match the root’s.
Sourcepub fn federated_average(
&self,
data: &mut [f32],
deadline: Duration,
) -> Result<usize, CollectiveError>
pub fn federated_average( &self, data: &mut [f32], deadline: Duration, ) -> Result<usize, CollectiveError>
Bounded-staleness federated averaging. Rank 0 collects each rank’s
vector but skips any that miss deadline, averages the arrivals
(always including its own), and broadcasts the result. Returns the
participant count on rank 0 (callers weight / log dropouts; non-roots
get the averaged data and a count of 0). This is how a slow or
offline edge client can’t stall a training round.
Trait Implementations§
Source§impl Clone for ProcessGroup
impl Clone for ProcessGroup
Source§fn clone(&self) -> ProcessGroup
fn clone(&self) -> ProcessGroup
1.0.0 (const: unstable) · Source§fn clone_from(&mut self, source: &Self)
fn clone_from(&mut self, source: &Self)
source. Read more