docs: Phase T8 — distributed data parallel

Design doc for the NCCL DDP path: comm bootstrap (rank-0 UniqueId + grouped
CommInitRank), thread-per-GPU launch model (Var is !Send), all-reduce-then-
local-step scheme (in-place fp32 AllReduce on .grad() + /world, each rank steps
its own GpuAdamW), why params stay consistent (NCCL bit-identical reduce + same
init/state), batch sharding math vs single-GPU, verification plan + scaling
table. Lists TP/PP/ZeRO/bf16-comm as out-of-scope follow-ups.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Cargo.lock

@@ -205,6 +205,18 @@ dependencies = [
  "cc",
 ]
 
+[[package]]
+name = "xtrain-distributed"
+version = "0.1.0"
+dependencies = [
+ "xtrain-autodiff",
+ "xtrain-cuda",
+ "xtrain-model",
+ "xtrain-optim",
+ "xtrain-tensor",
+ "xtrain-train",
+]
+
 [[package]]
 name = "xtrain-model"
 version = "0.1.0"

Cargo.toml

@@ -7,6 +7,7 @@ members = [
     "crates/xtrain-model",
     "crates/xtrain-optim",
     "crates/xtrain-train",
+    "crates/xtrain-distributed",
 ]
 
 [workspace.package]

crates/xtrain-distributed/Cargo.toml

@@ -0,0 +1,13 @@
+[package]
+name = "xtrain-distributed"
+version.workspace = true
+edition.workspace = true
+license.workspace = true
+
+[dependencies]
+xtrain-cuda = { path = "../xtrain-cuda" }
+xtrain-tensor = { path = "../xtrain-tensor" }
+xtrain-autodiff = { path = "../xtrain-autodiff" }
+xtrain-model = { path = "../xtrain-model" }
+xtrain-optim = { path = "../xtrain-optim" }
+xtrain-train = { path = "../xtrain-train" }

crates/xtrain-distributed/build.rs

@@ -0,0 +1,33 @@
+use std::env;
+use std::path::Path;
+use std::process::Command;
+
+// Mirror the per-crate convention (see xtrain-cuda/build.rs): with no nvcc/GPU
+// locally, emit `no_cuda` so the NCCL FFI + DDP code compiles (but is not linked
+// or run). On dash5, link NCCL exactly like xserv-distributed's build.rs.
+fn main() {
+    println!("cargo:rustc-check-cfg=cfg(no_cuda)");
+
+    let cuda_path = env::var("CUDA_HOME")
+        .or_else(|_| env::var("CUDA_PATH"))
+        .unwrap_or_else(|_| "/usr/local/cuda".to_string());
+
+    if !nvcc_available(&cuda_path) {
+        println!("cargo:warning=nvcc not found — skipping NCCL link (host-only build).");
+        println!("cargo:rustc-cfg=no_cuda");
+        return;
+    }
+
+    println!("cargo:rustc-link-search=native={cuda_path}/lib64");
+    // NCCL is installed as a system library on dash5.
+    println!("cargo:rustc-link-search=native=/usr/lib/x86_64-linux-gnu");
+    println!("cargo:rustc-link-lib=dylib=nccl");
+    println!("cargo:rustc-link-lib=dylib=cudart");
+}
+
+fn nvcc_available(cuda_path: &str) -> bool {
+    if Command::new("nvcc").arg("--version").output().is_ok() {
+        return true;
+    }
+    Path::new(&format!("{cuda_path}/bin/nvcc")).exists()
+}

crates/xtrain-distributed/src/bin/train_ddp.rs

