metrale_comm/
nccl_backend.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: `NcclBackend`, the multi-rank [`CommBackend`](crate::CommBackend)
4//! over NCCL. Bootstrap: rank 0 creates an NCCL unique id and sends it over
5//! TCP to every other rank, then all ranks call `ncclCommInitRank`.
6//!
7//! At `world_size == 2`, once the add kernel is set, all-reduce is a grouped
8//! `ncclSend`/`ncclRecv` into a registered receive buffer plus a local BF16
9//! add; otherwise it is `ncclAllReduce`. After each collective the backend
10//! queries `ncclCommGetAsyncError`; a broadcast waits for completion for at
11//! most `COLLECTIVE_TIMEOUT_SECS`, an idle command receive without a deadline,
12//! and a failure in either marks the communicator unhealthy. `attempt_reconnect`
13//! aborts the communicator and bootstraps again. `METRALE_COMM_DIAGNOSTICS=1`
14//! logs every host submission.
15//!
16//! Owner: metrale-comm.
17//! Invariants:
18//! - Every collective and point-to-point operation first refuses an unhealthy
19//!   communicator and validates its byte count and peer (`begin_submission`).
20//! - A 2-rank send/recv all-reduce never receives more than `recv_capacity`
21//!   bytes: `all_reduce_2rank` checks before any NCCL call.
22//!
23//! Each `unsafe` block wraps one NCCL or CUDA driver call. The declarations
24//! follow nccl.h and cuda.h. Callers must pass device pointers valid for
25//! `bytes` on the communicator's device and streams that outlive the enqueued
26//! work.
27
28use anyhow::{Context, Result};
29use parking_lot::Mutex;
30use std::ffi::c_void;
31use std::io::{Read, Write};
32use std::net::{TcpListener, TcpStream};
33use std::ptr;
34use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
35
36use metrale_gpu_sys::nccl::{self, NcclComm, NcclDataType, NcclResult, NcclUniqueId};
37
38// 2026-09-26: CUDA driver calls for the receive buffer and the add kernel.
39unsafe extern "C" {
40    fn cuMemAlloc_v2(dptr: *mut u64, bytesize: usize) -> i32;
41    fn cuMemFree_v2(dptr: u64) -> i32;
42    fn cuLaunchKernel(
43        f: u64,
44        gridDimX: u32,
45        gridDimY: u32,
46        gridDimZ: u32,
47        blockDimX: u32,
48        blockDimY: u32,
49        blockDimZ: u32,
50        sharedMemBytes: u32,
51        hStream: u64,
52        kernelParams: *mut *mut c_void,
53        extra: *mut *mut c_void,
54    ) -> i32;
55}
56
57mod recv_buffer;
58use recv_buffer::ensure_payload_fits;
59pub use recv_buffer::{ALL_REDUCE_DTYPE_BYTES, required_model_recv_bytes, required_recv_bytes};
60
61/// 2026-09-26: Deadline, in seconds, for a broadcast to complete. It bounds
62/// only the completion polling, not an NCCL or driver call that hangs.
63pub(super) const COLLECTIVE_TIMEOUT_SECS: u64 = 30;
64
65/// 2026-09-26: NCCL communicator plus the streams, events and receive buffer
66/// its operations use.
67pub struct NcclBackend {
68    /// 2026-09-26: The communicator handle. Operations copy it out under the
69    /// lock and release the lock before calling NCCL; `reconnect_inner`
70    /// holds the lock while it replaces the handle.
71    comm: Mutex<NcclComm>,
72    diagnostics: crate::collective_diagnostics::Diagnostics,
73    rank: usize,
74    world_size: usize,
75    /// 2026-09-26: Stream this backend creates for `all_reduce_async`.
76    comm_stream: u64,
77    /// 2026-09-26: Recorded on the caller's compute stream before an async
78    /// all-reduce.
79    compute_done_event: u64,
80    /// 2026-09-26: Recorded on `comm_stream` after an async all-reduce.
81    comm_done_event: u64,
82    /// 2026-09-26: The caller's stream from `new`, used by every operation
83    /// that takes no stream argument.
84    legacy_stream: u64,
85    /// 2026-09-26: Receive buffer of the 2-rank send/recv all-reduce; 0 when
86    /// `world_size != 2`.
87    recv_buffer: u64,
88    /// 2026-09-26: Size of `recv_buffer` in bytes; 0 when `world_size != 2`.
89    recv_capacity: usize,
90    /// 2026-09-26: Handles from `register_buffer`, deregistered in `Drop` and
91    /// dropped without deregistering on reconnect.
92    registered_handles: Mutex<Vec<*mut c_void>>,
93    /// 2026-09-26: `bf16_add_inplace` handle from `set_add_kernel`; 0 until
94    /// set.
95    add_kernel: AtomicU64,
96    /// 2026-09-26: Set by an async error or a failed completion, cleared by a
97    /// successful reconnect.
98    unhealthy: AtomicBool,
99    /// 2026-09-26: Successful reconnects, for the log.
100    reconnect_count: AtomicU64,
101    /// 2026-09-26: Bootstrap address, reused by reconnect.
102    master_addr: String,
103    master_port: u16,
104}
105
106// 2026-09-26: SAFETY: The raw pointers are an NCCL communicator handle and
107// registration handles, which the backend only copies and passes to NCCL
108// (`register_buffer` also returns a copy as `u64`). The mutex guards only
109// reading and replacing the communicator handle; NCCL calls run after the
110// guard is released and are not serialised by it.
111unsafe impl Send for NcclBackend {}
112unsafe impl Sync for NcclBackend {}
113
114impl NcclBackend {
115    /// 2026-09-26: Bootstrap the communicator. Rank 0 binds `0.0.0.0:master_port`
116    /// and sends the unique id to `world_size - 1` connections; other ranks
117    /// connect to `master_addr:master_port`, retrying once a second for up to
118    /// 600 attempts. `recv_capacity` is the largest 2-rank all-reduce payload
119    /// in bytes (serve passes [`required_model_recv_bytes`]); it is used only
120    /// when `world_size == 2`.
121    ///
122    /// # Errors
123    /// An invalid `METRALE_COMM_DIAGNOSTICS` value or rank/world pair, a failed
124    /// bootstrap or NCCL init, a failed stream or event creation, a zero
125    /// `recv_capacity` at `world_size == 2`, or a failed receive-buffer
126    /// allocation. A failed registration of the receive buffer is only logged.
127    pub fn new(
128        rank: usize,
129        world_size: usize,
130        master_addr: &str,
131        master_port: u16,
132        stream: u64,
133        recv_capacity: usize,
134    ) -> Result<Self> {
135        let diagnostic_env = std::env::var(crate::collective_diagnostics::ENV).ok();
136        let diagnostics = crate::collective_diagnostics::Diagnostics::new(
137            diagnostic_env.as_deref(),
138            rank,
139            world_size,
140        )?;
141        Self::log_nccl_env_vars();
142
143        let unique_id = if rank == 0 {
144            let id = Self::generate_unique_id()?;
145            Self::distribute_id(&id, master_addr, master_port, world_size)?;
146            id
147        } else {
148            Self::receive_id(master_addr, master_port)?
149        };
150
151        let mut comm: NcclComm = ptr::null_mut();
152        let result =
153            unsafe { nccl::ncclCommInitRank(&mut comm, world_size as i32, unique_id, rank as i32) };
154        nccl::check_nccl(result, "ncclCommInitRank")?;
155
156        let comm_stream = nccl::create_stream()?;
157        let compute_done_event = nccl::create_event()?;
158        let comm_done_event = nccl::create_event()?;
159
160        let mut recv_buffer: u64 = 0;
161        if world_size == 2 {
162            if recv_capacity == 0 {
163                anyhow::bail!(
164                    "world_size == 2 requires a non-zero receive-buffer capacity; \
165                     compute it with required_recv_bytes(max_batch_tokens, hidden_size, \
166                     ALL_REDUCE_DTYPE_BYTES)"
167                );
168            }
169            let status = unsafe { cuMemAlloc_v2(&mut recv_buffer, recv_capacity) };
170            if status != 0 {
171                anyhow::bail!(
172                    "cuMemAlloc_v2 for recv_buffer ({recv_capacity} bytes) failed: status {status}"
173                );
174            }
175            let mut handle: *mut c_void = ptr::null_mut();
176            let result = unsafe {
177                nccl::ncclCommRegister(comm, recv_buffer as *mut c_void, recv_capacity, &mut handle)
178            };
179            if result != NcclResult::Success {
180                tracing::warn!("ncclCommRegister for recv_buffer failed (non-fatal): {result:?}");
181            } else {
182                tracing::info!(
183                    "Registered recv_buffer ({} KB) with NCCL",
184                    recv_capacity / 1024
185                );
186            }
187        }
188
189        Ok(Self {
190            comm: Mutex::new(comm),
191            diagnostics,
192            rank,
193            world_size,
194            comm_stream,
195            compute_done_event,
196            comm_done_event,
197            legacy_stream: stream,
198            recv_buffer,
199            recv_capacity: if world_size == 2 { recv_capacity } else { 0 },
200            registered_handles: Mutex::new(Vec::new()),
201            add_kernel: AtomicU64::new(0),
202            unhealthy: AtomicBool::new(false),
203            reconnect_count: AtomicU64::new(0),
204            master_addr: master_addr.to_owned(),
205            master_port,
206        })
207    }
208
209    /// 2026-09-26: Log the listed NCCL environment variables at `info` when
210    /// set, `debug` when not.
211    fn log_nccl_env_vars() {
212        let vars = [
213            "NCCL_TIMEOUT",
214            "NCCL_WATCHDOG_TIMEOUT",
215            "NCCL_IB_TIMEOUT",
216            "NCCL_IB_RETRY_CNT",
217            "NCCL_SOCKET_IFNAME",
218            "NCCL_DEBUG",
219        ];
220        for var in &vars {
221            match std::env::var(var) {
222                Ok(val) => tracing::info!("NCCL env: {var}={val}"),
223                Err(_) => tracing::debug!("NCCL env: {var} not set"),
224            }
225        }
226    }
227
228    /// 2026-09-26: Query `ncclCommGetAsyncError`. Returns `true` when there is
229    /// no error; otherwise, or when the query itself fails, sets the unhealthy
230    /// flag and returns `false`.
231    fn check_async_error(&self, comm: NcclComm) -> bool {
232        let mut async_err = NcclResult::Success;
233        let result = unsafe { nccl::ncclCommGetAsyncError(comm, &mut async_err) };
234        if result != NcclResult::Success {
235            tracing::error!(
236                "ncclCommGetAsyncError call itself failed: {result:?} \
237                 — marking unhealthy"
238            );
239            self.unhealthy.store(true, Ordering::Release);
240            return false;
241        }
242        if async_err != NcclResult::Success {
243            tracing::error!("NCCL async error detected: {async_err:?} — marking unhealthy");
244            self.unhealthy.store(true, Ordering::Release);
245            return false;
246        }
247        true
248    }
249
250    /// 2026-09-26: Abort the communicator and bootstrap a new one on
251    /// `master_port + 1`, the same exchange as `new`, so every rank must call
252    /// it. Returns at once when the backend is not marked unhealthy. The
253    /// receive buffer is registered again; other registrations are dropped.
254    fn reconnect_inner(&self) -> Result<()> {
255        let mut comm_guard = self.comm.lock();
256
257        // 2026-09-26: Another thread may have reconnected while this one
258        // waited for the lock.
259        if !self.unhealthy.load(Ordering::Acquire) {
260            tracing::info!("NCCL communicator already recovered by another thread");
261            return Ok(());
262        }
263
264        let old_comm = *comm_guard;
265        let attempt = self.reconnect_count.load(Ordering::Relaxed) + 1;
266        tracing::warn!(
267            "NCCL reconnect: aborting old communicator \
268             (rank={}, reconnect #{})",
269            self.rank,
270            attempt,
271        );
272
273        if !old_comm.is_null() {
274            let result = unsafe { nccl::ncclCommAbort(old_comm) };
275            if result != NcclResult::Success {
276                tracing::warn!("ncclCommAbort returned {result:?} (proceeding anyway)");
277            }
278        }
279
280        // 2026-09-26: `master_port + 1`, so the bind does not collide with the
281        // first bootstrap's listener if it has not closed.
282        let reconnect_port = self.master_port.wrapping_add(1);
283        let unique_id = if self.rank == 0 {
284            let id = Self::generate_unique_id()?;
285            Self::distribute_id(&id, &self.master_addr, reconnect_port, self.world_size)?;
286            id
287        } else {
288            Self::receive_id(&self.master_addr, reconnect_port)?
289        };
290
291        let mut new_comm: NcclComm = ptr::null_mut();
292        let result = unsafe {
293            nccl::ncclCommInitRank(
294                &mut new_comm,
295                self.world_size as i32,
296                unique_id,
297                self.rank as i32,
298            )
299        };
300        nccl::check_nccl(result, "ncclCommInitRank (reconnect)")?;
301
302        if self.world_size == 2 && self.recv_buffer != 0 {
303            let mut handle: *mut c_void = ptr::null_mut();
304            let result = unsafe {
305                nccl::ncclCommRegister(
306                    new_comm,
307                    self.recv_buffer as *mut c_void,
308                    self.recv_capacity,
309                    &mut handle,
310                )
311            };
312            if result != NcclResult::Success {
313                tracing::warn!(
314                    "ncclCommRegister for recv_buffer after reconnect \
315                     failed: {result:?}"
316                );
317            } else {
318                tracing::info!("Re-registered recv_buffer after reconnect");
319            }
320        }
321
322        // 2026-09-26: The old handles belong to the aborted communicator, and
323        // only handles were stored, not (ptr, size), so they cannot be
324        // registered again.
325        let mut handles = self.registered_handles.lock();
326        if !handles.is_empty() {
327            tracing::warn!(
328                "Clearing {} stale registered buffer handles after reconnect",
329                handles.len()
330            );
331            handles.clear();
332        }
333        drop(handles);
334
335        *comm_guard = new_comm;
336        self.unhealthy.store(false, Ordering::Release);
337        self.reconnect_count.fetch_add(1, Ordering::Relaxed);
338
339        tracing::info!(
340            "NCCL reconnect successful (rank={}, total reconnects={})",
341            self.rank,
342            self.reconnect_count.load(Ordering::Relaxed),
343        );
344
345        Ok(())
346    }
347
348    /// 2026-09-26: 2-rank all-reduce: refuse a payload larger than
349    /// `recv_capacity`, then a grouped `ncclSend`/`ncclRecv` with the partner
350    /// rank, then `ptr[i] += recv_buffer[i]` with the BF16 add kernel on
351    /// `stream`. Errors when the add kernel is not set.
352    fn all_reduce_2rank(&self, ptr: u64, bytes: usize, stream: u64) -> Result<()> {
353        ensure_payload_fits(bytes, self.recv_capacity, self.rank, self.world_size)?;
354
355        // 2026-09-26: Nothing to reduce, and a zero-block launch is invalid.
356        // Both ranks skip the send/recv, provided both pass the same `bytes`.
357        if bytes == 0 {
358            return Ok(());
359        }
360
361        let count = bytes / ALL_REDUCE_DTYPE_BYTES;
362        let partner = (1 - self.rank) as i32;
363        let comm = *self.comm.lock();
364
365        let result = unsafe { nccl::ncclGroupStart() };
366        nccl::check_nccl(result, "ncclGroupStart")?;
367
368        let result = unsafe {
369            nccl::ncclSend(
370                ptr as *const c_void,
371                count,
372                NcclDataType::Bfloat16,
373                partner,
374                comm,
375                stream,
376            )
377        };
378        nccl::check_nccl(result, "ncclSend")?;
379
380        let result = unsafe {
381            nccl::ncclRecv(
382                self.recv_buffer as *mut c_void,
383                count,
384                NcclDataType::Bfloat16,
385                partner,
386                comm,
387                stream,
388            )
389        };
390        nccl::check_nccl(result, "ncclRecv")?;
391
392        let result = unsafe { nccl::ncclGroupEnd() };
393        nccl::check_nccl(result, "ncclGroupEnd")?;
394
395        self.check_async_error(comm);
396
397        let kernel = self.add_kernel.load(Ordering::Relaxed);
398        if kernel != 0 {
399            let threads: u32 = 256;
400            let blocks: u32 = (count as u32).div_ceil(threads);
401            let mut p_dst = ptr;
402            let mut p_src = self.recv_buffer;
403            let mut p_n = count as i32;
404            let mut params: [*mut c_void; 3] = [
405                &mut p_dst as *mut u64 as *mut c_void,
406                &mut p_src as *mut u64 as *mut c_void,
407                &mut p_n as *mut i32 as *mut c_void,
408            ];
409            let status = unsafe {
410                cuLaunchKernel(
411                    kernel,
412                    blocks,
413                    1,
414                    1,
415                    threads,
416                    1,
417                    1,
418                    0,
419                    stream,
420                    params.as_mut_ptr(),
421                    ptr::null_mut(),
422                )
423            };
424            if status != 0 {
425                anyhow::bail!("cuLaunchKernel (bf16_add_inplace) failed: status {status}");
426            }
427        } else {
428            anyhow::bail!("bf16_add_inplace kernel not set — call set_add_kernel() first");
429        }
430
431        Ok(())
432    }
433
434    fn generate_unique_id() -> Result<NcclUniqueId> {
435        let mut id = NcclUniqueId {
436            internal: [0u8; 128],
437        };
438        let result = unsafe { nccl::ncclGetUniqueId(&mut id) };
439        nccl::check_nccl(result, "ncclGetUniqueId")?;
440        Ok(id)
441    }
442
443    /// 2026-09-26: Rank 0: bind `0.0.0.0:port`, accept `world_size - 1`
444    /// connections and send each the unique id. `addr` is unused.
445    fn distribute_id(id: &NcclUniqueId, addr: &str, port: u16, world_size: usize) -> Result<()> {
446        let bind_addr = format!("0.0.0.0:{port}");
447        let listener = TcpListener::bind(&bind_addr)
448            .with_context(|| format!("Rank 0: failed to bind {bind_addr}"))?;
449        tracing::info!(
450            "Rank 0: waiting for {} worker(s) on {}",
451            world_size - 1,
452            bind_addr
453        );
454
455        for i in 0..(world_size - 1) {
456            let (mut stream, peer_addr) = listener.accept().context("Rank 0: accept failed")?;
457            stream
458                .write_all(&id.internal)
459                .context("Rank 0: failed to send unique ID")?;
460            tracing::info!("Rank 0: sent unique ID to worker {} ({})", i + 1, peer_addr);
461        }
462        let _ = addr;
463        Ok(())
464    }
465
466    /// 2026-09-26: Other ranks: connect to rank 0, retrying, and read the
467    /// unique id.
468    fn receive_id(addr: &str, port: u16) -> Result<NcclUniqueId> {
469        let target = format!("{addr}:{port}");
470        tracing::info!("Rank N: connecting to master at {target}");
471
472        // 2026-09-26: Serve bootstraps after loading weights, so rank 0 may
473        // open the port minutes after a worker starts connecting; the worker
474        // retries once a second for about ten minutes.
475        const MAX_ATTEMPTS: u32 = 600;
476        let mut stream = None;
477        for attempt in 0..MAX_ATTEMPTS {
478            match TcpStream::connect(&target) {
479                Ok(s) => {
480                    stream = Some(s);
481                    break;
482                }
483                Err(e) => {
484                    if attempt + 1 < MAX_ATTEMPTS {
485                        if attempt < 30 || attempt.is_multiple_of(30) {
486                            tracing::info!(
487                                "Connect attempt {}/{MAX_ATTEMPTS}: {e}, retrying in 1s (rank 0 may be loading weights)",
488                                attempt + 1
489                            );
490                        }
491                        std::thread::sleep(std::time::Duration::from_secs(1));
492                    } else {
493                        return Err(e).with_context(|| {
494                            format!(
495                                "Failed to connect to master at {target} \
496                                 after {MAX_ATTEMPTS} attempts (~{} minutes)",
497                                MAX_ATTEMPTS / 60
498                            )
499                        });
500                    }
501                }
502            }
503        }
504
505        let mut id = NcclUniqueId {
506            internal: [0u8; 128],
507        };
508        stream
509            .unwrap()
510            .read_exact(&mut id.internal)
511            .context("Failed to receive unique ID from rank 0")?;
512        tracing::info!("Received NCCL unique ID from master");
513        Ok(id)
514    }
515}
516
517impl Drop for NcclBackend {
518    fn drop(&mut self) {
519        let comm = *self.comm.lock();
520        let mut handles = self.registered_handles.lock();
521        for handle in handles.drain(..) {
522            unsafe { nccl::ncclCommDeregister(comm, handle) };
523        }
524        drop(handles);
525        if self.recv_buffer != 0 {
526            unsafe { cuMemFree_v2(self.recv_buffer) };
527        }
528        nccl::destroy_event(self.compute_done_event);
529        nccl::destroy_event(self.comm_done_event);
530        nccl::destroy_stream(self.comm_stream);
531        if !comm.is_null() {
532            unsafe { nccl::ncclCommDestroy(comm) };
533        }
534    }
535}
536
537mod comm_impl;
metrale_comm/nccl_backend.rs

metrale_comm/
nccl_backend.rs
