xtrain/crates/xtrain-distributed/src/lib.rs

//! Distributed data-parallel (DDP) primitives for xtrain (Phase T8).
//!
//! Launch model: **one OS thread per GPU** (same as xserv-distributed). Each
//! rank thread binds its device, builds its own model (xtrain's `Var` graph is
//! `Rc`-based and not `Send`, so it must be constructed thread-locally — only the
//! `UniqueId` and scalar config cross the thread boundary), processes a disjoint
//! shard of the global batch, then AllReduces every parameter's `.grad()` device
//! buffer in place, averages by world size, and runs its own `GpuAdamW.step`.
//! Identical init + identical optimizer state across ranks keeps the parameters
//! consistent without ever re-syncing the weights.
//!
//! NCCL is issued on the legacy null stream — every xtrain kernel launches on the
//! null stream (`std::ptr::null_mut()`), so the AllReduce stays correctly ordered
//! after the producing backward kernels and before the consuming optimizer step,
//! with no extra synchronization.

#![cfg(not(no_cuda))]

pub mod ddp;
pub mod ffi;

pub use ddp::{DdpConfig, build_model, launch, train_rank};

use std::ffi::c_void;

use ffi::{NcclComm, NcclUniqueId};
use xtrain_autodiff::tape::Var;
use xtrain_cuda::device;

pub use ffi::NcclUniqueId as UniqueId;

/// Generate a unique id on one rank (rank 0) and share the raw bytes to every
/// other rank out-of-band — across threads it is just a `Copy` struct moved into
/// each rank closure; across processes it would be written to a file/env.
pub fn get_unique_id() -> NcclUniqueId {
    let mut id = NcclUniqueId::default();
    ffi::check(unsafe { ffi::ncclGetUniqueId(&mut id) }, "ncclGetUniqueId");
    id
}

/// Per-rank data-parallel context: the NCCL communicator plus this rank's
/// identity. AllReduce is in-place on the null stream.
pub struct DdpContext {
    pub rank: usize,
    pub world: usize,
    pub device: u32,
    comm: NcclComm,
}

// The communicator is owned by exactly one rank thread.
unsafe impl Send for DdpContext {}

impl DdpContext {
    /// Initialize this rank. Must run on the thread that will own this rank's GPU
    /// work; binds the thread to `device` first. All ranks call this concurrently
    /// with the same `id` and `world` — the group wrapper lets the concurrent
    /// inits rendezvous without deadlock.
    pub fn init(rank: usize, world: usize, id: NcclUniqueId, device: u32) -> Self {
        device::set_device(device).expect("set_device");
        let mut comm: NcclComm = std::ptr::null_mut();
        ffi::check(unsafe { ffi::ncclGroupStart() }, "ncclGroupStart(init)");
        ffi::check(
            unsafe { ffi::ncclCommInitRank(&mut comm, world as i32, id, rank as i32) },
            "ncclCommInitRank",
        );
        ffi::check(unsafe { ffi::ncclGroupEnd() }, "ncclGroupEnd(init)");
        Self {
            rank,
            world,
            device,
            comm,
        }
    }

    /// In-place AllReduce(sum) over `count` F32 elements at a raw device pointer,
    /// issued on the null stream (so it orders with this rank's kernels). The
    /// reduction is asynchronous; a later sync (the caller's, or the next null-
    /// stream kernel) completes it.
    ///
    /// # Safety
    /// `ptr` must point to at least `count` valid F32 device elements on this
    /// rank's device. The reduction is in-place (send == recv).
    pub fn all_reduce_sum_f32_ptr(&self, ptr: *mut c_void, count: usize) {
        if self.world == 1 {
            return; // nothing to reduce
        }
        ffi::check(
            unsafe {
                ffi::ncclAllReduce(
                    ptr as *const c_void,
                    ptr,
                    count,
                    ffi::NCCL_FLOAT32,
                    ffi::NCCL_SUM,
                    self.comm,
                    std::ptr::null_mut(),
                )
            },
            "ncclAllReduce",
        );
    }

    /// AllReduce every parameter's `.grad()` across ranks and divide by `world`,
    /// the one collective DDP needs per step.
    ///
    /// Each rank ran forward+backward on its own shard of `b` sequences, so
    /// `.grad()` holds the SUM over that shard (the tape's fan-out rule). After
    /// `AllReduce(sum)` every rank holds `Σ_global` (the sum over all `world·b`
    /// sequences); dividing by `world` leaves `Σ_global / world`. The DDP train
    /// loop's clip pass then applies the remaining `1/b` (`pre_scale = 1/b_local`),
    /// giving `Σ_global / (world·b) = Σ_global / B_global` — bit-for-bit the same
    /// mean gradient the single-GPU loop computes from a batch of `B_global`.
    /// Params without a grad are skipped.
    ///
    /// A single-process group barrier is unnecessary: the all-reduces serialize
    /// on the comm, and the in-place scale runs on the same null stream after.
    pub fn all_reduce_average_grads(&self, params: &[Var]) {
        if self.world == 1 {
            return;
        }
        // 1. Sum every grad across ranks (in place, on the null stream).
        for p in params {
            if let Some(g) = p.grad() {
                let n = g.numel();
                self.all_reduce_sum_f32_ptr(g.data_ptr() as *mut c_void, n);
            }
        }
        // 2. Average: scale each summed grad by 1/world (null-stream kernel,
        //    ordered after the AllReduce that produced it).
        let inv_world = 1.0 / self.world as f32;
        for p in params {
            if let Some(g) = p.grad() {
                unsafe {
                    xtrain_cuda::ffi::launch_scale_inplace_f32(
                        g.data_ptr() as *mut f32,
                        inv_world,
                        g.numel() as i32,
                        std::ptr::null_mut(),
                    );
                }
            }
        }
        device::synchronize().expect("grad all-reduce sync failed");
    }
}

impl Drop for DdpContext {
    fn drop(&mut self) {
        if !self.comm.is_null() {
            unsafe { ffi::ncclCommDestroy(self.comm) };
        }
    }
}