@@ -0,0 +1,101 @@
+//! Multi-rank DDP training launcher (Phase T8): spawn one thread per GPU, NCCL
+//! all-reduce the gradients each step, and train the tiny transformer on
+//! TinyStories. Doubles as the throughput driver — run it with 1/2/4 GPUs and
+//! read the global tok/s line.
+//!
+//! Run on dash5 (pick idle GPUs — dash5 is shared):
+//!   export PATH=/usr/local/cuda/bin:/opt/wjh/.cargo/bin:$PATH
+//!   CUDA_VISIBLE_DEVICES=0,1 cargo run -p xtrain-distributed --release \
+//!       --bin train_ddp -- 100 64 16
+//! Args: [steps] [seq_len] [global_batch] [tokenizer.json] [corpus.txt]
+//! The launcher uses every GPU visible to it (CUDA_VISIBLE_DEVICES selects them),
+//! so the rank devices are always 0..N within the visible set.
+
+#[cfg(no_cuda)]
+fn main() {
+    eprintln!("train_ddp: built without CUDA (no_cuda); run on a GPU host (dash5).");
+}
+
+#[cfg(not(no_cuda))]
+fn main() {
+    use std::path::PathBuf;
+    use xtrain_cuda::device;
+    use xtrain_distributed::{DdpConfig, build_model, launch};
+    use xtrain_model::Config;
+    use xtrain_train::data::Corpus;
+    use xtrain_train::schedule::LrSchedule;
+
+    let args: Vec<String> = std::env::args().collect();
+    let steps: usize = args.get(1).and_then(|s| s.parse().ok()).unwrap_or(100);
+    let seq_len: usize = args.get(2).and_then(|s| s.parse().ok()).unwrap_or(64);
+    let batch: usize = args.get(3).and_then(|s| s.parse().ok()).unwrap_or(16);
+    let tok_path = args
+        .get(4)
+        .map(PathBuf::from)
+        .unwrap_or_else(|| PathBuf::from("/opt/wjh/models/gpt2/tokenizer.json"));
+    let corpus_path = args
+        .get(5)
+        .map(PathBuf::from)
+        .unwrap_or_else(|| PathBuf::from("data/tinystories-valid-3mb.txt"));
+
+    // Use every visible GPU as a rank (CUDA_VISIBLE_DEVICES selects the set;
+    // device ordinals are 0..count within it).
+    let count = device::device_count().expect("device_count") as u32;
+    assert!(count > 0, "no CUDA device visible");
+    let devices: Vec<u32> = (0..count).collect();
+    assert_eq!(
+        batch % devices.len(),
+        0,
+        "global batch {batch} not divisible by world {}",
+        devices.len()
+    );
+
+    println!(
+        "DDP: world={} devices={:?} | steps={steps} seq={seq_len} global_batch={batch}",
+        devices.len(),
+        devices
+    );
+
+    let corpus = Corpus::load(&tok_path, &corpus_path);
+    println!(
+        "corpus: {} tokens, vocab {}",
+        corpus.len(),
+        corpus.vocab_size
+    );
+
+    let mut cfg = Config::tiny();
+    cfg.vocab = corpus.vocab_size;
+    cfg.n_layers = 4;
+    println!(
+        "model: dim {} layers {} heads {} ffn {} → {} params",
+        cfg.dim,
+        cfg.n_layers,
+        cfg.n_heads,
+        cfg.ffn_hidden,
+        cfg.num_params()
+    );
+
+    let dcfg = DdpConfig {
+        seq_len,
+        batch_size: batch,
+        steps,
+        schedule: LrSchedule {
+            max_lr: 3e-3,
+            min_lr: 3e-4,
+            warmup: (steps / 20).max(5),
+            total: steps,
+        },
+        weight_decay: 0.1,
+        max_grad_norm: 1.0,
+        log_every: 10,
+        seed: 42,
+    };
+
+    let traces = launch(&devices, &corpus, &dcfg, move |device| {
+        build_model(cfg, device)
+    });
+    let trace = &traces[0];
+    let start = trace.first().copied().unwrap_or(0.0);
+    let end = trace.last().copied().unwrap_or(0.0);
+    println!("loss: start {start:.4} → end {end:.4}");
+}

crates/xtrain-distributed/src/ddp.rs

