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

Examples found in repository ?

examples/transport_bench.rs (line 44)

28fn main() {
29    let world = 2u32;
30    let listeners: Vec<TcpListener> = (0..world)
31        .map(|_| TcpListener::bind((Ipv4Addr::LOCALHOST, 0)).unwrap())
32        .collect();
33    let addrs: Vec<SocketAddr> = listeners.iter().map(|l| l.local_addr().unwrap()).collect();
34
35    println!("transport: loopback TCP, world_size = {world}");
36    let handles: Vec<_> = listeners
37        .into_iter()
38        .enumerate()
39        .map(|(rank, listener)| {
40            let addrs = addrs.clone();
41            thread::spawn(move || {
42                let t = NetTransport::from_listener(rank as u32, world, listener, addrs, 64 << 20)
43                    .unwrap();
44                bench(&ProcessGroup::new(Arc::new(t)));
45            })
46        })
47        .collect();
48    for h in handles {
49        h.join().unwrap();
50    }
51}

Source

pub fn rank(&self) -> u32

Examples found in repository ?

examples/transport_bench.rs (line 54)

53fn bench(g: &ProcessGroup) {
54    let rank = g.rank();
55    g.barrier().unwrap();
56
57    // ---- point-to-point round-trip latency (ping-pong) ----
58    let iters = 5000;
59    if rank == 0 {
60        let p = [1.0f32];
61        for _ in 0..200 {
62            g.send_f32(1, 0, &p).unwrap();
63            g.recv_f32(1, 0).unwrap();
64        }
65        let t0 = Instant::now();
66        for _ in 0..iters {
67            g.send_f32(1, 0, &p).unwrap();
68            g.recv_f32(1, 0).unwrap();
69        }
70        let us = t0.elapsed().as_secs_f64() * 1e6 / iters as f64;
71        println!(
72            "[latency]   ping-pong: {us:6.2} µs RTT   ({:.2} µs one-way)   [{iters} iters]",
73            us / 2.0
74        );
75    } else {
76        for _ in 0..(200 + iters) {
77            let x = g.recv_f32(0, 0).unwrap();
78            g.send_f32(0, 0, &x).unwrap();
79        }
80    }
81    g.barrier().unwrap();
82
83    // ---- streaming bandwidth across payload sizes ----
84    for &elems in &[256usize, 16_384, 262_144, 4_194_304] {
85        let bytes = elems * 4;
86        let reps = (512 * 1024 * 1024 / bytes).clamp(8, 200_000);
87        if rank == 0 {
88            let buf = vec![0.5f32; elems];
89            g.send_f32(1, 1, &[0.0]).unwrap(); // sync start
90            let t0 = Instant::now();
91            for _ in 0..reps {
92                g.send_f32(1, 2, &buf).unwrap();
93            }
94            g.recv_f32(1, 3).unwrap(); // ack: all received
95            let secs = t0.elapsed().as_secs_f64();
96            let gbps = (reps * bytes) as f64 / secs / 1e9;
97            println!(
98                "[bandwidth] {:>9} B × {:>6} = {:>5.1} GB  ->  {gbps:6.2} GB/s",
99                bytes,
100                reps,
101                (reps * bytes) as f64 / 1e9
102            );
103        } else {
104            g.recv_f32(0, 1).unwrap();
105            for _ in 0..reps {
106                g.recv_f32(0, 2).unwrap();
107            }
108            g.send_f32(0, 3, &[1.0]).unwrap();
109        }
110        g.barrier().unwrap();
111    }
112
113    // ---- collective latency for a hidden-state-sized vector ----
114    // e.g. batch1 × seq1 × d_model=4096 — the kind of all_reduce a
115    // tensor-parallel layer issues, or a broadcast the pipeline issues.
116    for &elems in &[4096usize, 4096 * 8] {
117        let mut data = vec![1.0f32; elems];
118        let iters = 2000;
119        for _ in 0..50 {
120            g.all_reduce(&mut data, ReduceKind::Sum).unwrap();
121        }
122        g.barrier().unwrap();
123        let t0 = Instant::now();
124        for _ in 0..iters {
125            g.all_reduce(&mut data, ReduceKind::Sum).unwrap();
126        }
127        let us = t0.elapsed().as_secs_f64() * 1e6 / iters as f64;
128        if rank == 0 {
129            println!("[collective] all_reduce {:>6} f32: {us:6.2} µs/op", elems);
130        }
131        g.barrier().unwrap();
132    }
133
134    g.barrier().unwrap();
135}

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>

Examples found in repository ?

examples/transport_bench.rs (line 55)

