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

//! Multi-rank DDP training launcher (Phase T8 / Scaling v2): spawn one thread per
//! GPU, NCCL all-reduce the gradients each step, and train the tiny transformer on
//! TinyStories. At parity with the single-GPU `bin/train`: CLI-tunable arch
//! (scaling-ladder rung), the cached token-id stream, held-out val-loss eval, LR
//! warmup→cosine, grad clip, and best-val checkpointing. 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=1,2 cargo run -p xtrain-distributed --release \
//!       --bin train_ddp -- /opt/wjh/models/gpt2/tokenizer.json \
//!       data/tinystories-train.txt \
//!       --dim 384 --heads 12 --head-dim 32 --layers 12 --ffn 1536 \
//!       --steps 6000 --batch 32 --seq 256 --max-lr 6e-4 \
//!       --val-tokens 1000000 --eval-every 500 --ckpt /tmp/xtrain_v2.ckpt
//!
//! Positional: <tokenizer.json> <corpus.txt>. Everything else is a flag with a
//! sane default. The launcher uses every GPU visible to it (CUDA_VISIBLE_DEVICES
//! selects them), so 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))]
use std::path::PathBuf;

// A flag like `--dim 384`: scan argv for `name`, parse the following token.
#[cfg(not(no_cuda))]
fn flag<T: std::str::FromStr>(args: &[String], name: &str, default: T) -> T {
    args.iter()
        .position(|a| a == name)
        .and_then(|i| args.get(i + 1))
        .and_then(|s| s.parse().ok())
        .unwrap_or(default)
}

