Gahow Wang
e625aa05dd
Config.dropout (default 0). TinyTransformer gets a Cell<bool> training switch (train()/eval()/with_training, default eval = safe) + a Cell<u64> step_seed bumped once per training forward. forward_batched derives a per-layer block_seed (pure fn of step_seed×layer) and block_forward derives two per-site seeds, inserting ops::dropout at the attn and ffn sub-block outputs (before each residual). The seed is a pure function of (step_seed, layer, site) so the checkpoint (T13) recompute re-derives the same masks → grads stay exact. p=0 or eval → no dropout node → graph bit-identical to pre-T18. train_loop: model.train() per step (restored after eval flips to eval); eval_loss runs model.eval(). bin/train: --dropout flag → cfg.dropout. Export/sampling run in eval (default), so exported weights are dropout-free (xserv closed loop unaffected). Model-level tests (dropout.rs): p=0 bit-identical to no-dropout (logits/loss/grads); eval(p>0) == p=0 identity; train differs from eval + finite; recompute-with-dropout grads match non-recompute (fp32 + bf16). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
274 lines
10 KiB
Rust
274 lines
10 KiB
Rust
//! 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);
|
||
// Dropout (Phase T18): 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).
|
||
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 —
|
||
// 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");
|
||
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 mut cfg = Config::from_arch(vocab, n_heads, head_dim, n_layers, ffn);
|
||
cfg.dropout = dropout;
|
||
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 dropout > 0.0 {
|
||
println!("dropout: ON (p={dropout}, residual-path, train-only inverted scaling)");
|
||
}
|
||
|
||
// 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}");
|
||
}
|
||
}
|