53fn bench(g: &ProcessGroup) {
54    let rank = g.rank();
55    g.barrier().unwrap();
56
57    // ---- point-to-point round-trip latency (ping-pong) ----
58    let iters = 5000;
59    if rank == 0 {
60        let p = [1.0f32];
61        for _ in 0..200 {
62            g.send_f32(1, 0, &p).unwrap();
63            g.recv_f32(1, 0).unwrap();
64        }
65        let t0 = Instant::now();
66        for _ in 0..iters {
67            g.send_f32(1, 0, &p).unwrap();
68            g.recv_f32(1, 0).unwrap();
69        }
70        let us = t0.elapsed().as_secs_f64() * 1e6 / iters as f64;
71        println!(
72            "[latency]   ping-pong: {us:6.2} µs RTT   ({:.2} µs one-way)   [{iters} iters]",
73            us / 2.0
74        );
75    } else {
76        for _ in 0..(200 + iters) {
77            let x = g.recv_f32(0, 0).unwrap();
78            g.send_f32(0, 0, &x).unwrap();
79        }
80    }
81    g.barrier().unwrap();
82
83    // ---- streaming bandwidth across payload sizes ----
84    for &elems in &[256usize, 16_384, 262_144, 4_194_304] {
85        let bytes = elems * 4;
86        let reps = (512 * 1024 * 1024 / bytes).clamp(8, 200_000);
87        if rank == 0 {
88            let buf = vec![0.5f32; elems];
89            g.send_f32(1, 1, &[0.0]).unwrap(); // sync start
90            let t0 = Instant::now();
91            for _ in 0..reps {
92                g.send_f32(1, 2, &buf).unwrap();
93            }
94            g.recv_f32(1, 3).unwrap(); // ack: all received
95            let secs = t0.elapsed().as_secs_f64();
96            let gbps = (reps * bytes) as f64 / secs / 1e9;
97            println!(
98                "[bandwidth] {:>9} B × {:>6} = {:>5.1} GB  ->  {gbps:6.2} GB/s",
99                bytes,
100                reps,
101                (reps * bytes) as f64 / 1e9
102            );
103        } else {
104            g.recv_f32(0, 1).unwrap();
105            for _ in 0..reps {
106                g.recv_f32(0, 2).unwrap();
107            }
108            g.send_f32(0, 3, &[1.0]).unwrap();
109        }
110        g.barrier().unwrap();
111    }
112
113    // ---- collective latency for a hidden-state-sized vector ----
114    // e.g. batch1 × seq1 × d_model=4096 — the kind of all_reduce a
115    // tensor-parallel layer issues, or a broadcast the pipeline issues.
116    for &elems in &[4096usize, 4096 * 8] {
117        let mut data = vec![1.0f32; elems];
118        let iters = 2000;
119        for _ in 0..50 {
120            g.all_reduce(&mut data, ReduceKind::Sum).unwrap();
121        }
122        g.barrier().unwrap();
123        let t0 = Instant::now();
124        for _ in 0..iters {
125            g.all_reduce(&mut data, ReduceKind::Sum).unwrap();
126        }
127        let us = t0.elapsed().as_secs_f64() * 1e6 / iters as f64;
128        if rank == 0 {
129            println!("[collective] all_reduce {:>6} f32: {us:6.2} µs/op", elems);
130        }
131        g.barrier().unwrap();
132    }
133
134    g.barrier().unwrap();
135}

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

Examples found in repository ?

examples/transport_bench.rs (line 62)