@@ -0,0 +1,204 @@
+//! The DDP training step + a single-process, thread-per-GPU launcher (Phase T8).
+//!
+//! Each rank owns one GPU and one thread. Per step it processes a DISJOINT shard
+//! of the global batch, all-reduce-averages the gradients, then runs its own
+//! `GpuAdamW.step`. Identical init + identical optimizer state across ranks keep
+//! the parameters consistent — verified by the cross-rank param-identity check in
+//! the tests.
+//!
+//! Sampling matches single-GPU bit-for-bit: every rank advances the SAME RNG and
+//! draws all `B_global` sequences of a step, but only runs forward+backward on
+//! the ones assigned to it (`global index % world == rank`). The union over ranks
+//! is exactly the single-GPU batch in the same order, so the all-reduced grad sum
+//! equals the single-GPU summed grad.
+
+use std::thread;
+use std::time::Instant;
+
+use xtrain_autodiff::tape::Var;
+use xtrain_model::{Config, TinyTransformer, ids_tensor};
+use xtrain_optim::GpuAdamW;
+use xtrain_tensor::Device;
+use xtrain_train::clip::clip_grad_norm_gpu;
+use xtrain_train::data::Corpus;
+use xtrain_train::schedule::LrSchedule;
+
+use crate::{DdpContext, get_unique_id};
+
+/// Per-rank DDP training config. `batch_size` is the GLOBAL batch (split across
+/// ranks); the rest mirror `xtrain_train::TrainConfig`.
+#[derive(Clone)]
+pub struct DdpConfig {
+    pub seq_len: usize,
+    /// Global batch size; must be divisible by the world size.
+    pub batch_size: usize,
+    pub steps: usize,
+    pub schedule: LrSchedule,
+    pub weight_decay: f32,
+    pub max_grad_norm: f32,
+    pub log_every: usize,
+    pub seed: u64,
+}
+
+/// Run `cfg.steps` DDP steps on this rank's `model`/`corpus`, using `ctx` for the
+/// gradient all-reduce. Returns this rank's per-step mean-loss trace (the mean
+/// over the GLOBAL batch — every rank computes the same value because losses are
+/// all-reduced alongside the grads). The optimizer step is identical on every
+/// rank, so the parameters stay in lockstep.
+pub fn train_rank(
+    ctx: &DdpContext,
+    model: &TinyTransformer,
+    device: Device,
+    corpus: &Corpus,
+    cfg: &DdpConfig,
+) -> Vec<f32> {
+    assert_eq!(
+        cfg.batch_size % ctx.world,
+        0,
+        "global batch {} not divisible by world {}",
+        cfg.batch_size,
+        ctx.world
+    );
+    let params = model.params();
+    let mut opt = GpuAdamW::new(cfg.weight_decay);
+    let mut rng = cfg.seed;
+    let mut losses = Vec::with_capacity(cfg.steps);
+    // Each rank reaches the global batch mean as (Σ_global / world) · (1/b_local),
+    // where b_local = batch_size / world (see DdpContext::all_reduce_average_grads).
+    let batch_local = cfg.batch_size / ctx.world;
+    let inv_batch_local = 1.0 / batch_local as f32;
+    let start = Instant::now();
+    let mut tokens_seen: u64 = 0;
+
+    for step in 0..cfg.steps {
+        let lr = cfg.schedule.lr(step);
+
+        // Draw the whole global batch from the shared RNG (same on every rank);
+        // run forward+backward only on this rank's shard. The tape SUMs the
+        // shard's grads; the union of shards == the single-GPU batch.
+        let mut local_loss_sum = 0.0f32;
+        for i in 0..cfg.batch_size {
+            let (input, target) = corpus.sample(cfg.seq_len, &mut rng);
+            if i % ctx.world != ctx.rank {
+                continue; // not this rank's sequence
+            }
+            let ids = ids_tensor(&input, device);
+            let targets = ids_tensor(&target, device);
+            let loss = model.loss(&ids, &targets);
+            local_loss_sum += read_scalar(&loss);
+            loss.backward();
+            tokens_seen += cfg.seq_len as u64;
+        }
+
+        // AllReduce(sum) + /world the grads → every rank holds Σ_global/world.
+        ctx.all_reduce_average_grads(&params);
+        // The reported loss is the global mean: average local sums across ranks.
+        let step_loss = all_reduce_loss(ctx, local_loss_sum) / cfg.batch_size as f32;
+        losses.push(step_loss);
+
+        // clip pre_scale = 1/b_local finishes the average to Σ_global/B_global,
+        // identical to the single-GPU clip(pre_scale = 1/B_global).
+        let gnorm = clip_grad_norm_gpu(&params, cfg.max_grad_norm, inv_batch_local);
+        opt.step(lr, &params);
+        for p in &params {
+            p.zero_grad();
+        }
+
+        if ctx.rank == 0 && (step % cfg.log_every == 0 || step == cfg.steps - 1) {
+            let elapsed = start.elapsed().as_secs_f32();
+            // Global tok/s = per-rank tok/s × world (each rank does 1/world of it).
+            let tps = (tokens_seen as f32 / elapsed.max(1e-6)) * ctx.world as f32;
+            println!(
+                "[rank0] step {step:5}/{}: loss {step_loss:.4} lr {lr:.2e} gnorm {gnorm:.3} \
+                 ({tps:.0} tok/s global, {} ranks)",
+                cfg.steps, ctx.world
+            );
+        }
+    }
+    losses
+}
+
+/// Spawn `world` rank threads (one per GPU in `devices`), init NCCL, build an
+/// identical model per rank via `make_model`, and run `train_rank`. Returns each
+/// rank's loss trace (all identical). The launcher owns the thread-per-GPU model:
+/// rank 0 mints the `UniqueId`, every thread `cudaSetDevice`s its GPU, builds its
+/// `Var` graph locally (the graph is `!Send`), and joins at the end.
+///
+/// `make_model(device)` must be deterministic — same params on every rank — for
+/// the parameters to stay consistent.
+pub fn launch<F>(devices: &[u32], corpus: &Corpus, cfg: &DdpConfig, make_model: F) -> Vec<Vec<f32>>
+where
+    F: Fn(Device) -> TinyTransformer + Send + Sync,
+{
+    let world = devices.len();
+    let id = get_unique_id();
+
+    thread::scope(|s| {
+        let handles: Vec<_> = devices
+            .iter()
+            .enumerate()
+            .map(|(rank, &dev)| {
+                let make_model = &make_model;
+                let cfg = cfg.clone();
+                s.spawn(move || {
+                    let ctx = DdpContext::init(rank, world, id, dev);
+                    let device = Device::Cuda(dev);
+                    let model = make_model(device);
+                    train_rank(&ctx, &model, device, corpus, &cfg)
+                })
+            })
+            .collect();
+        handles.into_iter().map(|h| h.join().unwrap()).collect()
+    })
+}
+
+/// AllReduce(sum) a single host scalar across ranks by round-tripping it through a
+/// one-element device buffer. Used only for the logged/returned loss, so the cost
+/// (one tiny collective per step) is negligible. Returns the summed value.
+fn all_reduce_loss(ctx: &DdpContext, local: f32) -> f32 {
+    use xtrain_tensor::{DType, Tensor};
+    if ctx.world == 1 {
+        return local;
+    }
+    let device = Device::Cuda(ctx.device);
+    let t = Tensor::from_slice(&[local], &[1]).to_device(device);
+    ctx.all_reduce_sum_f32_ptr(t.data_ptr() as *mut std::ffi::c_void, 1);
+    xtrain_cuda::device::synchronize().expect("loss all-reduce sync");
+    t.to_device(Device::Cpu).as_slice::<f32>()[0]
+}
+
+fn read_scalar(v: &Var) -> f32 {
+    v.value().to_device(Device::Cpu).as_slice::<f32>()[0]
+}
+
+/// Build a `TinyTransformer` on `device` with the SAME deterministic init the
+/// single-GPU `bin/train` uses (LCG fill, gammas ~1). Used by both the launcher
+/// and the correctness test so every rank — and the single-GPU baseline — start
+/// from bit-identical parameters. `cfg` must be identical on every call.
+pub fn build_model(cfg: Config, device: Device) -> TinyTransformer {
+    let mut seed = 1u64;
+    TinyTransformer::new(cfg, device, |shape| {
+        seed = seed.wrapping_add(1);
+        let n: usize = shape.iter().product();
+        if shape.len() == 1 {
+            fill(n, seed, 0.02).iter().map(|v| v + 1.0).collect()
+        } else {
+            fill(n, seed, 0.04)
+        }
+    })
+}
+
+// Deterministic LCG fill in [-scale, scale) — same scheme as bin/train's `fill`.
+fn fill(n: usize, seed: u64, scale: f32) -> Vec<f32> {
+    let mut state = seed
+        .wrapping_mul(2862933555777941757)
+        .wrapping_add(3037000493);
+    (0..n)
+        .map(|_| {
+            state = state
+                .wrapping_mul(6364136223846793005)
+                .wrapping_add(1442695040888963407);
+            (((state >> 33) as f32 / (1u64 << 31) as f32) - 0.5) * 2.0 * scale
+        })
+        .collect()
+}

