xtrain/crates/xtrain-train/src/bin/train.rs

//! End-to-end training entry point: load the GPT-2 BPE + a TinyStories corpus,
//! train the tiny transformer with hand-written AdamW for a BOUNDED budget,
//! evaluate held-out val loss, checkpoint the best, and print a few samples.
//!
//! The MODEL SIZE is a CLI-tunable scaling-ladder rung (v0 baseline = the
//! defaults; v1 = dim256/8L/8h via flags), not a hardcoded tiny config.
//!
//! Run on dash5 (needs a GPU + the corpus + tokenizer.json):
//!   export PATH=/usr/local/cuda/bin:/opt/wjh/.cargo/bin:$PATH
//!   cargo run -p xtrain-train --release --bin train -- \
//!       /opt/wjh/models/gpt2/tokenizer.json data/tinystories-train.txt \
//!       --dim 256 --heads 8 --head-dim 32 --layers 8 --ffn 1024 \
//!       --steps 3000 --batch 16 --seq 128 --max-lr 6e-4 \
//!       --val-tokens 200000 --eval-every 250 --ckpt /tmp/xtrain_v1.ckpt
//!
//! Positional: <tokenizer.json> <corpus.txt>. Everything else is a flag with a
//! sane default (defaults reproduce the v0-baseline tiny config).

// On a GPU-less host (no_cuda) the whole training body is unavailable; keep a
// stub `main` so the crate still builds for `cargo check`.
#[cfg(no_cuda)]
fn main() {
    eprintln!("xtrain train: built without CUDA (no_cuda); run on a GPU host (dash5).");
}

#[cfg(not(no_cuda))]
use std::path::{Path, PathBuf};

#[cfg(not(no_cuda))]
use xtrain_cuda::device;
#[cfg(not(no_cuda))]
use xtrain_model::{Config, TinyTransformer};
#[cfg(not(no_cuda))]
use xtrain_tensor::DType;
#[cfg(not(no_cuda))]
use xtrain_tensor::Device;
#[cfg(not(no_cuda))]
use xtrain_train::data::Corpus;
#[cfg(not(no_cuda))]
use xtrain_train::sample::generate;
#[cfg(not(no_cuda))]
use xtrain_train::schedule::LrSchedule;
#[cfg(not(no_cuda))]
use xtrain_train::{TrainConfig, train};

// Deterministic LCG fill in [-scale, scale) — same init scheme as the T5 tests.
#[cfg(not(no_cuda))]
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()
}

// A flag like `--dim 256`: 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() {
    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.
    let steps: usize = flag(&args, "--steps", 2000);
    let batch_size: usize = flag(&args, "--batch", 8);
    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 —
    // exact grads, lower peak activation memory (lets dim1024 batch32 fit). Opt-in;
    // default off stores every activation (unchanged numerics).
    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: PathBuf = PathBuf::from(
        args.iter()
            .position(|a| a == "--ckpt")
            .and_then(|i| args.get(i + 1))
            .cloned()
            .unwrap_or_else(|| "/tmp/xtrain_tinystories.ckpt".to_string()),
    );

    assert!(device::device_count().unwrap() > 0, "no CUDA device");
    device::set_device(0).unwrap();
    let device = Device::Cuda(0);

    println!(
        "loading tokenizer {} + corpus {} (cached id stream)",
        tok_path.display(),
        corpus_path.display()
    );
    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 (if requested and the corpus is big).
    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 mut seed = 1u64;
    let mut model = TinyTransformer::new(cfg, device, |shape| {
        seed = seed.wrapping_add(1);
        let n: usize = shape.iter().product();
        if shape.len() == 1 {
            // RMSNorm gammas → ~1.
            fill(n, seed, 0.02).iter().map(|v| v + 1.0).collect()
        } else {
            // Small fan-in-ish scale; keeps early logits tame.
            fill(n, seed, 0.04)
        }
    });
    if bf16 {
        model = model.with_compute_dtype(DType::BF16);
        println!("bf16 mixed precision: ON (fp32 master weights)");
    }
    if recompute {
        model = model.with_recompute(true);
        println!("activation recompute: ON (per-block gradient checkpointing)");
    }
    if flash {
        model = model.with_flash(true);
        println!("flash-attention: ON (fused SDPA kernel, no materialized scores)");
    }

    // Eval-only mode: load a checkpoint and score it on the held-out val set, then
    // exit. Used to put an EXISTING model (e.g. v0) and a new one on the same
    // metric — the v0-vs-v1 val-loss comparison. The arch flags must match the ckpt.
    if let Some(p) = args.iter().position(|a| a == "--eval-ckpt") {
        let ckpt_path = PathBuf::from(args.get(p + 1).expect("--eval-ckpt <path>"));
        xtrain_train::checkpoint::load_into(&ckpt_path, &model.params())
            .expect("load eval checkpoint");
        let v = valid.expect("--eval-ckpt needs --val-tokens > 0");
        let vl = xtrain_train::eval_loss(&model, device, &v, seq_len, eval_batches);
        println!("eval-only: {} → val loss {vl:.4}", ckpt_path.display());
        sample_some(&model, device, &tok_path);
        return;
    }

    let tcfg = TrainConfig {
        seq_len,
        batch_size,
        steps,
        schedule: LrSchedule {
            max_lr,
            min_lr,
            warmup: (steps / 20).max(20),
            total: steps,
        },
        weight_decay,
        max_grad_norm,
        log_every: 50,
        ckpt_path: Some(ckpt.clone()),
        ckpt_every: 500,
        eval_every,
        eval_batches,
        seed: 42,
    };

    println!(
        "training: {} steps, seq {}, batch {}, lr {:.1e}→{:.1e}, eval every {}",
        tcfg.steps,
        tcfg.seq_len,
        tcfg.batch_size,
        tcfg.schedule.max_lr,
        tcfg.schedule.min_lr,
        tcfg.eval_every
    );
    let result = train(&model, device, &train_corpus, valid.as_ref(), &tcfg);
    let start = result.train_losses.first().copied().unwrap_or(0.0);
    let end = result.train_losses.last().copied().unwrap_or(0.0);
    println!("train loss: start {start:.4} → end {end:.4}");
    if let Some(best) = result.best_val {
        println!("best val loss: {best:.4}");
    }
    if let Some((s, v)) = result.evals.last() {
        println!("final val loss (step {s}): {v:.4}");
    }

    sample_some(&model, device, &tok_path);
}

#[cfg(not(no_cuda))]
fn sample_some(model: &TinyTransformer, device: Device, tok_path: &Path) {
    use xserv_tokenizer::Tokenizer;
    let tok = Tokenizer::from_file(tok_path);
    let prompts = ["Once upon a time", "The little", "One day"];
    println!("\n--- samples (greedy) ---");
    for p in prompts {
        let ids: Vec<i32> = tok.encode(p).into_iter().map(|t| t as i32).collect();
        let mut rng = 7u64;
        let out = generate(model, device, &ids, 40, 0.0, &mut rng);
        let text = tok.decode(&out.iter().map(|&t| t as u32).collect::<Vec<_>>());
        println!("[{p}] → {text}");
    }
    println!("\n--- samples (temperature 0.8) ---");
    for p in prompts {
        let ids: Vec<i32> = tok.encode(p).into_iter().map(|t| t as i32).collect();
        let mut rng = 13u64;
        let out = generate(model, device, &ids, 40, 0.8, &mut rng);
        let text = tok.decode(&out.iter().map(|&t| t as u32).collect::<Vec<_>>());
        println!("[{p}] → {text}");
    }
}