data: full TinyStories + tokenized-id cache, val loss, CLI arch
- Corpus::load_cached: tokenize the (large) corpus ONCE, cache the id stream to
<corpus>.u16.bin (gpt2 vocab 50257 < 65536 → exact u16), read cache on reruns.
- Corpus::split_tail: hold out a tail slice as a validation corpus.
- train(): take an optional valid corpus + eval_every/eval_batches; periodic
deterministic val-loss eval that checkpoints the BEST val model; returns
TrainResult{train_losses, evals, best_val}. T6 fixed-cadence path preserved.
- bin/train + bin/export_safetensors: read architecture (--heads/--head-dim/
--layers/--ffn) + opt knobs (--steps/--batch/--seq/--max-lr/--val-tokens/
--eval-every) from CLI flags; defaults reproduce the v0-baseline tiny config.
- gitignore the multi-GB corpus + *.u16.bin caches + *.ckpt (dash5-only).
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
@@ -1,14 +1,20 @@
|
||||
//! End-to-end training entry point (Phase T6): load the GPT-2 BPE + TinyStories
|
||||
//! corpus, train the tiny transformer with hand-written AdamW for a BOUNDED
|
||||
//! budget, checkpoint it, and print a few generated samples.
|
||||
//! 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-valid-3mb.txt
|
||||
//! /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
|
||||
//!
|
||||
//! Optional 3rd/4th args: number of steps, checkpoint path.
|
||||
//! 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`.
|
||||
@@ -51,51 +57,101 @@ fn fill(n: usize, seed: u64, scale: f32) -> Vec<f32> {
|
||||
.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();
|
||||
let tok_path = args
|
||||
.get(1)
|
||||
.map(PathBuf::from)
|
||||
// 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 = args
|
||||
.get(2)
|
||||
.map(PathBuf::from)
|
||||
let corpus_path = positionals
|
||||
.get(1)
|
||||
.map(|s| PathBuf::from(s.as_str()))
|
||||
.unwrap_or_else(|| PathBuf::from("data/tinystories-valid-3mb.txt"));
|
||||
let steps: usize = args.get(3).and_then(|s| s.parse().ok()).unwrap_or(2000);
|
||||
let ckpt: PathBuf = args
|
||||
.get(4)
|
||||
.map(PathBuf::from)
|
||||
.unwrap_or_else(|| PathBuf::from("/tmp/xtrain_tinystories.ckpt"));
|
||||
|
||||
// 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);
|
||||
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 {}",
|
||||
"loading tokenizer {} + corpus {} (cached id stream)",
|
||||
tok_path.display(),
|
||||
corpus_path.display()
|
||||
);
|
||||
let corpus = Corpus::load(&tok_path, &corpus_path);
|
||||
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)
|
||||
};
|
||||
|
||||
// Tiny model sized to the BPE vocab. A real (but small) config: wider than
|
||||
// the overfit test so it has capacity to learn English structure.
|
||||
let mut cfg = Config::tiny();
|
||||
cfg.vocab = corpus.vocab_size;
|
||||
cfg.n_layers = 4;
|
||||
let cfg = Config::from_arch(vocab, n_heads, head_dim, n_layers, ffn);
|
||||
println!(
|
||||
"model: dim {} layers {} heads {} ffn {} → {} params",
|
||||
"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.num_params()
|
||||
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;
|
||||
@@ -111,33 +167,45 @@ fn main() {
|
||||
}
|
||||
});
|
||||
|
||||
let seq_len = 64;
|
||||
let tcfg = TrainConfig {
|
||||
seq_len,
|
||||
batch_size: 8,
|
||||
batch_size,
|
||||
steps,
|
||||
schedule: LrSchedule {
|
||||
max_lr: 3e-3,
|
||||
min_lr: 3e-4,
|
||||
max_lr,
|
||||
min_lr,
|
||||
warmup: (steps / 20).max(20),
|
||||
total: steps,
|
||||
},
|
||||
weight_decay: 0.1,
|
||||
max_grad_norm: 1.0,
|
||||
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}",
|
||||
tcfg.steps, tcfg.seq_len, tcfg.batch_size, tcfg.schedule.max_lr, tcfg.schedule.min_lr
|
||||
"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 losses = train(&model, device, &corpus, &tcfg);
|
||||
let start = losses.first().copied().unwrap_or(0.0);
|
||||
let end = losses.last().copied().unwrap_or(0.0);
|
||||
println!("loss: start {start:.4} → end {end:.4}");
|
||||
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);
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user