crates/xtrain-distributed/src/ffi.rs

@@ -0,0 +1,76 @@
+//! Minimal NCCL FFI bindings (hand-written, like the CUDA bindings in
+//! xtrain-cuda). Only the collectives data-parallel training needs:
+//! unique-id creation, communicator init/destroy, and AllReduce. Mirrors
+//! xserv-distributed's FFI.
+
+use std::ffi::c_void;
+use std::os::raw::c_char;
+use xtrain_cuda::ffi::CudaStream;
+
+/// Opaque NCCL communicator handle (`ncclComm_t`).
+pub type NcclComm = *mut c_void;
+
+/// `ncclUniqueId` is a 128-byte opaque blob shared from rank 0 to every rank.
+#[repr(C)]
+#[derive(Clone, Copy)]
+pub struct NcclUniqueId {
+    pub internal: [c_char; 128],
+}
+
+impl Default for NcclUniqueId {
+    fn default() -> Self {
+        Self { internal: [0; 128] }
+    }
+}
+
+// ncclDataType_t (subset) — DDP all-reduces fp32 gradients.
+pub const NCCL_FLOAT32: i32 = 7;
+
+// ncclRedOp_t
+pub const NCCL_SUM: i32 = 0;
+
+// ncclResult_t
+pub const NCCL_SUCCESS: i32 = 0;
+
+unsafe extern "C" {
+    pub fn ncclGetUniqueId(uid: *mut NcclUniqueId) -> i32;
+    // ncclUniqueId is passed BY VALUE (a 128-byte struct) per the NCCL ABI.
+    pub fn ncclCommInitRank(
+        comm: *mut NcclComm,
+        nranks: i32,
+        commid: NcclUniqueId,
+        rank: i32,
+    ) -> i32;
+    pub fn ncclCommDestroy(comm: NcclComm) -> i32;
+    pub fn ncclAllReduce(
+        sendbuff: *const c_void,
+        recvbuff: *mut c_void,
+        count: usize,
+        datatype: i32,
+        op: i32,
+        comm: NcclComm,
+        stream: CudaStream,
+    ) -> i32;
+    pub fn ncclGroupStart() -> i32;
+    pub fn ncclGroupEnd() -> i32;
+    pub fn ncclGetErrorString(result: i32) -> *const c_char;
+}
+
+pub fn err_string(result: i32) -> String {
+    unsafe {
+        let p = ncclGetErrorString(result);
+        if p.is_null() {
+            return format!("nccl error {result}");
+        }
+        std::ffi::CStr::from_ptr(p).to_string_lossy().into_owned()
+    }
+}
+
+pub fn check(result: i32, what: &str) {
+    assert_eq!(
+        result,
+        NCCL_SUCCESS,
+        "{what} failed: {}",
+        err_string(result)
+    );
+}

crates/xtrain-distributed/src/lib.rs

@@ -0,0 +1,153 @@
+//! 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) };
+        }
+    }
+}

crates/xtrain-distributed/tests/ddp_correctness.rs

