Skip to main content

ProcessGroup

Struct ProcessGroup 

Source
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

Source

pub fn new(transport: Arc<dyn Transport>) -> ProcessGroup

Source

pub fn rank(&self) -> u32

Source

pub fn world_size(&self) -> u32

Source

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.

Source

pub fn transport(&self) -> &Arc<dyn Transport>

Source

pub fn barrier(&self) -> Result<(), CollectiveError>

Source

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).

Source

pub fn recv_f32(&self, from: u32, tag: u32) -> Result<Vec<f32>, CollectiveError>

Receive a tensor from from. Length is taken from the wire.

Source

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.

Source

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.

Source

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.

Source

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().

Source

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.

Source

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

Source§

fn clone(&self) -> ProcessGroup

Returns a duplicate of the value. Read more
1.0.0 (const: unstable) · Source§

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source. Read more

Auto Trait Implementations§

Blanket Implementations§

Source§

impl<T> Any for T
where T: 'static + ?Sized,

Source§

fn type_id(&self) -> TypeId

Gets the TypeId of self. Read more
Source§

impl<T> Borrow<T> for T
where T: ?Sized,

Source§

fn borrow(&self) -> &T

Immutably borrows from an owned value. Read more
Source§

impl<T> BorrowMut<T> for T
where T: ?Sized,

Source§

fn borrow_mut(&mut self) -> &mut T

Mutably borrows from an owned value. Read more
Source§

impl<T> CloneToUninit for T
where T: Clone,

Source§

unsafe fn clone_to_uninit(&self, dest: *mut u8)

🔬This is a nightly-only experimental API. (clone_to_uninit)
Performs copy-assignment from self to dest. Read more
Source§

impl<T> From<T> for T

Source§

fn from(t: T) -> T

Returns the argument unchanged.

Source§

impl<T, U> Into<U> for T
where U: From<T>,

Source§

fn into(self) -> U

Calls U::from(self).

That is, this conversion is whatever the implementation of From<T> for U chooses to do.

Source§

impl<T> IntoEither for T

Source§

fn into_either(self, into_left: bool) -> Either<Self, Self>

Converts 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 more
Source§

fn into_either_with<F>(self, into_left: F) -> Either<Self, Self>
where F: FnOnce(&Self) -> bool,

Converts 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
Source§

impl<T> Pointable for T

Source§

const ALIGN: usize

The alignment of pointer.
Source§

type Init = T

The type for initializers.
Source§

unsafe fn init(init: <T as Pointable>::Init) -> usize

Initializes a with the given initializer. Read more
Source§

unsafe fn deref<'a>(ptr: usize) -> &'a T

Dereferences the given pointer. Read more
Source§

unsafe fn deref_mut<'a>(ptr: usize) -> &'a mut T

Mutably dereferences the given pointer. Read more
Source§

unsafe fn drop(ptr: usize)

Drops the object pointed to by the given pointer. Read more
Source§

impl<T> ToOwned for T
where T: Clone,

Source§

type Owned = T

The resulting type after obtaining ownership.
Source§

fn to_owned(&self) -> T

Creates owned data from borrowed data, usually by cloning. Read more
Source§

fn clone_into(&self, target: &mut T)

Uses borrowed data to replace owned data, usually by cloning. Read more
Source§

impl<T, U> TryFrom<U> for T
where U: Into<T>,

Source§

type Error = Infallible

The type returned in the event of a conversion error.
Source§

fn try_from(value: U) -> Result<T, <T as TryFrom<U>>::Error>

Performs the conversion.
Source§

impl<T, U> TryInto<U> for T
where U: TryFrom<T>,

Source§

type Error = <U as TryFrom<T>>::Error

The type returned in the event of a conversion error.
Source§

fn try_into(self) -> Result<U, <U as TryFrom<T>>::Error>

Performs the conversion.
Source§

impl<T> WasmNotSend for T
where T: Send,

Source§

impl<T> WasmNotSendSync for T

Source§

impl<T> WasmNotSync for T
where T: Sync,