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>) -> ProcessGroup
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<ProcessGroup>,
data: Vec<f32>,
op: ReduceKind,
) -> JoinHandle<Vec<f32>>
pub fn spawn_all_reduce( self: &Arc<ProcessGroup>, 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 moreAuto Trait Implementations§
impl !RefUnwindSafe for ProcessGroup
impl !UnwindSafe for ProcessGroup
impl Freeze for ProcessGroup
impl Send for ProcessGroup
impl Sync for ProcessGroup
impl Unpin for ProcessGroup
impl UnsafeUnpin for ProcessGroup
Blanket Implementations§
Source§impl<T> BorrowMut<T> for Twhere
T: ?Sized,
impl<T> BorrowMut<T> for Twhere
T: ?Sized,
Source§fn borrow_mut(&mut self) -> &mut T
fn borrow_mut(&mut self) -> &mut T
Source§impl<T> CloneToUninit for Twhere
T: Clone,
impl<T> CloneToUninit for Twhere
T: Clone,
Source§impl<T> IntoEither for T
impl<T> IntoEither for T
Source§fn into_either(self, into_left: bool) -> Either<Self, Self>
fn into_either(self, into_left: bool) -> Either<Self, Self>
self into a Left variant of Either<Self, Self>
if into_left is true.
Converts self into a Right variant of Either<Self, Self>
otherwise. Read moreSource§fn into_either_with<F>(self, into_left: F) -> Either<Self, Self>
fn into_either_with<F>(self, into_left: F) -> Either<Self, Self>
self into a Left variant of Either<Self, Self>
if into_left(&self) returns true.
Converts self into a Right variant of Either<Self, Self>
otherwise. Read more