@@ -0,0 +1,236 @@
+//! DDP acceptance (Phase T8). Gated to a GPU host; skips when fewer than 2 GPUs.
+//!
+//! 1. **Correctness**: K steps single-GPU (world=1, global batch B) vs 2-rank DDP
+//!    (B/2 of the SAME data in the same order each) → loss trajectories match
+//!    within tight fp tolerance (it's just gradient averaging), and the two
+//!    ranks' parameters are identical after the run.
+//! 2. **Throughput**: 1 / 2 / 4 GPU global tok/s on the SAME per-GPU workload →
+//!    near-linear scaling. Prints the table (run with `--nocapture`).
+
+#![cfg(not(no_cuda))]
+
+use std::time::Instant;
+
+use xtrain_cuda::device;
+use xtrain_distributed::{DdpConfig, DdpContext, build_model, get_unique_id, launch, train_rank};
+use xtrain_model::{Config, ids_tensor};
+use xtrain_optim::GpuAdamW;
+use xtrain_tensor::Device;
+use xtrain_train::clip::clip_grad_norm_gpu;
+use xtrain_train::data::Corpus;
+use xtrain_train::schedule::LrSchedule;
+
+// A self-contained synthetic corpus so the test needs no tokenizer/data files.
+fn synth_corpus(vocab: usize, n_tokens: usize) -> Corpus {
+    let tokens: Vec<i32> = (0..n_tokens)
+        .map(|i| (i * 7 + 3) as i32 % vocab as i32)
+        .collect();
+    Corpus {
+        tokens,
+        vocab_size: vocab,
+    }
+}
+
+fn test_config(vocab: usize) -> Config {
+    let mut cfg = Config::tiny();
+    cfg.vocab = vocab;
+    cfg.n_layers = 2;
+    cfg
+}
+
+// Single-GPU baseline: the SAME loop as the DDP rank but world=1, so the global
+// batch is processed on one device. Returns (loss trace, final params on host).
+fn run_single_gpu(cfg: Config, corpus: &Corpus, dcfg: &DdpConfig) -> (Vec<f32>, Vec<Vec<f32>>) {
+    device::set_device(0).unwrap();
+    let device = Device::Cuda(0);
+    let model = build_model(cfg, device);
+    let params = model.params();
+    let mut opt = GpuAdamW::new(dcfg.weight_decay);
+    let mut rng = dcfg.seed;
+    let inv_batch = 1.0 / dcfg.batch_size as f32;
+    let mut losses = Vec::new();
+
+    for step in 0..dcfg.steps {
+        let lr = dcfg.schedule.lr(step);
+        let mut loss_sum = 0.0f32;
+        for _ in 0..dcfg.batch_size {
+            let (input, target) = corpus.sample(dcfg.seq_len, &mut rng);
+            let ids = ids_tensor(&input, device);
+            let targets = ids_tensor(&target, device);
+            let loss = model.loss(&ids, &targets);
+            loss_sum += loss.value().to_device(Device::Cpu).as_slice::<f32>()[0];
+            loss.backward();
+        }
+        losses.push(loss_sum * inv_batch);
+        clip_grad_norm_gpu(&params, dcfg.max_grad_norm, inv_batch);
+        opt.step(lr, &params);
+        for p in &params {
+            p.zero_grad();
+        }
+    }
+
+    let host = params
+        .iter()
+        .map(|p| p.value().to_device(Device::Cpu).as_slice::<f32>().to_vec())
+        .collect();
+    (losses, host)
+}
+
+#[test]
+fn ddp_matches_single_gpu_and_params_consistent() {
+    let world = 2usize;
+    if device::device_count().unwrap_or(0) < world as i32 {
+        eprintln!("skip: need >= {world} GPUs");
+        return;
+    }
+
+    let vocab = 64usize;
+    let cfg = test_config(vocab);
+    let corpus = synth_corpus(vocab, 4096);
+    let steps = 20usize;
+    let dcfg = DdpConfig {
+        seq_len: 32,
+        batch_size: 8, // global; 4 per rank with world=2
+        steps,
+        schedule: LrSchedule {
+            max_lr: 3e-3,
+            min_lr: 3e-4,
+            warmup: 3,
+            total: steps,
+        },
+        weight_decay: 0.1,
+        max_grad_norm: 1.0,
+        log_every: 1_000_000, // silence per-step logging in the test
+        seed: 7,
+    };
+
+    // Single-GPU baseline (world=1) over the global batch.
+    let (single_losses, single_params) = run_single_gpu(cfg, &corpus, &dcfg);
+
+    // 2-rank DDP over the SAME corpus/config; returns per-rank (losses, params).
+    let devices = [0u32, 1u32];
+    let id = get_unique_id();
+    let results: Vec<(Vec<f32>, Vec<Vec<f32>>)> = std::thread::scope(|s| {
+        let handles: Vec<_> = devices
+            .iter()
+            .enumerate()
+            .map(|(rank, &dev)| {
+                let dcfg = dcfg.clone();
+                let corpus = &corpus;
+                s.spawn(move || {
+                    let ctx = DdpContext::init(rank, world, id, dev);
+                    let device = Device::Cuda(dev);
+                    let model = build_model(cfg, device);
+                    let losses = train_rank(&ctx, &model, device, corpus, &dcfg);
+                    let host = model
+                        .params()
+                        .iter()
+                        .map(|p| p.value().to_device(Device::Cpu).as_slice::<f32>().to_vec())
+                        .collect::<Vec<_>>();
+                    (losses, host)
+                })
+            })
+            .collect();
+        handles.into_iter().map(|h| h.join().unwrap()).collect()
+    });
+
+    let (ddp_losses, ddp_p0) = &results[0];
+    let (_, ddp_p1) = &results[1];
+
+    // (a) DDP loss trajectory matches single-GPU within tight tolerance.
+    let mut max_rel = 0.0f32;
+    for (s, d) in single_losses.iter().zip(ddp_losses) {
+        let rel = (s - d).abs() / s.abs().max(1e-6);
+        max_rel = max_rel.max(rel);
+    }
+    println!(
+        "DDP vs single-GPU loss: single[last]={:.6} ddp[last]={:.6} max_rel={max_rel:.2e}",
+        single_losses.last().unwrap(),
+        ddp_losses.last().unwrap()
+    );
+    assert!(
+        max_rel < 1e-3,
+        "DDP loss trajectory diverged from single-GPU: max_rel {max_rel:.3e}"
+    );
+
+    // (b) Cross-rank parameter identity (same init + same averaged grad + same
+    //     optimizer state ⇒ identical params).
+    let mut max_pdiff = 0.0f32;
+    for (a, b) in ddp_p0.iter().zip(ddp_p1) {
+        for (x, y) in a.iter().zip(b) {
+            max_pdiff = max_pdiff.max((x - y).abs());
+        }
+    }
+    println!("cross-rank max |param diff| = {max_pdiff:.3e}");
+    assert_eq!(max_pdiff, 0.0, "ranks' params drifted apart");
+
+    // (c) DDP final params match single-GPU final params within fp tolerance.
+    let mut max_sdiff = 0.0f32;
+    for (a, b) in ddp_p0.iter().zip(&single_params) {
+        for (x, y) in a.iter().zip(b) {
+            max_sdiff = max_sdiff.max((x - y).abs() / y.abs().max(1e-6));
+        }
+    }
+    println!("DDP vs single-GPU max rel |param diff| = {max_sdiff:.3e}");
+    assert!(max_sdiff < 1e-3, "DDP params diverged from single-GPU");
+}
+
+#[test]
+fn ddp_throughput_scaling() {
+    let max_gpus = device::device_count().unwrap_or(0) as usize;
+    if max_gpus < 1 {
+        eprintln!("skip: no GPU");
+        return;
+    }
+    // Same PER-GPU workload at each world size (batch scales with world), so the
+    // per-rank cost is fixed and global tok/s should scale ~linearly.
+    let per_gpu_batch = 8usize;
+    let vocab = 256usize;
+    let cfg = test_config(vocab);
+    let corpus = synth_corpus(vocab, 8192);
+    let steps = 30usize;
+    let seq_len = 64usize;
+
+    let worlds: Vec<usize> = [1, 2, 4].into_iter().filter(|&w| w <= max_gpus).collect();
+    println!("\n=== DDP throughput scaling (per-GPU batch {per_gpu_batch}, seq {seq_len}) ===");
+    println!(
+        "{:>6} | {:>14} | {:>8}",
+        "GPUs", "tok/s (global)", "speedup"
+    );
+
+    let mut base = 0.0f64;
+    for &world in &worlds {
+        let devices: Vec<u32> = (0..world as u32).collect();
+        let dcfg = DdpConfig {
+            seq_len,
+            batch_size: per_gpu_batch * world,
+            steps,
+            schedule: LrSchedule {
+                max_lr: 1e-3,
+                min_lr: 1e-3,
+                warmup: 1,
+                total: steps,
+            },
+            weight_decay: 0.0,
+            max_grad_norm: 1.0,
+            log_every: 1_000_000,
+            seed: 1,
+        };
+        let total_tokens = (steps * dcfg.batch_size * seq_len) as f64;
+        let t = Instant::now();
+        let _ = launch(&devices, &corpus, &dcfg, move |device| {
+            build_model(cfg, device)
+        });
+        let secs = t.elapsed().as_secs_f64();
+        let tps = total_tokens / secs;
+        if world == 1 {
+            base = tps;
+        }
+        println!(
+            "{:>6} | {:>14.0} | {:>7.2}x",
+            world,
+            tps,
+            tps / base.max(1e-9)
+        );
+    }
+}