53fn bench(g: &ProcessGroup) {
54    let rank = g.rank();
55    g.barrier().unwrap();
56
57    // ---- point-to-point round-trip latency (ping-pong) ----
58    let iters = 5000;
59    if rank == 0 {
60        let p = [1.0f32];
61        for _ in 0..200 {
62            g.send_f32(1, 0, &p).unwrap();
63            g.recv_f32(1, 0).unwrap();
64        }
65        let t0 = Instant::now();
66        for _ in 0..iters {
67            g.send_f32(1, 0, &p).unwrap();
68            g.recv_f32(1, 0).unwrap();
69        }
70        let us = t0.elapsed().as_secs_f64() * 1e6 / iters as f64;
71        println!(
72            "[latency]   ping-pong: {us:6.2} µs RTT   ({:.2} µs one-way)   [{iters} iters]",
73            us / 2.0
74        );
75    } else {
76        for _ in 0..(200 + iters) {
77            let x = g.recv_f32(0, 0).unwrap();
78            g.send_f32(0, 0, &x).unwrap();
79        }
80    }
81    g.barrier().unwrap();
82
83    // ---- streaming bandwidth across payload sizes ----
84    for &elems in &[256usize, 16_384, 262_144, 4_194_304] {
85        let bytes = elems * 4;
86        let reps = (512 * 1024 * 1024 / bytes).clamp(8, 200_000);
87        if rank == 0 {
88            let buf = vec![0.5f32; elems];
89            g.send_f32(1, 1, &[0.0]).unwrap(); // sync start
90            let t0 = Instant::now();
91            for _ in 0..reps {
92                g.send_f32(1, 2, &buf).unwrap();
93            }
94            g.recv_f32(1, 3).unwrap(); // ack: all received
95            let secs = t0.elapsed().as_secs_f64();
96            let gbps = (reps * bytes) as f64 / secs / 1e9;
97            println!(
98                "[bandwidth] {:>9} B × {:>6} = {:>5.1} GB  ->  {gbps:6.2} GB/s",
99                bytes,
100                reps,
101                (reps * bytes) as f64 / 1e9
102            );
103        } else {
104            g.recv_f32(0, 1).unwrap();
105            for _ in 0..reps {
106                g.recv_f32(0, 2).unwrap();
107            }
108            g.send_f32(0, 3, &[1.0]).unwrap();
109        }
110        g.barrier().unwrap();
111    }
112
113    // ---- collective latency for a hidden-state-sized vector ----
114    // e.g. batch1 × seq1 × d_model=4096 — the kind of all_reduce a
115    // tensor-parallel layer issues, or a broadcast the pipeline issues.
116    for &elems in &[4096usize, 4096 * 8] {
117        let mut data = vec![1.0f32; elems];
118        let iters = 2000;
119        for _ in 0..50 {
120            g.all_reduce(&mut data, ReduceKind::Sum).unwrap();
121        }
122        g.barrier().unwrap();
123        let t0 = Instant::now();
124        for _ in 0..iters {
125            g.all_reduce(&mut data, ReduceKind::Sum).unwrap();
126        }
127        let us = t0.elapsed().as_secs_f64() * 1e6 / iters as f64;
128        if rank == 0 {
129            println!("[collective] all_reduce {:>6} f32: {us:6.2} µs/op", elems);
130        }
131        g.barrier().unwrap();
132    }
133
134    g.barrier().unwrap();
135}

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.

Examples found in repository ?

examples/transport_bench.rs (line 63)

53fn bench(g: &ProcessGroup) {
54    let rank = g.rank();
55    g.barrier().unwrap();
56
57    // ---- point-to-point round-trip latency (ping-pong) ----
58    let iters = 5000;
59    if rank == 0 {
60        let p = [1.0f32];
61        for _ in 0..200 {
62            g.send_f32(1, 0, &p).unwrap();
63            g.recv_f32(1, 0).unwrap();
64        }
65        let t0 = Instant::now();
66        for _ in 0..iters {
67            g.send_f32(1, 0, &p).unwrap();
68            g.recv_f32(1, 0).unwrap();
69        }
70        let us = t0.elapsed().as_secs_f64() * 1e6 / iters as f64;
71        println!(
72            "[latency]   ping-pong: {us:6.2} µs RTT   ({:.2} µs one-way)   [{iters} iters]",
73            us / 2.0
74        );
75    } else {
76        for _ in 0..(200 + iters) {
77            let x = g.recv_f32(0, 0).unwrap();
78            g.send_f32(0, 0, &x).unwrap();
79        }
80    }
81    g.barrier().unwrap();
82
83    // ---- streaming bandwidth across payload sizes ----
84    for &elems in &[256usize, 16_384, 262_144, 4_194_304] {
85        let bytes = elems * 4;
86        let reps = (512 * 1024 * 1024 / bytes).clamp(8, 200_000);
87        if rank == 0 {
88            let buf = vec![0.5f32; elems];
89            g.send_f32(1, 1, &[0.0]).unwrap(); // sync start
90            let t0 = Instant::now();
91            for _ in 0..reps {
92                g.send_f32(1, 2, &buf).unwrap();
93            }
94            g.recv_f32(1, 3).unwrap(); // ack: all received
95            let secs = t0.elapsed().as_secs_f64();
96            let gbps = (reps * bytes) as f64 / secs / 1e9;
97            println!(
98                "[bandwidth] {:>9} B × {:>6} = {:>5.1} GB  ->  {gbps:6.2} GB/s",
99                bytes,
100                reps,
101                (reps * bytes) as f64 / 1e9
102            );
103        } else {
104            g.recv_f32(0, 1).unwrap();
105            for _ in 0..reps {
106                g.recv_f32(0, 2).unwrap();
107            }
108            g.send_f32(0, 3, &[1.0]).unwrap();
109        }
110        g.barrier().unwrap();
111    }
112
113    // ---- collective latency for a hidden-state-sized vector ----
114    // e.g. batch1 × seq1 × d_model=4096 — the kind of all_reduce a
115    // tensor-parallel layer issues, or a broadcast the pipeline issues.
116    for &elems in &[4096usize, 4096 * 8] {
117        let mut data = vec![1.0f32; elems];
118        let iters = 2000;
119        for _ in 0..50 {
120            g.all_reduce(&mut data, ReduceKind::Sum).unwrap();
121        }
122        g.barrier().unwrap();
123        let t0 = Instant::now();
124        for _ in 0..iters {
125            g.all_reduce(&mut data, ReduceKind::Sum).unwrap();
126        }
127        let us = t0.elapsed().as_secs_f64() * 1e6 / iters as f64;
128        if rank == 0 {
129            println!("[collective] all_reduce {:>6} f32: {us:6.2} µs/op", elems);
130        }
131        g.barrier().unwrap();
132    }
133
134    g.barrier().unwrap();
135}

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.

