Gahow Wang
f26db882e5
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> # Conflicts: # README.md # docs/evolution.md
222 lines
8.4 KiB
Rust
222 lines
8.4 KiB
Rust
//! 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);
|
||
// `--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);
|
||
// 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 cfg = Config::from_arch(vocab, n_heads, head_dim, n_layers, ffn);
|
||
println!(
|
||
"model: dim {} layers {} heads {} head_dim {} ffn {} → core {:.3}M params \
|
||
(+ embed/lm {:.2}M = {:.2}M total)",
|
||
cfg.dim,
|
||
cfg.n_layers,
|
||
cfg.n_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)");
|
||
}
|
||
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());
|
||
}
|
||
}
|