docs/07-distributed.md

@@ -0,0 +1,157 @@
+# Phase T8: Distributed Data Parallel (DDP) — Design Document
+
+## Goal
+
+T7 把单卡训练榨到 **~8500 tok/s**（梯度常驻 device，AdamW + grad-norm 都在 device 侧）。
+T8 的目标是**把训练 scale 到 dash5 的多张 RTX 5090**：用 **NCCL 数据并行**——每个 rank 跑 global batch
+的一个分片，反向后把每个参数的梯度在显存里 **all-reduce 求和、按 world size 取均值**，然后各 rank 用
+**自己**那份（已同步的）优化器跑 `GpuAdamW.step`。同 init + 同优化器超参/状态 ⇒ 参数永远逐位一致，
+无需再同步权重。
+
+验收（两条都要）：
+
+1. **正确性**：DDP 必须对住单卡。单卡跑 K 步 global batch B，DDP 用 2 rank 各跑同一批数据同序的 B/2 →
+   loss 轨迹与单卡在紧容差内吻合（纯梯度平均，理论上近乎一致）；且一步后断言两个 rank 的参数完全相同。
+2. **吞吐**：固定每卡 workload，测 1/2/4 卡的 tok/s → 近线性 scaling，给表。
+
+**本 Phase 范围内只做 DDP**。tensor/pipeline 并行、sharded optimizer 是可选 bonus，**不做**（model 太小，
+当前学习目标是把 DDP 这条最朴素也最核心的多卡链路吃透；切分式并行留待 model 放大）。导出回流 xserv 是 T9。
+
+## Module Layout
+
+```
+crates/xtrain-distributed/                  # 新 crate（镜像 xserv-distributed）
+├── build.rs                                # nvcc 探测 → no_cuda cfg；有 nvcc 时链 -lnccl -lcudart（同 xserv build.rs）
+├── Cargo.toml                              # 依赖 cuda/tensor/autodiff/model/optim/train
+├── src/ffi.rs                              # NCCL FFI：GetUniqueId/CommInitRank/AllReduce/CommDestroy/Group{Start,End}
+├── src/lib.rs                              # DdpContext（comm bootstrap + all_reduce_average_grads）+ get_unique_id
+├── src/ddp.rs                              # DDP 训练 step（train_rank）+ thread-per-GPU launcher + 确定性 build_model
+└── src/bin/train_ddp.rs                    # 多 rank 启动器 / 吞吐 driver（CUDA_VISIBLE_DEVICES 选卡）
+crates/xtrain-distributed/tests/
+└── ddp_correctness.rs                      # 2 卡 vs 单卡 loss 对拍 + 跨 rank 参数一致 + 1/2/4 卡吞吐表
+Cargo.toml                                  # workspace members += xtrain-distributed
+docs/07-distributed.md                      # 本文
+```
+
+整个 crate 用 `#![cfg(not(no_cuda))]` 在 crate 根门控：本地（无 nvcc）crate 编译为空，不引用任何 NCCL 符号，
+`cargo check` 直接过；dash5 上 `not(no_cuda)` 全量编译并链接 NCCL。bin 另带一个 `#[cfg(no_cuda)]` stub main。
+
+## Key Design Decisions
+
+### ① NCCL comm bootstrap（rank 0 发 UniqueId）
+
+NCCL 的握手：**rank 0 调 `ncclGetUniqueId`** 生成一个 128 字节不透明 id，**带外**分发给所有 rank，每个 rank
+用 `ncclCommInitRank(comm, world, id, rank)` 加入通信组。`ncclUniqueId` 按 NCCL ABI **按值传递**（128 字节 struct），
+FFI 里 `#[repr(C)] struct NcclUniqueId { internal: [c_char;128] }`，`ncclCommInitRank` 的 `commid` 参数直接传值。
+
+并发 init 必须用 `ncclGroupStart()/ncclGroupEnd()` 包住，否则多个线程同时 `CommInitRank` 会互相等待握手而死锁——
+group 让它们 rendezvous。这套与 xserv-distributed 的 `TpContext::init` 完全一致。
+
+### ② Launch model：thread-per-GPU（与 xserv 一致）
+
+**单进程、每卡一个 OS 线程**，线程启动先 `cudaSetDevice(dev)` 绑定本卡。选这个模型的硬约束：xtrain 的
+autograd `Var` 是 `Rc<RefCell<…>>`，**不是 `Send`**，整张计算图不能跨线程。所以**每个 rank 线程在闭包内本地
+`build_model` 自己的 `Var` 图**——跨线程边界的只有 `UniqueId`（`Copy`）、标量 config、和 `&Corpus`（只读共享，`Sync`）。
+
+`launch()` 用 `std::thread::scope`：rank 0 发 id → 每线程 `DdpContext::init` → 本地建模 → `train_rank` → join。
+`DdpContext` 持有 `comm`（裸指针），`unsafe impl Send` 因为它被恰好一个 rank 线程独占。
+
+为什么不用多进程（torchrun 式）：单机多卡，单进程多线程改动最小、无 IPC；UniqueId 跨线程就是 move 一个 `Copy`
+struct，跨进程才需要写文件/env 分发。多进程留待真正跨节点。
+
+### ③ all-reduce-then-local-step：梯度在显存里平均，各 rank 各自 step
+
+T7 起**梯度常驻 device**（`Var::grad()` 是 device tensor）。DDP 的通信因此极简——**直接对每个参数的 `.grad()`
+device buffer 做 in-place AllReduce**，零 host 往返：
+
+```rust
+// DdpContext::all_reduce_average_grads(params)
+for p in params { if let Some(g) = p.grad() {
+    ncclAllReduce(g.data_ptr(), g.data_ptr(), g.numel(), NCCL_FLOAT32, NCCL_SUM, comm, null_stream);
+}}
+for p in params { if let Some(g) = p.grad() {
+    launch_scale_inplace_f32(g.data_ptr(), 1.0/world, g.numel(), null_stream);  // 复用 T7 的 kernel
+}}
+```
+
+AllReduce 用 **fp32 + Sum**（梯度就是 fp32），发在 **null/legacy stream** 上——xtrain 每个 kernel 都在 null stream
+（`std::ptr::null_mut()`）launch，所以 AllReduce 自然排在产出梯度的 backward kernel 之后、消费它的 scale/optimizer
+kernel 之前，**无需额外 barrier**。平均直接复用 T7 的 `launch_scale_inplace_f32`（grad-clip 用的同一个 kernel）。
+
+之后**每个 rank 各自跑 `GpuAdamW.step`**——不再同步权重。
+
+### ④ 为什么参数保持一致
+
+三个充分条件：(a) **同 init**：所有 rank（含单卡 baseline）用同一个确定性 `build_model`（同 LCG 种子），起点逐位相同；
+(b) **同梯度**：NCCL AllReduce 对参与的所有 rank **返回逐位相同的归约结果**（这是 NCCL 的保证），平均、clip-norm
+（device reduction，确定性）、AdamW kernel 都是确定性的 → 每步喂给优化器的梯度逐位相同；(c) **同优化器状态**：
+每个 rank 各自维护 m/v，但因为起点相同、每步输入相同、超参相同，状态演化也相同。
+
+⇒ 各 rank 参数**逐位一致**（测试里断言 `max|p0-p1| == 0.0`），不需要任何权重再同步。
+
+> 对比单卡：DDP 与单卡的最终参数只在 **fp 容差**内一致（`<1e-3` rel），不逐位——因为求和顺序不同
+> （单卡 tape 顺序 SUM B 个；DDP 各 rank 先 SUM 分片再 NCCL 跨 rank SUM），fp 加法不结合。跨 rank 才逐位。
+
+### ⑤ Batch sharding：对住单卡的均值
+
+单卡 loop：`batch_size=B` 个序列各 forward+backward，tape **SUM** 出 `Σ_B`，再 `clip(pre_scale=1/B)` → batch 均值。
+
+DDP 严格对齐（保证 all-reduce 后的和与单卡逐序列一致）：**每个 rank 推进同一个 RNG、抽出整批 B 个序列，但只对
+分到自己的那些（`i % world == rank`）跑 forward+backward**。各 rank 的并集 == 单卡那一批同序，于是：
+
+```
+rank 本地:  Σ_local（分片内 b = B/world 个的和，tape SUM）
+AllReduce:  Σ_global = Σ_ranks Σ_local        （= 单卡的 Σ_B，只是求和顺序不同）
+/world:     Σ_global / world
+clip pre_scale = 1/b_local = world/B:  Σ_global/world · world/B = Σ_global / B  ← 与单卡 clip(1/B) 同一个量
+```
+
+「按 world 取均值」放在通信原语里（语义清晰），剩下的 `1/b_local` 交给本就存在的 clip 前置缩放完成——
+最终喂给 AdamW 的就是 **global batch 均值梯度**，与单卡完全对齐。loss 日志同理：本地 loss 和跨 rank AllReduce 后
+除以 B，得 global 均值（每个 rank 拿到相同值）。
+
+## 验证方法
+
+测试 `tests/ddp_correctness.rs`（`#[cfg(not(no_cuda))]`，<2 卡自动 skip），用合成语料（无需 tokenizer/数据文件）：
+
+### 正确性：`ddp_matches_single_gpu_and_params_consistent`
+
+同一 config + 合成语料，跑 20 步：
+
+- **(a) loss 对拍**：单卡 baseline（world=1，global batch 8）vs 2-rank DDP（各 4）→ 整条 loss 轨迹 `max_rel < 1e-3`。
+- **(b) 跨 rank 参数一致**：断言 `max|p_rank0 - p_rank1| == 0.0`（逐位，NCCL 保证 + 确定性 step）。
+- **(c) DDP vs 单卡参数**：`max rel|Δp| < 1e-3`（fp 求和顺序差，不逐位）。
+
+### 吞吐：`ddp_throughput_scaling`
+
+固定**每卡 batch=8**（batch 随 world 放大，每 rank 成本不变），测 world ∈ {1,2,4} 的 global tok/s，打印：
+
+```
+  GPUs | tok/s (global) | speedup
+     1 |          ...   |   1.00x
+     2 |          ...   |   ~2x
+     4 |          ...   |   ~4x
+```
+
+近线性即过（tiny model latency-bound + 跨卡通信开销会让大 world 略低于理想）。
+
+### dash5 实跑
+
+```bash
+export PATH=/usr/local/cuda/bin:/opt/wjh/.cargo/bin:$PATH
+# 选空闲卡（dash5 共享，先看 nvidia-smi）：
+CUDA_VISIBLE_DEVICES=<idle GPUs> \
+  cargo test -p xtrain-distributed --release -- --nocapture --test-threads=1
+# 真训练 / 吞吐 driver：
+CUDA_VISIBLE_DEVICES=0,1 cargo run -p xtrain-distributed --release --bin train_ddp -- \
+  100 64 16 /opt/wjh/models/gpt2/tokenizer.json data/tinystories-valid-3mb.txt
+```
+
+实测数字回填见 xtrain.md T8 note / commit。
+
+## 不做（本 Phase 范围外，记为 follow-up）
+
+- **Tensor / pipeline 并行**：要切权重、layer 内/跨 layer 通信，对 tiny model 收益小、改动大 → model 放大再做。
+- **Sharded optimizer（ZeRO）**：每卡只存 1/world 的优化器状态。当前 model 优化器状态很小，不是瓶颈。
+- **bf16 AllReduce / 通信压缩**：fp32 已够，且会动数值正确性（与 T7 同理由延后）。
+- **多进程 / 跨节点 bootstrap**：单机多卡用线程模型足够；跨节点再上文件/env 分发 UniqueId。