xtrain/crates/xtrain-train/tests/real_training.rs

// Real-training acceptance (Phase T6): train the tiny transformer on the
// TinyStories corpus (tokenized with the reused GPT-2 BPE) for a BOUNDED budget
// and assert the loss decreases substantially — the end-to-end signal that the
// whole stack (data pipeline, AdamW, LR schedule, grad clip) learns. Prints the
// loss curve and a couple of greedy samples.
//
// Needs the corpus + tokenizer present, so it is #[ignore] (run with --ignored)
// and gated #![cfg(not(no_cuda))]. Paths are overridable via env vars.
//
// Run:  cargo test -p xtrain-train --release --test real_training \
//         -- --ignored --nocapture
#![cfg(not(no_cuda))]

use std::path::PathBuf;

use xtrain_cuda::device;
use xtrain_model::{Config, TinyTransformer};
use xtrain_tensor::Device;
use xtrain_train::data::Corpus;
use xtrain_train::sample::generate;
use xtrain_train::schedule::LrSchedule;
use xtrain_train::{TrainConfig, train};

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()
}

#[test]
#[ignore = "real training; needs corpus + tokenizer; run with --ignored --release"]
fn trains_on_tinystories() {
    assert!(device::device_count().unwrap() > 0, "no CUDA device");
    device::set_device(0).unwrap();
    let device = Device::Cuda(0);

    let tok_path = PathBuf::from(
        std::env::var("XTRAIN_TOKENIZER")
            .unwrap_or_else(|_| "/opt/wjh/models/gpt2/tokenizer.json".into()),
    );
    // Default resolves relative to the repo root (cargo runs tests with cwd =
    // crate dir, so `../../data/...` from crates/xtrain-train); override with
    // XTRAIN_CORPUS for any other location.
    let corpus_path = PathBuf::from(std::env::var("XTRAIN_CORPUS").unwrap_or_else(|_| {
        format!(
            "{}/../../data/tinystories-valid-3mb.txt",
            env!("CARGO_MANIFEST_DIR")
        )
    }));

    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;
    let mut seed = 1u64;
    let model = 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)
        }
    });

    let steps = std::env::var("XTRAIN_STEPS")
        .ok()
        .and_then(|s| s.parse().ok())
        .unwrap_or(800usize);
    let tcfg = TrainConfig {
        seq_len: 64,
        batch_size: 8,
        accum_steps: 1,
        steps,
        schedule: LrSchedule {
            max_lr: 3e-3,
            min_lr: 3e-4,
            warmup: (steps / 20).max(20),
            total: steps,
        },
        weight_decay: 0.1,
        max_grad_norm: 1.0,
        log_every: 50,
        ckpt_path: None,
        ckpt_every: 0,
        eval_every: 0,
        eval_batches: 0,
        seed: 42,
    };

    let losses = train(&model, device, &corpus, None, &tcfg).train_losses;
    // Average the first/last few steps to smooth per-step noise.
    let head: f32 =
        losses[..10.min(losses.len())].iter().sum::<f32>() / 10.0_f32.min(losses.len() as f32);
    let tail_n = 10.min(losses.len());
    let tail: f32 = losses[losses.len() - tail_n..].iter().sum::<f32>() / tail_n as f32;
    println!("loss: start(avg10) {head:.4} → end(avg10) {tail:.4}");

    // A couple of greedy samples (should show English structure, not gibberish).
    use xserv_tokenizer::Tokenizer;
    let tok = Tokenizer::from_file(&tok_path);
    for p in ["Once upon a time", "The little"] {
        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!("sample [{p}] → {text}");
    }

    // Bounded run: expect a substantial drop (not full convergence).
    assert!(
        tail < head - 0.5,
        "loss did not decrease substantially: {head:.4} → {tail:.4}"
    );
    assert!(tail < 6.5, "final loss implausibly high: {tail:.4}");
}