#[cfg(not(no_cuda))]
fn main() {
    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();
    // First two non-flag positionals: tokenizer.json, corpus.txt.
    let positionals: Vec<&String> = args[1..].iter().filter(|a| !a.starts_with("--")).collect();
    let tok_path = positionals
        .first()
        .map(|s| PathBuf::from(s.as_str()))
        .unwrap_or_else(|| PathBuf::from("/opt/wjh/models/gpt2/tokenizer.json"));
    let corpus_path = positionals
        .get(1)
        .map(|s| PathBuf::from(s.as_str()))
        .unwrap_or_else(|| PathBuf::from("data/tinystories-valid-3mb.txt"));

    // Architecture (scaling-ladder rung). Defaults = v0-baseline tiny config.
    let n_heads = flag(&args, "--heads", 2usize);
    let head_dim = flag(&args, "--head-dim", 16usize);
    let n_layers = flag(&args, "--layers", 4usize);
    let ffn = flag(&args, "--ffn", 64usize);
    // GQA (Phase T15): num K/V heads (must divide --heads). Default = --heads (MHA).
    let kv_heads = flag(&args, "--kv-heads", n_heads);
    // `--dim` is informational; dim is always n_heads*head_dim. Warn on mismatch.
    let dim_flag = flag(&args, "--dim", 0usize);
    if dim_flag != 0 && dim_flag != n_heads * head_dim {
        eprintln!(
            "warning: --dim {dim_flag} != heads*head_dim {}; using {}",
            n_heads * head_dim,
            n_heads * head_dim
        );
    }

    // Optimization knobs (mirror bin/train).
    let steps: usize = flag(&args, "--steps", 100);
    let batch: usize = flag(&args, "--batch", 16);
    // Micro-batch gradient accumulation (Phase T16): effective global batch =
    // accum_steps × batch, all-reducing only at the accumulation boundary. Default
    // 1 = no accumulation (bit-identical to the pre-T16 DDP path).
    let accum_steps: usize = flag(&args, "--accum-steps", 1).max(1);
    let seq_len: usize = flag(&args, "--seq", 64);
    let max_lr: f32 = flag(&args, "--max-lr", 3e-3);
    let min_lr: f32 = flag(&args, "--min-lr", max_lr * 0.1);
    let weight_decay: f32 = flag(&args, "--wd", 0.1);
    let max_grad_norm: f32 = flag(&args, "--clip", 1.0);
    let val_tokens: usize = flag(&args, "--val-tokens", 0);
    let eval_every: usize = flag(&args, "--eval-every", 0);
    let eval_batches: usize = flag(&args, "--eval-batches", 64);
    // Dropout (Phase T18/T21): residual-path dropout prob, active at training time
    // only (inverted scaling), identity at eval/sampling/export. Default 0 = off
    // (forward graph bit-identical to the no-dropout path). Mirrors bin/train; the
    // train_rank loop calls model.train() each step so dropout is actually live
    // under DDP (T21 wired this — the launcher previously never set training mode).
    let dropout: f32 = flag(&args, "--dropout", 0.0f32);
    // bf16 mixed precision (Phase T12): fp32 master weights, bf16 linears +
    // activations. Opt-in; default fp32 reproduces v0–v4 numerics.
    let bf16 = args.iter().any(|a| a == "--bf16");
    // Activation recomputation (Phase T13): per-block gradient checkpointing — each
    // rank checkpoints its own forward/backward; exact grads, lower peak activation
    // memory (lets dim1024 batch32 fit). Opt-in; default off.
    let recompute = args.iter().any(|a| a == "--recompute");
    // Fused flash-attention (Phase T14): single fused SDPA kernel, online softmax,
    // no materialized [bh,S,S] scores. Opt-in; default off keeps the composed path.
    let flash = args.iter().any(|a| a == "--flash");
    let ckpt: Option<PathBuf> = args
        .iter()
        .position(|a| a == "--ckpt")
        .and_then(|i| args.get(i + 1))
        .map(PathBuf::from);

    // 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
    );

    // Reuse the cached token-id stream (v1's u16 cache); never re-tokenize 2GB.
    let corpus = Corpus::load_cached(&tok_path, &corpus_path);
    println!(
        "corpus: {} tokens, vocab {}",
        corpus.len(),
        corpus.vocab_size
    );
    let vocab = corpus.vocab_size;
    // Hold out a tail slice for validation (rank 0 evaluates on it).
    let (train_corpus, valid) = if val_tokens > 0 {
        let (t, v) = corpus.split_tail(val_tokens);
        println!("split: {} train tokens / {} val tokens", t.len(), v.len());
        (t, Some(v))
    } else {
        (corpus, None)
    };

    let mut cfg =
        Config::from_arch(vocab, n_heads, head_dim, n_layers, ffn).with_kv_heads(kv_heads);
    cfg.dropout = dropout;
    println!(
        "model: dim {} layers {} heads {} kv_heads {} head_dim {} ffn {} → core {:.3}M params \
         (+ embed/lm {:.2}M = {:.2}M total)",
        cfg.dim,
        cfg.n_layers,
        cfg.n_heads,
        cfg.num_kv_heads,
        cfg.head_dim,
        cfg.ffn_hidden,
        cfg.core_params() as f32 / 1e6,
        (cfg.num_params() - cfg.core_params()) as f32 / 1e6,
        cfg.num_params() as f32 / 1e6,
    );

    let dcfg = DdpConfig {
        seq_len,
        batch_size: batch,
        accum_steps,
        steps,
        schedule: LrSchedule {
            max_lr,
            min_lr,
            warmup: (steps / 20).max(5),
            total: steps,
        },
        weight_decay,
        max_grad_norm,
        log_every: 50,
        seed: 42,
        eval_every,
        eval_batches,
        ckpt_path: ckpt.clone(),
    };

    println!(
        "training: {steps} steps, seq {seq_len}, global batch {batch} × accum {accum_steps} = \
         effective global batch {}, lr {max_lr:.1e}→{min_lr:.1e}, eval every {eval_every}",
        batch * accum_steps
    );

    if bf16 {
        println!("bf16 mixed precision: ON (fp32 master weights)");
    }
    if recompute {
        println!("activation recompute: ON (per-block gradient checkpointing)");
    }
    if flash {
        println!("flash-attention: ON (fused SDPA kernel, no materialized scores)");
    }
    if dropout > 0.0 {
        println!("dropout: ON (p={dropout}, residual-path, train-only inverted scaling)");
    }
    let results = launch(
        &devices,
        &train_corpus,
        valid.as_ref(),
        &dcfg,
        move |device| {
            let mut m = build_model(cfg, device);
            if bf16 {
                m = m.with_compute_dtype(xtrain_tensor::DType::BF16);
            }
            if recompute {
                m = m.with_recompute(true);
            }
            if flash {
                m = m.with_flash(true);
            }
            m
        },
    );
    let r0 = &results[0];
    let start = r0.losses.first().copied().unwrap_or(0.0);
    let end = r0.losses.last().copied().unwrap_or(0.0);
    println!("train loss: start {start:.4} → end {end:.4}");
    if let Some(best) = r0.best_val {
        println!("best val loss: {best:.4}");
    }
    if let Some((s, v)) = r0.evals.last() {
        println!("final val loss (step {s}): {v:.4}");
    }
    if let Some(path) = &ckpt {
        println!("best-val checkpoint → {}", path.display());
    }
}