Examples found in repository ?

examples/transport_bench.rs (line 120)

53fn bench(g: &ProcessGroup) {
54    let rank = g.rank();
55    g.barrier().unwrap();
56
57    // ---- point-to-point round-trip latency (ping-pong) ----
58    let iters = 5000;
59    if rank == 0 {
60        let p = [1.0f32];
61        for _ in 0..200 {
62            g.send_f32(1, 0, &p).unwrap();
63            g.recv_f32(1, 0).unwrap();
64        }
65        let t0 = Instant::now();
66        for _ in 0..iters {
67            g.send_f32(1, 0, &p).unwrap();
68            g.recv_f32(1, 0).unwrap();
69        }
70        let us = t0.elapsed().as_secs_f64() * 1e6 / iters as f64;
71        println!(
72            "[latency]   ping-pong: {us:6.2} µs RTT   ({:.2} µs one-way)   [{iters} iters]",
73            us / 2.0
74        );
75    } else {
76        for _ in 0..(200 + iters) {
77            let x = g.recv_f32(0, 0).unwrap();
78            g.send_f32(0, 0, &x).unwrap();
79        }
80    }
81    g.barrier().unwrap();
82
83    // ---- streaming bandwidth across payload sizes ----
84    for &elems in &[256usize, 16_384, 262_144, 4_194_304] {
85        let bytes = elems * 4;
86        let reps = (512 * 1024 * 1024 / bytes).clamp(8, 200_000);
87        if rank == 0 {
88            let buf = vec![0.5f32; elems];
89            g.send_f32(1, 1, &[0.0]).unwrap(); // sync start
90            let t0 = Instant::now();
91            for _ in 0..reps {
92                g.send_f32(1, 2, &buf).unwrap();
93            }
94            g.recv_f32(1, 3).unwrap(); // ack: all received
95            let secs = t0.elapsed().as_secs_f64();
96            let gbps = (reps * bytes) as f64 / secs / 1e9;
97            println!(
98                "[bandwidth] {:>9} B × {:>6} = {:>5.1} GB  ->  {gbps:6.2} GB/s",
99                bytes,
100                reps,
101                (reps * bytes) as f64 / 1e9
102            );
103        } else {
104            g.recv_f32(0, 1).unwrap();
105            for _ in 0..reps {
106                g.recv_f32(0, 2).unwrap();
107            }
108            g.send_f32(0, 3, &[1.0]).unwrap();
109        }
110        g.barrier().unwrap();
111    }
112
113    // ---- collective latency for a hidden-state-sized vector ----
114    // e.g. batch1 × seq1 × d_model=4096 — the kind of all_reduce a
115    // tensor-parallel layer issues, or a broadcast the pipeline issues.
116    for &elems in &[4096usize, 4096 * 8] {
117        let mut data = vec![1.0f32; elems];
118        let iters = 2000;
119        for _ in 0..50 {
120            g.all_reduce(&mut data, ReduceKind::Sum).unwrap();
121        }
122        g.barrier().unwrap();
123        let t0 = Instant::now();
124        for _ in 0..iters {
125            g.all_reduce(&mut data, ReduceKind::Sum).unwrap();
126        }
127        let us = t0.elapsed().as_secs_f64() * 1e6 / iters as f64;
128        if rank == 0 {
129            println!("[collective] all_reduce {:>6} f32: {us:6.2} µs/op", elems);
130        }
131        g.barrier().unwrap();
132    }
133
134    g.barrier().unwrap();
135}

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.