docs: Phase T6 — training loop

Design doc for the T6 training stack: Goal / Module Layout / Key Design
Decisions (AdamW math + decoupled WD, LR schedule, global-norm grad clip with
batch averaging, checkpoint format, data pipeline + xserv tokenizer reuse,
sampler) / 验证方法 (AdamW parity, checkpoint round-trip, real training, host
unit tests).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Cargo.lock

@@ -2,6 +2,15 @@
 # It is not intended for manual editing.
 version = 4
 
+[[package]]
+name = "aho-corasick"
+version = "1.1.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301"
+dependencies = [
+ "memchr",
+]
+
 [[package]]
 name = "cc"
 version = "1.2.64"
@@ -41,6 +50,18 @@ dependencies = [
  "zerocopy",
 ]
 
+[[package]]
+name = "itoa"
+version = "1.0.18"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682"
+
+[[package]]
+name = "memchr"
+version = "2.8.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "88904434abc2901f197fe8cc55f0445e7ded921dba5911dad2e2b39b48e663c4"
+
 [[package]]
 name = "proc-macro2"
 version = "1.0.106"
@@ -59,6 +80,78 @@ dependencies = [
  "proc-macro2",
 ]
 
+[[package]]
+name = "regex"
+version = "1.12.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f1292b7759ae1cb9ec195452d1390a074f0cd8541ab7a5a8c31cd6db45d4a6ba"
+dependencies = [
+ "aho-corasick",
+ "memchr",
+ "regex-automata",
+ "regex-syntax",
+]
+
+[[package]]
+name = "regex-automata"
+version = "0.4.14"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6e1dd4122fc1595e8162618945476892eefca7b88c52820e74af6262213cae8f"
+dependencies = [
+ "aho-corasick",
+ "memchr",
+ "regex-syntax",
+]
+
+[[package]]
+name = "regex-syntax"
+version = "0.8.11"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4"
+
+[[package]]
+name = "serde"
+version = "1.0.228"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e"
+dependencies = [
+ "serde_core",
+ "serde_derive",
+]
+
+[[package]]
+name = "serde_core"
+version = "1.0.228"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad"
+dependencies = [
+ "serde_derive",
+]
+
+[[package]]
+name = "serde_derive"
+version = "1.0.228"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn",
+]
+
+[[package]]
+name = "serde_json"
+version = "1.0.150"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e8014e44b4736ed0538adeecded0fce2a272f22dc9578a7eb6b2d9993c74cfb9"
+dependencies = [
+ "itoa",
+ "memchr",
+ "serde",
+ "serde_core",
+ "zmij",
+]
+
 [[package]]
 name = "shlex"
 version = "2.0.1"
@@ -88,6 +181,15 @@ version = "1.0.24"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75"
 
+[[package]]
+name = "xserv-tokenizer"
+version = "0.1.0"
+dependencies = [
+ "regex",
+ "serde",
+ "serde_json",
+]
+
 [[package]]
 name = "xtrain-autodiff"
 version = "0.1.0"
@@ -112,6 +214,15 @@ dependencies = [
  "xtrain-tensor",
 ]
 
+[[package]]
+name = "xtrain-optim"
+version = "0.1.0"
+dependencies = [
+ "xtrain-autodiff",
+ "xtrain-cuda",
+ "xtrain-tensor",
+]
+
 [[package]]
 name = "xtrain-tensor"
 version = "0.1.0"
@@ -122,6 +233,18 @@ dependencies = [
  "xtrain-cuda",
 ]
 
+[[package]]
+name = "xtrain-train"
+version = "0.1.0"
+dependencies = [
+ "xserv-tokenizer",
+ "xtrain-autodiff",
+ "xtrain-cuda",
+ "xtrain-model",
+ "xtrain-optim",
+ "xtrain-tensor",
+]
+
 [[package]]
 name = "zerocopy"
 version = "0.8.52"
@@ -141,3 +264,9 @@ dependencies = [
  "quote",
  "syn",
 ]
+
+[[package]]
+name = "zmij"
+version = "1.0.21"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa"

Cargo.toml

@@ -5,6 +5,8 @@ members = [
     "crates/xtrain-tensor",
     "crates/xtrain-autodiff",
     "crates/xtrain-model",
+    "crates/xtrain-optim",
+    "crates/xtrain-train",
 ]
 
 [workspace.package]

crates/xtrain-optim/Cargo.toml

@@ -0,0 +1,12 @@
+[package]
+name = "xtrain-optim"
+version.workspace = true
+edition.workspace = true
+
+[dependencies]
+xtrain-tensor = { path = "../xtrain-tensor" }
+xtrain-autodiff = { path = "../xtrain-autodiff" }
+
+[dev-dependencies]
+# Acceptance tests drive the GPU (device selection) directly.
+xtrain-cuda = { path = "../xtrain-cuda" }

crates/xtrain-optim/build.rs

@@ -0,0 +1,26 @@
+use std::env;
+use std::path::Path;
+use std::process::Command;
+
+// Per-crate convention (see the other crates): the AdamW *math* is host-only and
+// always compiles, but `AdamW::step(&[Var])` round-trips parameter values/grads
+// through GPU tensors, so that call site is gated behind `not(no_cuda)`. cfg does
+// not propagate across crates, so this crate re-detects nvcc. No CUDA is compiled.
+fn main() {
+    println!("cargo:rustc-check-cfg=cfg(no_cuda)");
+
+    let cuda_path = env::var("CUDA_HOME")
+        .or_else(|_| env::var("CUDA_PATH"))
+        .unwrap_or_else(|_| "/usr/local/cuda".to_string());
+
+    if !nvcc_available(&cuda_path) {
+        println!("cargo:rustc-cfg=no_cuda");
+    }
+}
+
+fn nvcc_available(cuda_path: &str) -> bool {
+    if Command::new("nvcc").arg("--version").output().is_ok() {
+        return true;
+    }
+    Path::new(&format!("{cuda_path}/bin/nvcc")).exists()
+}

crates/xtrain-optim/src/lib.rs

@@ -0,0 +1,154 @@
+//! Hand-written AdamW optimizer (Phase T6).
+//!
+//! AdamW = Adam with **decoupled** weight decay (Loshchilov & Hutter, 2019): the
+//! weight-decay term is applied directly to the parameter, NOT folded into the
+//! gradient (so it does not interact with the adaptive `v` denominator). This
+//! matches `torch.optim.AdamW`.
+//!
+//! Update for parameter `θ` at step `t` (1-indexed), with gradient `g`:
+//! ```text
+//!   m ← β1·m + (1−β1)·g
+//!   v ← β2·v + (1−β2)·g²
+//!   m̂ ← m / (1 − β1ᵗ)          (bias correction)
+//!   v̂ ← v / (1 − β2ᵗ)
+//!   θ ← θ − lr·( m̂ / (√v̂ + ε) + wd·θ )
+//! ```
+//! The `lr·wd·θ` term is the decoupled decay. Note PyTorch applies decay as
+//! `θ ← θ·(1 − lr·wd)` then the Adam step; both are algebraically the same
+//! first-order update — we fold decay into the single subtraction above, which
+//! is what PyTorch's default (`maximize=False`, no `amsgrad`) computes.
+//!
+//! The math operates on flat host `f32` buffers ([`AdamW::step_host`]) so it is
+//! unit-testable on a GPU-less host; [`AdamW::step`] is a thin wrapper that
+//! round-trips each parameter's value/grad through the GPU tensor and is gated
+//! behind `not(no_cuda)`.
+
+/// Per-parameter optimizer state: the first (`m`) and second (`v`) moment
+/// estimates, one f32 per element, kept flat (matching the parameter layout).
+struct ParamState {
+    m: Vec<f32>,
+    v: Vec<f32>,
+}
+
+/// Decoupled-weight-decay Adam. One instance owns the moment state for a fixed
+/// list of parameters, keyed by their index in the slice passed to `step`
+/// (the model's stable `params()` order).
+pub struct AdamW {
+    pub lr: f32,
+    beta1: f32,
+    beta2: f32,
+    eps: f32,
+    weight_decay: f32,
+    /// Global step count (shared across all params for bias correction).
+    t: u64,
+    /// Lazily sized to the parameter list on the first `step`.
+    state: Vec<ParamState>,
+}
+
+impl AdamW {
+    /// PyTorch-default hyperparameters except `lr`/`weight_decay`, which you set
+    /// (β1=0.9, β2=0.999, ε=1e-8).
+    pub fn new(lr: f32, weight_decay: f32) -> Self {
+        Self::with_betas(lr, weight_decay, 0.9, 0.999, 1e-8)
+    }
+
+    pub fn with_betas(lr: f32, weight_decay: f32, beta1: f32, beta2: f32, eps: f32) -> Self {
+        Self {
+            lr,
+            beta1,
+            beta2,
+            eps,
+            weight_decay,
+            t: 0,
+            state: Vec::new(),
+        }
+    }
+
+    /// Current global step (number of `step` calls so far).
+    pub fn step_count(&self) -> u64 {
+        self.t
+    }
+
+    /// Pure-host AdamW step over flat parameter/gradient buffers. `params[i]` is
+    /// updated in place using `grads[i]`; both are the i-th parameter's elements
+    /// in the model's stable order. Lazily allocates moment state on first call.
+    ///
+    /// This is the testable core — no GPU, no autograd. `lr` is passed per call
+    /// so a schedule can vary it each step.
+    pub fn step_host(&mut self, lr: f32, params: &mut [Vec<f32>], grads: &[Vec<f32>]) {
+        assert_eq!(params.len(), grads.len(), "param/grad count mismatch");
+        if self.state.is_empty() {
+            self.state = params
+                .iter()
+                .map(|p| ParamState {
+                    m: vec![0.0; p.len()],
+                    v: vec![0.0; p.len()],
+                })
+                .collect();
+        }
+        assert_eq!(self.state.len(), params.len(), "param count changed");
+
+        self.t += 1;
+        let bc1 = 1.0 - self.beta1.powi(self.t as i32);
+        let bc2 = 1.0 - self.beta2.powi(self.t as i32);
+
+        for (i, (p, g)) in params.iter_mut().zip(grads).enumerate() {
+            assert_eq!(p.len(), g.len(), "param/grad len mismatch at {i}");
+            let st = &mut self.state[i];
+            for j in 0..p.len() {
+                let gj = g[j];
+                st.m[j] = self.beta1 * st.m[j] + (1.0 - self.beta1) * gj;
+                st.v[j] = self.beta2 * st.v[j] + (1.0 - self.beta2) * gj * gj;
+                let mhat = st.m[j] / bc1;
+                let vhat = st.v[j] / bc2;
+                // Decoupled weight decay: decay term uses the *current* param,
+                // matching PyTorch's `p ← p − lr·wd·p` applied alongside the step.
+                p[j] -= lr * (mhat / (vhat.sqrt() + self.eps) + self.weight_decay * p[j]);
+            }
+        }
+    }
+}
+
+#[cfg(not(no_cuda))]
+mod gpu {
+    use super::AdamW;
+    use xtrain_autodiff::tape::Var;
+    use xtrain_tensor::{Device, Tensor};
+
+    impl AdamW {
+        /// Apply one AdamW step to every parameter `Var`, using `lr` for this step
+        /// (so an LR schedule can vary it). Pulls each param's value and `.grad()`
+        /// to the host, runs [`AdamW::step_host`], and writes the updated value
+        /// back with `set_value`. A param with no grad is fed a zero grad, so the
+        /// Adam term vanishes and only decoupled weight decay applies (the model's
+        /// params all receive grads each step, so this is just a safety default).
+        ///
+        /// Does NOT zero grads — the caller does that (matching the GD-step
+        /// template in the T5 overfit test).
+        pub fn step(&mut self, lr: f32, params: &[Var]) {
+            let device = params[0].value().device();
+            let shapes: Vec<Vec<usize>> =
+                params.iter().map(|p| p.value().shape().to_vec()).collect();
+
+            let mut host_params: Vec<Vec<f32>> = params
+                .iter()
+                .map(|p| p.value().to_device(Device::Cpu).as_slice::<f32>().to_vec())
+                .collect();
+            let host_grads: Vec<Vec<f32>> = params
+                .iter()
+                .zip(&host_params)
+                .map(|(p, hp)| match p.grad() {
+                    Some(g) => g.to_device(Device::Cpu).as_slice::<f32>().to_vec(),
+                    None => vec![0.0; hp.len()], // no grad → no update this step
+                })
+                .collect();
+
+            self.step_host(lr, &mut host_params, &host_grads);
+
+            for ((p, data), shape) in params.iter().zip(&host_params).zip(&shapes) {
+                let t = Tensor::from_slice(data, shape).to_device(device);
+                p.set_value(t);
+            }
+        }
+    }
+}

crates/xtrain-optim/tests/adamw_host.rs

@@ -0,0 +1,99 @@
+// Host-only unit test for the AdamW *math* (no GPU). Verifies the update against
+// an independent, hand-rolled reference implementation of the same recurrence for
+// several steps with non-trivial weight decay — catching bias-correction and
+// decoupled-decay mistakes. The rigorous vs-PyTorch parity (end-to-end on a real
+// model) lives in xtrain-train; this is the fast local guard on the formula.
+
+use xtrain_optim::AdamW;
+
+// Independent reference: the textbook AdamW recurrence, kept separate from the
+// implementation so a shared bug can't hide.
+struct RefAdamW {
+    b1: f32,
+    b2: f32,
+    eps: f32,
+    wd: f32,
+    t: i32,
+    m: Vec<f32>,
+    v: Vec<f32>,
+}
+
+impl RefAdamW {
+    fn new(n: usize, wd: f32) -> Self {
+        Self {
+            b1: 0.9,
+            b2: 0.999,
+            eps: 1e-8,
+            wd,
+            t: 0,
+            m: vec![0.0; n],
+            v: vec![0.0; n],
+        }
+    }
+    fn step(&mut self, lr: f32, p: &mut [f32], g: &[f32]) {
+        self.t += 1;
+        let bc1 = 1.0 - self.b1.powi(self.t);
+        let bc2 = 1.0 - self.b2.powi(self.t);
+        for i in 0..p.len() {
+            self.m[i] = self.b1 * self.m[i] + (1.0 - self.b1) * g[i];
+            self.v[i] = self.b2 * self.v[i] + (1.0 - self.b2) * g[i] * g[i];
+            let mhat = self.m[i] / bc1;
+            let vhat = self.v[i] / bc2;
+            p[i] -= lr * (mhat / (vhat.sqrt() + self.eps) + self.wd * p[i]);
+        }
+    }
+}
+
+#[test]
+fn adamw_matches_reference_recurrence() {
+    let lr = 0.01;
+    let wd = 0.1;
+    let mut opt = AdamW::new(lr, wd);
+
+    // Two parameters of different sizes (exercises per-param state keying).
+    let mut p_impl = vec![vec![0.5f32, -1.0, 2.0, 0.0], vec![1.5f32, -0.25]];
+    let mut p_ref = p_impl.clone();
+    let mut r0 = RefAdamW::new(4, wd);
+    let mut r1 = RefAdamW::new(2, wd);
+
+    // Deterministic pseudo-grads that change every step.
+    let grad = |step: usize, idx: usize, j: usize| -> f32 {
+        let s = (step * 13 + idx * 7 + j * 3) as f32;
+        (s * 0.123).sin() * 0.5
+    };
+
+    for step in 0..20 {
+        let grads = vec![
+            (0..4).map(|j| grad(step, 0, j)).collect::<Vec<_>>(),
+            (0..2).map(|j| grad(step, 1, j)).collect::<Vec<_>>(),
+        ];
+        opt.step_host(lr, &mut p_impl, &grads);
+        r0.step(lr, &mut p_ref[0], &grads[0]);
+        r1.step(lr, &mut p_ref[1], &grads[1]);
+    }
+
+    assert_eq!(opt.step_count(), 20);
+    for (pi, pr) in p_impl.iter().zip(&p_ref) {
+        for (a, b) in pi.iter().zip(pr) {
+            assert!((a - b).abs() < 1e-6, "impl {a} != ref {b}");
+        }
+    }
+}
+
+#[test]
+fn zero_grad_only_decays() {
+    // With g=0 and wd>0, the step must reduce to pure decoupled decay:
+    // θ ← θ − lr·wd·θ (Adam term is 0/eps = 0).
+    let lr = 0.1;
+    let wd = 0.5;
+    let mut opt = AdamW::new(lr, wd);
+    let mut p = vec![vec![2.0f32]];
+    let g = vec![vec![0.0f32]];
+    opt.step_host(lr, &mut p, &g);
+    let expected = 2.0 - lr * wd * 2.0;
+    assert!(
+        (p[0][0] - expected).abs() < 1e-6,
+        "{} != {expected}",
+        p[0][0]
+    );
+}

crates/xtrain-train/Cargo.toml

@@ -0,0 +1,20 @@
+[package]
+name = "xtrain-train"
+version.workspace = true
+edition.workspace = true
+
+[dependencies]
+xtrain-tensor = { path = "../xtrain-tensor" }
+xtrain-autodiff = { path = "../xtrain-autodiff" }
+xtrain-model = { path = "../xtrain-model" }
+xtrain-optim = { path = "../xtrain-optim" }
+xtrain-cuda = { path = "../xtrain-cuda" }
+# Reuse xserv's from-scratch GPT-2/Qwen BPE (project decision). This relative
+# path resolves on both ~/projects (local) and /opt/wjh/projects (dash5). The
+# crate inherits xserv's workspace for its own deps (serde/regex) — Cargo reads
+# the target package's workspace, not ours.
+xserv-tokenizer = { path = "../../../xserv/crates/xserv-tokenizer" }
+
+[[bin]]
+name = "train"
+path = "src/bin/train.rs"

crates/xtrain-train/build.rs

@@ -0,0 +1,26 @@
+use std::env;
+use std::path::Path;
+use std::process::Command;
+
+// Per-crate convention: the training loop / sampler / checkpoint all drive GPU
+// ops through the model + tensor layers, so the bulk of this crate is gated
+// behind `not(no_cuda)`. The LR schedule and the grad-clip *math* are host-only
+// and always compile. cfg does not propagate across crates, so re-detect nvcc.
+fn main() {
+    println!("cargo:rustc-check-cfg=cfg(no_cuda)");
+
+    let cuda_path = env::var("CUDA_HOME")
+        .or_else(|_| env::var("CUDA_PATH"))
+        .unwrap_or_else(|_| "/usr/local/cuda".to_string());
+
+    if !nvcc_available(&cuda_path) {
+        println!("cargo:rustc-cfg=no_cuda");
+    }
+}
+
+fn nvcc_available(cuda_path: &str) -> bool {
+    if Command::new("nvcc").arg("--version").output().is_ok() {
+        return true;
+    }
+    Path::new(&format!("{cuda_path}/bin/nvcc")).exists()
+}

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

@@ -0,0 +1,166 @@
+//! 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.
+//!
+//! 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
+//!
+//! Optional 3rd/4th args: number of steps, checkpoint path.
+
+// 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::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()
+}
+
+#[cfg(not(no_cuda))]
+fn main() {
+    let args: Vec<String> = std::env::args().collect();
+    let tok_path = args
+        .get(1)
+        .map(PathBuf::from)
+        .unwrap_or_else(|| PathBuf::from("/opt/wjh/models/gpt2/tokenizer.json"));
+    let corpus_path = args
+        .get(2)
+        .map(PathBuf::from)
+        .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"));
+
+    assert!(device::device_count().unwrap() > 0, "no CUDA device");
+    device::set_device(0).unwrap();
+    let device = Device::Cuda(0);
+
+    println!(
+        "loading tokenizer {} + corpus {}",
+        tok_path.display(),
+        corpus_path.display()
+    );
+    let corpus = Corpus::load(&tok_path, &corpus_path);
+    println!(
+        "corpus: {} tokens, vocab {}",
+        corpus.len(),
+        corpus.vocab_size
+    );
+
+    // 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;
+    println!(
+        "model: dim {} layers {} heads {} ffn {} → {} params",
+        cfg.dim,
+        cfg.n_layers,
+        cfg.n_heads,
+        cfg.ffn_hidden,
+        cfg.num_params()
+    );
+
+    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 {
+            // 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)
+        }
+    });
+
+    let seq_len = 64;
+    let tcfg = TrainConfig {
+        seq_len,
+        batch_size: 8,
+        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: Some(ckpt.clone()),
+        ckpt_every: 500,
+        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
+    );
+    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}");
+
+    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}");
+    }
+}

crates/xtrain-train/src/checkpoint.rs

@@ -0,0 +1,90 @@
+//! Checkpoint save/load. Dumps the model's `params()` (in their stable order) to
+//! a flat binary file and reloads them into a model with matching architecture.
+//!
+//! Format (little-endian):
+//! ```text
+//!   magic   : u32  = 0x58545254  ("XTRT")
+//!   version : u32  = 1
+//!   n_params: u32
+//!   repeat n_params times:
+//!     ndim  : u32
+//!     dims  : [u32; ndim]
+//!     data  : [f32; prod(dims)]
+//! ```
+//! Architecture/config is NOT stored here — the caller rebuilds the model from
+//! the same `Config` and `load_into`s the params (the round-trip and resume both
+//! know their config). Gated behind `not(no_cuda)` (it round-trips GPU tensors).
+
+#![cfg(not(no_cuda))]
+
+use std::fs::File;
+use std::io::{BufReader, BufWriter, Read, Write};
+use std::path::Path;
+use xtrain_autodiff::tape::Var;
+use xtrain_tensor::{Device, Tensor};
+
+const MAGIC: u32 = 0x5854_5254;
+const VERSION: u32 = 1;
+
+/// Write every parameter (value) to `path` in `params()` order.
+pub fn save(path: &Path, params: &[Var]) -> std::io::Result<()> {
+    let mut w = BufWriter::new(File::create(path)?);
+    w.write_all(&MAGIC.to_le_bytes())?;
+    w.write_all(&VERSION.to_le_bytes())?;
+    w.write_all(&(params.len() as u32).to_le_bytes())?;
+    for p in params {
+        let v = p.value().to_device(Device::Cpu);
+        let shape = v.shape();
+        w.write_all(&(shape.len() as u32).to_le_bytes())?;
+        for &d in shape {
+            w.write_all(&(d as u32).to_le_bytes())?;
+        }
+        for &x in v.as_slice::<f32>() {
+            w.write_all(&x.to_le_bytes())?;
+        }
+    }
+    w.flush()
+}
+
+/// Read a checkpoint and overwrite each parameter's value in `params` (in order).
+/// Shapes must match the saved ones. Tensors are placed on each param's device.
+pub fn load_into(path: &Path, params: &[Var]) -> std::io::Result<()> {
+    let mut r = BufReader::new(File::open(path)?);
+    assert_eq!(read_u32(&mut r)?, MAGIC, "bad checkpoint magic");
+    assert_eq!(read_u32(&mut r)?, VERSION, "unsupported checkpoint version");
+    let n = read_u32(&mut r)? as usize;
+    assert_eq!(n, params.len(), "checkpoint param count != model");
+
+    for p in params {
+        let ndim = read_u32(&mut r)? as usize;
+        let mut dims = Vec::with_capacity(ndim);
+        for _ in 0..ndim {
+            dims.push(read_u32(&mut r)? as usize);
+        }
+        let numel: usize = dims.iter().product();
+        let mut data = vec![0.0f32; numel];
+        for slot in data.iter_mut() {
+            *slot = read_f32(&mut r)?;
+        }
+        let device = p.value().device();
+        assert_eq!(
+            p.value().shape(),
+            dims.as_slice(),
+            "checkpoint shape mismatch"
+        );
+        p.set_value(Tensor::from_slice(&data, &dims).to_device(device));
+    }
+    Ok(())
+}
+
+fn read_u32<R: Read>(r: &mut R) -> std::io::Result<u32> {
+    let mut b = [0u8; 4];
+    r.read_exact(&mut b)?;
+    Ok(u32::from_le_bytes(b))
+}
+
+fn read_f32<R: Read>(r: &mut R) -> std::io::Result<f32> {
+    let mut b = [0u8; 4];
+    r.read_exact(&mut b)?;
+    Ok(f32::from_le_bytes(b))
+}

crates/xtrain-train/src/clip.rs

@@ -0,0 +1,93 @@
+//! Global-norm gradient clipping. The norm is computed across *all* parameter
+//! gradients jointly (the same as `torch.nn.utils.clip_grad_norm_`): if the total
+//! L2 norm exceeds `max_norm`, every gradient is scaled by `max_norm / total`.
+//!
+//! The norm math is host-only and testable; [`clip_grad_norm`] wraps it over the
+//! parameter `Var`s (GPU round-trip), gated behind `not(no_cuda)`.
+
+/// Compute the global L2 norm over a set of flat gradient buffers.
+pub fn global_l2_norm(grads: &[Vec<f32>]) -> f32 {
+    let mut sumsq = 0.0f64;
+    for g in grads {
+        for &x in g {
+            sumsq += (x as f64) * (x as f64);
+        }
+    }
+    sumsq.sqrt() as f32
+}
+
+/// The scale factor to apply for clipping to `max_norm` (1.0 if already under).
+pub fn clip_scale(total_norm: f32, max_norm: f32) -> f32 {
+    if total_norm > max_norm && total_norm > 0.0 {
+        max_norm / total_norm
+    } else {
+        1.0
+    }
+}
+
+#[cfg(not(no_cuda))]
+mod gpu {
+    use super::{clip_scale, global_l2_norm};
+    use xtrain_autodiff::tape::Var;
+    use xtrain_tensor::{Device, Tensor};
+
+    /// First multiply every parameter's `.grad()` by `pre_scale` (use `1/batch`
+    /// to turn accumulated summed grads into a batch mean; `1.0` for no-op), then
+    /// clip the result to a joint global L2 norm of `max_norm`, writing the final
+    /// grads back via the tape's grad slot. Returns the post-pre_scale total norm
+    /// (the value the clip threshold is compared against, handy for logging).
+    /// Parameters without a grad contribute 0.
+    pub fn clip_grad_norm(params: &[Var], max_norm: f32, pre_scale: f32) -> f32 {
+        let device = params[0].value().device();
+        let grads: Vec<Option<Vec<f32>>> = params
+            .iter()
+            .map(|p| {
+                p.grad()
+                    .map(|g| g.to_device(Device::Cpu).as_slice::<f32>().to_vec())
+            })
+            .collect();
+
+        // Norm is measured on the (pre_scale-applied) grads — what the optimizer
+        // will actually see if no clipping is needed.
+        let scaled_present: Vec<Vec<f32>> = grads
+            .iter()
+            .flatten()
+            .map(|g| g.iter().map(|x| x * pre_scale).collect())
+            .collect();
+        let total = global_l2_norm(&scaled_present);
+        let factor = pre_scale * clip_scale(total, max_norm);
+        if (factor - 1.0).abs() < f32::EPSILON {
+            return total; // pre_scale==1 and under threshold → grads untouched
+        }
+
+        for (p, g) in params.iter().zip(&grads) {
+            if let Some(g) = g {
+                let shape = p.grad().unwrap().shape().to_vec();
+                let scaled: Vec<f32> = g.iter().map(|x| x * factor).collect();
+                let t = Tensor::from_slice(&scaled, &shape).to_device(device);
+                p.zero_grad();
+                Var::push_grad(p, t);
+            }
+        }
+        total
+    }
+}
+
+#[cfg(not(no_cuda))]
+pub use gpu::clip_grad_norm;
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn norm_and_scale() {
+        // grads = [3,4] → norm 5.
+        let g = vec![vec![3.0f32], vec![4.0]];
+        assert!((global_l2_norm(&g) - 5.0).abs() < 1e-6);
+        // Clip to 2.5 → scale 0.5.
+        assert!((clip_scale(5.0, 2.5) - 0.5).abs() < 1e-6);
+        // Already under → no scaling.
+        assert!((clip_scale(5.0, 10.0) - 1.0).abs() < 1e-6);
+    }
+}

crates/xtrain-train/src/data.rs

@@ -0,0 +1,75 @@
+//! Data pipeline: load the GPT-2 BPE (reusing xserv's from-scratch tokenizer),
+//! tokenize a text corpus into one flat token stream, and sample fixed-length
+//! `(input, target)` windows for next-token prediction. Host-only (no GPU).
+
+use std::path::Path;
+use xserv_tokenizer::Tokenizer;
+
+/// A tokenized corpus: one flat stream of token ids, plus the vocab size.
+pub struct Corpus {
+    pub tokens: Vec<i32>,
+    pub vocab_size: usize,
+}
+
+impl Corpus {
+    /// Load `tokenizer.json` (GPT-2 BPE) and tokenize the UTF-8 text at
+    /// `corpus_path` into a single id stream. TinyStories separates stories with
+    /// `<|endoftext|>`; the GPT-2 tokenizer emits that as a single special token,
+    /// so document boundaries are preserved in the stream.
+    pub fn load(tokenizer_path: &Path, corpus_path: &Path) -> Self {
+        let tok = Tokenizer::from_file(tokenizer_path);
+        let text = std::fs::read_to_string(corpus_path)
+            .unwrap_or_else(|e| panic!("failed to read corpus {}: {e}", corpus_path.display()));
+        // The range-fetched corpus may start/end mid-story; drop a leading partial
+        // line and a trailing partial story so we only train on whole sentences.
+        let text = trim_to_whole_stories(&text);
+        let ids: Vec<i32> = tok.encode(text).into_iter().map(|t| t as i32).collect();
+        Self {
+            tokens: ids,
+            vocab_size: tok.vocab_size(),
+        }
+    }
+
+    /// Total number of tokens.
+    pub fn len(&self) -> usize {
+        self.tokens.len()
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.tokens.is_empty()
+    }
+
+    /// Sample one `(input, target)` pair of length `seq` for next-token
+    /// prediction: a window `[s, s+seq+1)` → input `[s, s+seq)`, target shifted
+    /// by one. `rng_state` is advanced (a tiny LCG, so sampling is reproducible
+    /// from a seed without pulling in an RNG crate).
+    pub fn sample(&self, seq: usize, rng_state: &mut u64) -> (Vec<i32>, Vec<i32>) {
+        assert!(self.tokens.len() > seq + 1, "corpus shorter than a window");
+        let max_start = self.tokens.len() - seq - 1;
+        let start = (next_rand(rng_state) % (max_start as u64 + 1)) as usize;
+        let input = self.tokens[start..start + seq].to_vec();
+        let target = self.tokens[start + 1..start + seq + 1].to_vec();
+        (input, target)
+    }
+}
+
+/// Drop a leading partial line (before the first newline) and everything after
+/// the last `<|endoftext|>` marker, so a byte-range download still yields only
+/// complete stories. Falls back to the raw text if no marker is present.
+fn trim_to_whole_stories(text: &str) -> &str {
+    let start = text.find('\n').map(|i| i + 1).unwrap_or(0);
+    let body = &text[start..];
+    match body.rfind("<|endoftext|>") {
+        Some(end) => &body[..end + "<|endoftext|>".len()],
+        None => body,
+    }
+}
+
+/// Tiny LCG (same constants as the model tests' deterministic fill) so dataset
+/// sampling is reproducible from a single u64 seed.
+fn next_rand(state: &mut u64) -> u64 {
+    *state = state
+        .wrapping_mul(6364136223846793005)
+        .wrapping_add(1442695040888963407);
+    *state >> 16
+}

crates/xtrain-train/src/lib.rs

@@ -0,0 +1,22 @@
+//! Training stack (Phase T6): LR schedule, global-norm grad clipping, checkpoint
+//! save/load, the GPT-2 BPE data pipeline (reusing xserv's tokenizer), an
+//! autoregressive sampler, and the training loop that wires them onto the T5
+//! `TinyTransformer` + the hand-written AdamW (`xtrain-optim`).
+//!
+//! Host-only pieces (LR schedule, grad-norm math) always compile so the crate
+//! `cargo check`s on a GPU-less host; everything that touches GPU tensors is
+//! gated behind `not(no_cuda)`.
+
+pub mod clip;
+pub mod data;
+pub mod schedule;
+
+#[cfg(not(no_cuda))]
+pub mod checkpoint;
+#[cfg(not(no_cuda))]
+pub mod sample;
+#[cfg(not(no_cuda))]
+mod train_loop;
+
+#[cfg(not(no_cuda))]
+pub use train_loop::{TrainConfig, train};

crates/xtrain-train/src/sample.rs

@@ -0,0 +1,76 @@
+//! Autoregressive text sampling from the trained model. The model is
+//! single-sequence with RoPE position = row index, so generation re-runs the
+//! forward on the growing prefix each step and reads the last row's logits — the
+//! simplest correct approach (no KV cache; that is an inference/perf concern).
+//!
+//! Greedy when `temperature == 0`, else temperature sampling over the softmax.
+
+#![cfg(not(no_cuda))]
+
+use xtrain_model::{TinyTransformer, ids_tensor};
+use xtrain_tensor::Device;
+
+/// Generate `max_new` tokens continuing `prompt`. `temperature == 0` → greedy
+/// argmax; otherwise sample from softmax(logits / temperature). `rng_state` is a
+/// reproducible LCG seed (only used when temperature > 0).
+pub fn generate(
+    model: &TinyTransformer,
+    device: Device,
+    prompt: &[i32],
+    max_new: usize,
+    temperature: f32,
+    rng_state: &mut u64,
+) -> Vec<i32> {
+    let vocab = model.config().vocab;
+    let mut ids: Vec<i32> = prompt.to_vec();
+
+    for _ in 0..max_new {
+        let ids_t = ids_tensor(&ids, device);
+        let logits = model.forward(&ids_t).value().to_device(Device::Cpu);
+        let lg = logits.as_slice::<f32>();
+        // Last row = next-token distribution for the current prefix.
+        let last = &lg[(ids.len() - 1) * vocab..ids.len() * vocab];
+
+        let next = if temperature <= 0.0 {
+            argmax(last)
+        } else {
+            sample_temperature(last, temperature, rng_state)
+        };
+        ids.push(next as i32);
+    }
+    ids
+}
+
+fn argmax(row: &[f32]) -> usize {
+    row.iter()
+        .enumerate()
+        .max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
+        .unwrap()
+        .0
+}
+
+fn sample_temperature(row: &[f32], temperature: f32, rng_state: &mut u64) -> usize {
+    // Softmax with temperature (numerically stable).
+    let max = row.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
+    let exps: Vec<f32> = row
+        .iter()
+        .map(|&x| ((x - max) / temperature).exp())
+        .collect();
+    let sum: f32 = exps.iter().sum();
+    let r = (next_rand(rng_state) as f32 / u32::MAX as f32) * sum;
+    let mut acc = 0.0;
+    for (i, &e) in exps.iter().enumerate() {
+        acc += e;
+        if acc >= r {
+            return i;
+        }
+    }
+    exps.len() - 1
+}
+
+fn next_rand(state: &mut u64) -> u32 {
+    *state = state
+        .wrapping_mul(6364136223846793005)
+        .wrapping_add(1442695040888963407);
+    (*state >> 32) as u32
+}

crates/xtrain-train/src/schedule.rs

@@ -0,0 +1,58 @@
+//! Learning-rate schedule: linear warmup → cosine decay to a floor. Host-only
+//! and pure (just arithmetic over the step index), so it unit-tests locally.
+
+/// Warmup-then-cosine LR schedule.
+///
+/// - steps `0..warmup`: linear ramp `0 → max_lr`
+/// - steps `warmup..total`: cosine decay `max_lr → min_lr`
+/// - steps `>= total`: clamped at `min_lr`
+#[derive(Clone, Copy, Debug)]
+pub struct LrSchedule {
+    pub max_lr: f32,
+    pub min_lr: f32,
+    pub warmup: usize,
+    pub total: usize,
+}
+
+impl LrSchedule {
+    /// LR for a 0-indexed step.
+    pub fn lr(&self, step: usize) -> f32 {
+        if step < self.warmup {
+            // Linear warmup; +1 so step 0 already takes a (tiny) nonzero LR.
+            return self.max_lr * (step + 1) as f32 / self.warmup.max(1) as f32;
+        }
+        if step >= self.total {
+            return self.min_lr;
+        }
+        let progress = (step - self.warmup) as f32 / (self.total - self.warmup).max(1) as f32;
+        let cosine = 0.5 * (1.0 + (std::f32::consts::PI * progress).cos()); // 1 → 0
+        self.min_lr + (self.max_lr - self.min_lr) * cosine
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn warmup_then_cosine_shape() {
+        let s = LrSchedule {
+            max_lr: 1.0,
+            min_lr: 0.1,
+            warmup: 10,
+            total: 100,
+        };
+        // Warmup ramps up and reaches the peak at the end of warmup.
+        assert!(s.lr(0) < s.lr(5));
+        assert!(s.lr(5) < s.lr(9));
+        assert!((s.lr(9) - 1.0).abs() < 1e-6);
+        // Just after warmup, near the peak; midway, near the midpoint.
+        assert!(s.lr(10) > 0.95);
+        let mid = s.lr(55); // progress ~0.5 → cosine ~0.5
+        assert!((mid - (0.1 + 0.9 * 0.5)).abs() < 0.02);
+        // Decays monotonically to the floor by the end.
+        assert!(s.lr(99) < s.lr(55));
+        assert!(s.lr(100) <= s.min_lr + 1e-6);
+        assert!(s.lr(1_000) <= s.min_lr + 1e-6);
+    }
+}

crates/xtrain-train/src/train_loop.rs

@@ -0,0 +1,107 @@
+//! The training loop: sample sequences → forward `loss` → backward → grad clip
+//! (with batch averaging) → AdamW step → zero grads; with an LR schedule,
+//! periodic loss logging, and periodic checkpointing.
+//!
+//! The T5 model is single-sequence, so a "batch" of `batch_size` sequences is
+//! handled by running forward+backward on each and letting the tape SUM their
+//! grads (its fan-out rule); the clip pass then multiplies by `1/batch_size` to
+//! recover the batch-mean gradient before clipping + the optimizer step.
+
+#![cfg(not(no_cuda))]
+
+use std::path::PathBuf;
+use std::time::Instant;
+
+use xtrain_model::{TinyTransformer, ids_tensor};
+use xtrain_optim::AdamW;
+use xtrain_tensor::Device;
+
+use crate::checkpoint;
+use crate::clip::clip_grad_norm;
+use crate::data::Corpus;
+use crate::schedule::LrSchedule;
+
+/// Knobs for a training run.
+pub struct TrainConfig {
+    pub seq_len: usize,
+    pub batch_size: usize,
+    pub steps: usize,
+    pub schedule: LrSchedule,
+    pub weight_decay: f32,
+    pub max_grad_norm: f32,
+    pub log_every: usize,
+    /// Optional checkpoint path written every `ckpt_every` steps (and at the end).
+    pub ckpt_path: Option<PathBuf>,
+    pub ckpt_every: usize,
+    /// Seed for reproducible sequence sampling.
+    pub seed: u64,
+}
+
+/// Train `model` on `corpus` for `cfg.steps` AdamW steps. Returns the per-step
+/// loss trace (one mean loss per step, read from the first sequence of the
+/// batch — cheap and representative). Logs progress and checkpoints as configured.
+pub fn train(
+    model: &TinyTransformer,
+    device: Device,
+    corpus: &Corpus,
+    cfg: &TrainConfig,
+) -> Vec<f32> {
+    let params = model.params();
+    let mut opt = AdamW::new(cfg.schedule.max_lr, cfg.weight_decay);
+    let mut rng = cfg.seed;
+    let mut losses = Vec::with_capacity(cfg.steps);
+    let inv_batch = 1.0 / cfg.batch_size as f32;
+    let start = Instant::now();
+    let mut tokens_seen: u64 = 0;
+
+    for step in 0..cfg.steps {
+        let lr = cfg.schedule.lr(step);
+
+        // Accumulate grads over `batch_size` sequences (tape SUMs them).
+        let mut step_loss = 0.0f32;
+        for _ in 0..cfg.batch_size {
+            let (input, target) = corpus.sample(cfg.seq_len, &mut rng);
+            let ids = ids_tensor(&input, device);
+            let targets = ids_tensor(&target, device);
+            let loss = model.loss(&ids, &targets);
+            step_loss += read_scalar(&loss);
+            loss.backward();
+            tokens_seen += cfg.seq_len as u64;
+        }
+        step_loss *= inv_batch;
+        losses.push(step_loss);
+
+        // Average the summed grads (×1/batch) and clip to the global norm.
+        let gnorm = clip_grad_norm(&params, cfg.max_grad_norm, inv_batch);
+        opt.step(lr, &params);
+        for p in &params {
+            p.zero_grad();
+        }
+
+        if step % cfg.log_every == 0 || step == cfg.steps - 1 {
+            let elapsed = start.elapsed().as_secs_f32();
+            let tps = tokens_seen as f32 / elapsed.max(1e-6);
+            println!(
+                "step {step:5}/{}: loss {step_loss:.4} lr {lr:.2e} gnorm {gnorm:.3} \
+                 ({tps:.0} tok/s)",
+                cfg.steps
+            );
+        }
+
+        if let Some(path) = &cfg.ckpt_path {
+            if cfg.ckpt_every > 0 && (step + 1) % cfg.ckpt_every == 0 {
+                checkpoint::save(path, &params).expect("checkpoint save");
+            }
+        }
+    }
+
+    if let Some(path) = &cfg.ckpt_path {
+        checkpoint::save(path, &params).expect("final checkpoint save");
+        println!("saved checkpoint → {}", path.display());
+    }
+    losses
+}
+
+fn read_scalar(v: &xtrain_autodiff::tape::Var) -> f32 {
+    v.value().to_device(Device::Cpu).as_slice::<f32>()[0]
+}

crates/xtrain-train/tests/adamw_parity.py

@@ -0,0 +1,180 @@
+#!/usr/bin/env python3
+"""AdamW-vs-PyTorch parity (Phase T6).
+
+Loads the model dumped by tests/adamw_parity_dump.rs (config, ids, initial
+params, the loss trajectory, and final params), rebuilds the IDENTICAL tiny
+transformer in PyTorch from the same initial weights, and runs the SAME number
+of `torch.optim.AdamW` steps with matched hyperparameters (lr, weight_decay,
+betas, eps) on the same fixed batch. It then compares:
+
+  * the per-step loss trajectory (Rust AdamW vs torch AdamW), and
+  * the final parameters,
+
+within a relative tolerance. A correct hand-written AdamW (bias correction +
+decoupled weight decay) tracks torch's optimizer step-for-step.
+
+Usage:  python3 adamw_parity.py /tmp/xtrain_adamw
+"""
+import sys
+import os
+import math
+import torch
+
+DIR = sys.argv[1] if len(sys.argv) > 1 else "/tmp/xtrain_adamw"
+
+
+def read_vec(name):
+    path = os.path.join(DIR, name)
+    shape = None
+    vals = []
+    with open(path) as f:
+        for line in f:
+            line = line.strip()
+            if line.startswith("# shape"):
+                shape = [int(x) for x in line.split()[2].split(",") if x]
+            elif line:
+                vals.append(float(line))
+    t = torch.tensor(vals, dtype=torch.float64)
+    if shape:
+        t = t.reshape(shape)
+    return t
+
+
+def read_cfg():
+    cfg = {}
+    with open(os.path.join(DIR, "config.txt")) as f:
+        for line in f:
+            k, v = line.split()
+            cfg[k] = v
+    return cfg
+
+
+def read_ids(name):
+    with open(os.path.join(DIR, name)) as f:
+        return [int(x) for x in f.read().split()]
+
+
+cfg = read_cfg()
+DIM = int(cfg["dim"])
+NL = int(cfg["n_layers"])
+NH = int(cfg["n_heads"])
+HD = int(cfg["head_dim"])
+EPS = float(cfg["eps"])
+THETA = float(cfg["rope_theta"])
+LR = float(cfg["lr"])
+WD = float(cfg["wd"])
+N_STEPS = int(cfg["n_steps"])
+
+ids = read_ids("ids.txt")
+targets = read_ids("targets.txt")
+SEQ = len(ids)
+
+NAMES = ["embed"]
+for l in range(NL):
+    for p in ["attn_norm", "wq", "wk", "wv", "wo",
+              "ffn_norm", "w_gate", "w_up", "w_down"]:
+        NAMES.append(f"l{l}_{p}")
+NAMES += ["final_norm", "lm_head"]
+
+# Load the IDENTICAL initial weights as leaf params (float64 reference).
+P = {n: read_vec(f"w0_{n}.txt").clone().requires_grad_(True) for n in NAMES}
+
+
+def rms_norm(x, gamma):
+    ms = x.pow(2).mean(dim=-1, keepdim=True)
+    return x * torch.rsqrt(ms + EPS) * gamma
+
+
+def rope(x):  # x: [seq, nh, hd], position = token index
+    half = HD // 2
+    out = torch.empty_like(x)
+    i = torch.arange(half, dtype=torch.float64)
+    freq = THETA ** (-(2.0 * i) / HD)
+    pos = torch.arange(SEQ, dtype=torch.float64).reshape(SEQ, 1)
+    ang = pos * freq
+    c = torch.cos(ang).reshape(SEQ, 1, half)
+    s = torch.sin(ang).reshape(SEQ, 1, half)
+    x0, x1 = x[..., :half], x[..., half:]
+    out[..., :half] = x0 * c - x1 * s
+    out[..., half:] = x1 * c + x0 * s
+    return out
+
+
+idx = torch.tensor(ids, dtype=torch.long)
+tgt = torch.tensor(targets, dtype=torch.long)
+mask = torch.triu(torch.full((SEQ, SEQ), -1.0e9, dtype=torch.float64), diagonal=1)
+
+
+def forward():
+    h = P["embed"][idx]
+    for l in range(NL):
+        x = rms_norm(h, P[f"l{l}_attn_norm"])
+        q = (x @ P[f"l{l}_wq"]).reshape(SEQ, NH, HD)
+        k = (x @ P[f"l{l}_wk"]).reshape(SEQ, NH, HD)
+        v = (x @ P[f"l{l}_wv"]).reshape(SEQ, NH, HD)
+        q = rope(q).transpose(0, 1)
+        k = rope(k).transpose(0, 1)
+        v = v.transpose(0, 1)
+        scale = 1.0 / math.sqrt(HD)
+        scores = (q @ k.transpose(-1, -2)) * scale + mask
+        probs = torch.softmax(scores, dim=-1)
+        out = (probs @ v).transpose(0, 1).reshape(SEQ, DIM)
+        h = h + out @ P[f"l{l}_wo"]
+        x = rms_norm(h, P[f"l{l}_ffn_norm"])
+        act = torch.nn.functional.silu(x @ P[f"l{l}_w_gate"]) * (x @ P[f"l{l}_w_up"])
+        h = h + act @ P[f"l{l}_w_down"]
+    h = rms_norm(h, P["final_norm"])
+    return h @ P["lm_head"]
+
+
+# Match the Rust optimizer: torch.optim.AdamW with the same lr/wd/betas/eps.
+opt = torch.optim.AdamW(list(P.values()), lr=LR, betas=(0.9, 0.999),
+                        eps=1e-8, weight_decay=WD)
+
+torch_losses = []
+for _ in range(N_STEPS):
+    opt.zero_grad()
+    logits = forward()
+    loss = torch.nn.functional.cross_entropy(logits, tgt, reduction="mean")
+    torch_losses.append(loss.item())
+    loss.backward()
+    opt.step()
+
+
+def relerr(a, b):
+    a, b = a.double(), b.double()
+    denom = b.abs().clamp(min=1e-6)
+    return ((a - b).abs() / denom).max().item()
+
+
+rust_losses = read_vec("losses.txt")
+print("step  rust_loss     torch_loss    relerr")
+worst_loss = 0.0
+for i in range(N_STEPS):
+    rl, tl = rust_losses[i].item(), torch_losses[i]
+    e = abs(rl - tl) / max(abs(tl), 1e-6)
+    worst_loss = max(worst_loss, e)
+    if i < 5 or i == N_STEPS - 1:
+        print(f"{i:4d}  {rl:.6e}  {tl:.6e}  {e:.2e}")
+print(f"loss trajectory: worst relerr = {worst_loss:.2e}")
+
+RTOL = 2e-2
+worst_p, worst_name = 0.0, ""
+fails = []
+for n in NAMES:
+    ref = read_vec(f"wN_{n}.txt")
+    e = relerr(P[n].detach(), ref)
+    if e > worst_p:
+        worst_p, worst_name = e, n
+    if e > RTOL:
+        fails.append((n, e))
+print(f"final params: {len(NAMES)} checked, worst = {worst_name} @ {worst_p:.2e} (rtol={RTOL})")
+
+if worst_loss > RTOL or fails:
+    print("FAIL:")
+    if worst_loss > RTOL:
+        print(f"  loss trajectory relerr {worst_loss:.3e} > {RTOL}")
+    for n, e in fails:
+        print(f"  param[{n}]: relerr={e:.3e}")
+    sys.exit(1)
+print("ADAMW PARITY OK: loss trajectory + final params match torch.optim.AdamW within rtol")

crates/xtrain-train/tests/adamw_parity_dump.rs

@@ -0,0 +1,165 @@
+// AdamW-vs-PyTorch parity, step 1 of 2: build the tiny transformer with a fixed
+// deterministic init, then run N steps of the hand-written AdamW on a FIXED
+// (input, target) batch — recording the loss at each step and the final
+// parameters. tests/adamw_parity.py rebuilds the identical model + torch.optim
+// .AdamW with matched hyperparameters and compares the loss trajectory and final
+// params within rtol. This is the rigorous correctness check for the optimizer.
+//
+// Run:  XTRAIN_ADAMW_DIR=/tmp/xtrain_adamw cargo test -p xtrain-train \
+//         --test adamw_parity_dump -- --nocapture --ignored
+// then: python3 crates/xtrain-train/tests/adamw_parity.py /tmp/xtrain_adamw
+//
+// Marked #[ignore] (fixture generator) and gated #![cfg(not(no_cuda))].
+#![cfg(not(no_cuda))]
+
+use std::fs;
+use std::io::Write;
+use std::path::PathBuf;
+
+use xtrain_cuda::device;
+use xtrain_model::{Config, TinyTransformer, ids_tensor, param_to_host};
+use xtrain_optim::AdamW;
+use xtrain_tensor::Device;
+
+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()
+}
+
+fn write_vec(dir: &PathBuf, name: &str, data: &[f32], shape: &[usize]) {
+    let mut f = fs::File::create(dir.join(name)).unwrap();
+    let shape_str: Vec<String> = shape.iter().map(|d| d.to_string()).collect();
+    writeln!(f, "# shape {}", shape_str.join(",")).unwrap();
+    for v in data {
+        writeln!(f, "{v:.8e}").unwrap();
+    }
+}
+
+const LR: f32 = 0.01;
+const WD: f32 = 0.1;
+const N_STEPS: usize = 30;
+
+#[test]
+#[ignore = "fixture generator for AdamW PyTorch parity; run with --ignored"]
+fn dump_adamw_trajectory() {
+    assert!(device::device_count().unwrap() > 0, "no CUDA device");
+    device::set_device(0).unwrap();
+    let device = Device::Cuda(0);
+
+    let dir = PathBuf::from(
+        std::env::var("XTRAIN_ADAMW_DIR").unwrap_or_else(|_| "/tmp/xtrain_adamw".to_string()),
+    );
+    fs::create_dir_all(&dir).unwrap();
+
+    let mut cfg = Config::tiny();
+    cfg.vocab = 12;
+    let ids: Vec<i32> = vec![3, 1, 4, 1, 5, 9, 2, 6];
+    let targets: Vec<i32> = vec![1, 4, 1, 5, 9, 2, 6, 0];
+
+    // Same deterministic init the parity dump uses (so the torch side can reuse it).
+    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.08)
+        }
+    });
+
+    // Dump config + ids + initial params (named for adamw_parity.py).
+    {
+        let mut f = fs::File::create(dir.join("config.txt")).unwrap();
+        writeln!(f, "vocab {}", cfg.vocab).unwrap();
+        writeln!(f, "dim {}", cfg.dim).unwrap();
+        writeln!(f, "n_layers {}", cfg.n_layers).unwrap();
+        writeln!(f, "n_heads {}", cfg.n_heads).unwrap();
+        writeln!(f, "head_dim {}", cfg.head_dim).unwrap();
+        writeln!(f, "ffn_hidden {}", cfg.ffn_hidden).unwrap();
+        writeln!(f, "eps {:e}", cfg.eps).unwrap();
+        writeln!(f, "rope_theta {:e}", cfg.rope_theta).unwrap();
+        writeln!(f, "lr {LR:e}").unwrap();
+        writeln!(f, "wd {WD:e}").unwrap();
+        writeln!(f, "n_steps {N_STEPS}").unwrap();
+        let mut g = fs::File::create(dir.join("ids.txt")).unwrap();
+        for v in &ids {
+            writeln!(g, "{v}").unwrap();
+        }
+        let mut g = fs::File::create(dir.join("targets.txt")).unwrap();
+        for v in &targets {
+            writeln!(g, "{v}").unwrap();
+        }
+    }
+
+    let names = param_names(&cfg);
+    let params = model.params();
+    for (name, p) in names.iter().zip(&params) {
+        let shape = p.value().shape().to_vec();
+        write_vec(&dir, &format!("w0_{name}.txt"), &param_to_host(p), &shape);
+    }
+
+    // Train N steps of AdamW with a CONSTANT lr (no schedule) on the fixed batch.
+    let ids_t = ids_tensor(&ids, device);
+    let targets_t = ids_tensor(&targets, device);
+    let mut opt = AdamW::new(LR, WD);
+    let mut losses = Vec::with_capacity(N_STEPS);
+    for _ in 0..N_STEPS {
+        let loss = model.loss(&ids_t, &targets_t);
+        losses.push(param_to_host(&loss)[0]);
+        loss.backward();
+        opt.step(LR, &params);
+        for p in &params {
+            p.zero_grad();
+        }
+    }
+
+    {
+        let mut f = fs::File::create(dir.join("losses.txt")).unwrap();
+        for l in &losses {
+            writeln!(f, "{l:.8e}").unwrap();
+        }
+    }
+    for (name, p) in names.iter().zip(&params) {
+        let shape = p.value().shape().to_vec();
+        write_vec(&dir, &format!("wN_{name}.txt"), &param_to_host(p), &shape);
+    }
+
+    println!(
+        "adamw parity: dumped to {} (loss {:.6e} → {:.6e} over {N_STEPS} steps)",
+        dir.display(),
+        losses.first().unwrap(),
+        losses.last().unwrap()
+    );
+}
+
+fn param_names(cfg: &Config) -> Vec<String> {
+    let mut names = vec!["embed".to_string()];
+    for l in 0..cfg.n_layers {
+        for p in [
+            "attn_norm",
+            "wq",
+            "wk",
+            "wv",
+            "wo",
+            "ffn_norm",
+            "w_gate",
+            "w_up",
+            "w_down",
+        ] {
+            names.push(format!("l{l}_{p}"));
+        }
+    }
+    names.push("final_norm".to_string());
+    names.push("lm_head".to_string());
+    names
+}

crates/xtrain-train/tests/checkpoint_roundtrip.rs

@@ -0,0 +1,113 @@
+// Checkpoint round-trip acceptance (Phase T6): train a few AdamW steps on a fixed
+// batch, save the params, build a FRESH model (different init), load the
+// checkpoint into it, and assert it produces identical logits + loss on a fixed
+// input. This verifies the on-disk format dumps/reloads `params()` in order with
+// exact f32 fidelity. Gated #![cfg(not(no_cuda))] (runs on dash5).
+#![cfg(not(no_cuda))]
+
+use xtrain_cuda::device;
+use xtrain_model::{Config, TinyTransformer, ids_tensor, param_to_host};
+use xtrain_optim::AdamW;
+use xtrain_tensor::Device;
+use xtrain_train::checkpoint;
+
+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()
+}
+
+fn make_model(device: Device, vocab: usize, init_seed: u64) -> TinyTransformer {
+    let mut cfg = Config::tiny();
+    cfg.vocab = vocab;
+    let mut seed = init_seed;
+    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.08)
+        }
+    })
+}
+
+#[test]
+fn checkpoint_roundtrip_identical_logits() {
+    assert!(device::device_count().unwrap() > 0, "no CUDA device");
+    device::set_device(0).unwrap();
+    let device = Device::Cuda(0);
+
+    let vocab = 12;
+    let ids: Vec<i32> = vec![3, 1, 4, 1, 5, 9, 2, 6];
+    let targets: Vec<i32> = vec![1, 4, 1, 5, 9, 2, 6, 0];
+    let ids_t = ids_tensor(&ids, device);
+    let targets_t = ids_tensor(&targets, device);
+
+    // --- Train a few steps so the params are non-trivial (not the init). ---
+    let model = make_model(device, vocab, 1);
+    let params = model.params();
+    let mut opt = AdamW::new(0.01, 0.1);
+    for _ in 0..5 {
+        let loss = model.loss(&ids_t, &targets_t);
+        loss.backward();
+        opt.step(0.01, &params);
+        for p in &params {
+            p.zero_grad();
+        }
+    }
+
+    let path = std::env::temp_dir().join(format!("xtrain_ckpt_{}.bin", std::process::id()));
+    checkpoint::save(&path, &params).unwrap();
+
+    let ref_logits = param_to_host(&model.forward(&ids_t));
+    let ref_loss = param_to_host(&model.loss(&ids_t, &targets_t))[0];
+
+    // --- Fresh model with a DIFFERENT init; loading must overwrite it exactly. ---
+    let fresh = make_model(device, vocab, 999);
+    let fresh_params = fresh.params();
+    // Sanity: before load, the fresh model disagrees.
+    let pre = param_to_host(&fresh.forward(&ids_t));
+    let pre_diff: f32 = pre
+        .iter()
+        .zip(&ref_logits)
+        .map(|(a, b)| (a - b).abs())
+        .fold(0.0, f32::max);
+    assert!(
+        pre_diff > 1e-4,
+        "fresh model unexpectedly matched before load"
+    );
+
+    checkpoint::load_into(&path, &fresh_params).unwrap();
+
+    let got_logits = param_to_host(&fresh.forward(&ids_t));
+    let got_loss = param_to_host(&fresh.loss(&ids_t, &targets_t))[0];
+    let _ = std::fs::remove_file(&path);
+
+    // Exact f32 round-trip → bit-for-bit identical forward (same kernels, same
+    // inputs). Allow only float noise from re-running the forward.
+    let max_logit_diff: f32 = got_logits
+        .iter()
+        .zip(&ref_logits)
+        .map(|(a, b)| (a - b).abs())
+        .fold(0.0, f32::max);
+    println!(
+        "checkpoint round-trip: max logit diff = {max_logit_diff:.3e}, loss {ref_loss:.6} vs {got_loss:.6}"
+    );
+    assert!(
+        max_logit_diff < 1e-5,
+        "logits differ after reload: {max_logit_diff:e}"
+    );
+    assert!(
+        (got_loss - ref_loss).abs() < 1e-5,
+        "loss differs after reload: {ref_loss} vs {got_loss}"
+    );
+}

crates/xtrain-train/tests/real_training.rs

@@ -0,0 +1,121 @@
+// 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()),
+    );
+    let corpus_path = PathBuf::from(
+        std::env::var("XTRAIN_CORPUS").unwrap_or_else(|_| "data/tinystories-valid-3mb.txt".into()),
+    );
+
+    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,
+        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,
+        seed: 42,
+    };
+
+    let losses = train(&model, device, &corpus, &tcfg);
+    // 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}");
+}

docs/05-training-loop.md

@@ -0,0 +1,114 @@
+# Phase T6: Training Loop + AdamW + Real Training — Design Document
+
+## Goal
+
+在 T5 的 `TinyTransformer`（`params()` / `forward` / `loss` + `Var::{value,grad,set_value,zero_grad}`）之上，搭起**真正的训练栈**，并在**真实文本语料**上把 loss 训下来：
+
+1. **手写 AdamW**：per-param 一/二阶矩（m、v），bias correction，**decoupled weight decay**，对拍 `torch.optim.AdamW` 数值一致。
+2. **训练 loop**：语料 → 采样定长序列 → `forward(loss)` → `backward` → **global-norm grad clip** → **AdamW step** → `zero_grad`；**LR schedule**（warmup + cosine）；周期性 loss 日志；**checkpoint 存/取**。
+3. **采样器**：greedy / temperature，训练中/后吐文本看「在不在学」。
+4. **数据**：**复用 xserv 的 GPT-2 BPE** tokenizer（path-dep），语料 = **TinyStories** 子集。
+
+**不做**（留后续 Phase）：性能（cuBLAS 切换 / bf16 / 激活重计算 = T7）、分布式（NCCL 数据并行 = T8）。本 Phase 只要**正确性 + 清晰的学习信号**，训练预算有界（几分钟 / 几千步，非完全收敛）。
+
+## Module Layout
+
+```
+crates/xtrain-optim/                   # 新 crate：优化器
+├── build.rs                           # 检测 nvcc → no_cuda cfg（逐 crate）
+├── src/lib.rs                         # AdamW：step_host(纯 host 数学) + step(&[Var]) GPU 包装
+└── tests/adamw_host.rs               # host 单测：对独立参考递推 + 纯 decay 边界（本地可跑，无 GPU）
+
+crates/xtrain-train/                   # 新 crate：训练基建 + 入口
+├── build.rs                           # 检测 nvcc → no_cuda cfg
+├── Cargo.toml                         # path-dep: ../../../xserv/crates/xserv-tokenizer（本地/dash5 都解析）
+├── src/
+│   ├── lib.rs                         # 模块导出（host-only 与 GPU 件分门控）
+│   ├── schedule.rs                    # LrSchedule：warmup + cosine（host-only，可本地单测）
+│   ├── clip.rs                        # global L2 norm + clip_scale（host 数学）+ clip_grad_norm(&[Var])（GPU 门控）
+│   ├── data.rs                        # Corpus：load tokenizer+语料 → token 流 → sample(input,target) 窗口
+│   ├── checkpoint.rs                  # save / load_into：按 params() 顺序 dump/reload（GPU 门控）
+│   ├── sample.rs                      # generate：greedy / temperature 自回归采样（GPU 门控）
+│   ├── train_loop.rs                  # TrainConfig + train()：把以上接到 model+AdamW（GPU 门控）
+│   └── bin/train.rs                   # 真训练入口：load 数据 → train → checkpoint → 采样
+└── tests/
+    ├── adamw_parity_dump.rs           # AdamW 对拍 fixture：固定 init 跑 N 步 AdamW，dump loss 轨迹 + 终参
+    ├── adamw_parity.py                # 等价 PyTorch 模型 + torch.optim.AdamW，对比轨迹 + 终参
+    ├── checkpoint_roundtrip.rs        # 训几步→save→载入新模型→logits/loss 逐位一致
+    └── real_training.rs              # TinyStories 有界训练：loss 大幅下降 + 采样在学
+
+data/tinystories-valid-3mb.txt         # 语料子集（committed，~3MB，TinyStories-valid 前 3MB，整故事截断）
+```
+
+**为什么拆两个 crate**：对齐 xserv 的分层（优化器与训练编排分开）。`xtrain-optim` 只管参数更新数学；`xtrain-train` 管数据/调度/checkpoint/采样/loop。AdamW 数学独立可测，不依赖 model。
+
+**host / GPU 门控约定**（沿用全仓）：纯算术（`LrSchedule`、grad-norm 数学、AdamW 的 `step_host`）**始终编译**，本地 `cargo check` + 单测即可验证；凡 round-trip GPU 张量的（`step(&[Var])`、`clip_grad_norm(&[Var])`、checkpoint、采样、loop）一律 `#[cfg(not(no_cuda))]`，链接+实跑在 dash5。每 crate 的 `build.rs` 各自检测 nvcc（cfg 不跨 crate 传播）。
+
+## Key Design Decisions
+
+### AdamW：手写数学 + decoupled weight decay
+
+第 `t` 步（1-indexed），参数 `θ`、梯度 `g`：
+
+```text
+  m ← β1·m + (1−β1)·g
+  v ← β2·v + (1−β2)·g²
+  m̂ ← m / (1 − β1ᵗ)          (bias correction)
+  v̂ ← v / (1 − β2ᵗ)
+  θ ← θ − lr·( m̂ / (√v̂ + ε) + wd·θ )
+```
+
+- **decoupled weight decay**（Loshchilov & Hutter 2019）：`wd·θ` 直接作用在参数上，**不**并进梯度（不进入自适应 `√v̂` 分母）——这正是 `torch.optim.AdamW` 的定义，区别于「L2 正则把 `wd·θ` 加到 `g`」的 Adam。
+- 默认超参对齐 PyTorch：β1=0.9，β2=0.999，ε=1e-8。
+- **状态 keyed by 参数在 `params()` 中的下标**（稳定序），首次 `step` 惰性按各参数 numel 分配 `m,v`。`t` 全局共享（所有参数同一 bias correction，和 PyTorch 一致）。
+
+**实现分层**：`step_host(lr, &mut [Vec<f32>], &[Vec<f32>])` 是纯 host f32 数学（无 GPU、无 autograd，本地单测）；`step(lr, &[Var])` 把每参数的 `value()`/`grad()` 拉到 host、调 `step_host`、`set_value` 写回。这条路子（host 算优化器）对 tiny 模型完全够用，且让 AdamW 数学**脱离 GPU 可严格对拍**——T6 是正确性 Phase，不做 GPU 优化器 kernel（那是性能向，超范围）。`lr` 每步传入，给 schedule 留口。
+
+### LR schedule：warmup + cosine
+
+`step ∈ [0,warmup)` 线性 `0→max_lr`；`[warmup,total)` cosine `max_lr→min_lr`；`≥total` 钳到 `min_lr`。纯函数（只吃 step 下标），本地单测形状。
+
+### grad clip：global L2 norm（+ batch 平均）
+
+跨**所有**参数梯度联合算 L2 norm（同 `torch.nn.utils.clip_grad_norm_`）：`total > max_norm` 则全体 `×(max_norm/total)`。
+
+模型是**单序列**（无 batch 维），一个 `batch_size` 的「batch」靠**跑 `batch_size` 次 forward+backward**、让 tape 的 fan-out 规则**把梯度 SUM** 起来实现。为得到 batch 均值梯度，clip 这一趟 host pass 里**先 `×1/batch_size`** 再算 norm/裁剪——`clip_grad_norm(params, max_norm, pre_scale)` 把「平均」与「裁剪」融成一次 host 往返（省一趟拷贝）。batch 是 T7/边角关切，这里只求正确。
+
+### checkpoint 格式
+
+按 `params()` 顺序 dump 每个参数的 value 到扁平二进制：
+
+```text
+  magic u32 = "XTRT" | version u32 | n_params u32
+  ×n_params: ndim u32 | dims[ndim] u32 | data[Πdims] f32   (小端)
+```
+
+不存架构/config——调用方用同一 `Config` 重建模型再 `load_into`（round-trip 与 resume 都自知 config）。`load_into` 校验 magic/version/数量/逐参数 shape，按各参数 device 写回 `set_value`。f32 精确往返 → 重载后 forward 逐位一致（同 kernel 同输入）。
+
+### 数据管线 + tokenizer 复用
+
+- **tokenizer = 复用 xserv 的 from-scratch GPT-2/Qwen BPE**：`Cargo.toml` path-dep `../../../xserv/crates/xserv-tokenizer`，该相对路径在本地 `~/projects` 与 dash5 `/opt/wjh/projects` 都解析；Cargo 按目标 crate 自身的 workspace（xserv 的）解析它的 `serde/regex` 依赖，不需要 xtrain 复制 workspace dep。加载 `/opt/wjh/models/gpt2/tokenizer.json`。
+- **语料 = TinyStories 子集**：dash5 经 `hf-mirror.com` 取 `TinyStories-valid.txt` 前 ~3MB（HF 直连不可达，proxy 脚本只起后台 SOCKS；hf-mirror 直连 200），committed 进 `data/`。`Corpus::load` 整篇 tokenize 成一条 token 流（TinyStories 用 `<|endoftext|>` 分故事，GPT-2 BPE 正好出成单个 special token，文档边界保留）；range 下载会掐头去尾，故先丢首个不完整行、截到最后一个 `<|endoftext|>`，只训整故事。`sample(seq)` 随机取窗口 `[s,s+seq+1)` → input `[s,s+seq)` / target 右移一位（next-token），LCG 种子可复现，不引 RNG crate。
+
+### 采样器
+
+模型单序列、RoPE pos=行号，故自回归生成**每步对增长前缀重跑 forward、取末行 logits**（最简正确法；KV cache 是推理/性能向，超范围）。`temperature==0` greedy argmax，否则按 `softmax(logits/T)` 采样。
+
+### 训练 loop（`train`）
+
+每步：采 `batch_size` 序列各自 forward `loss` + backward（tape SUM 梯度）→ `clip_grad_norm(×1/batch + 裁剪)` → `AdamW::step(lr)` → 全参数 `zero_grad`；按 `log_every` 打 `loss/lr/gnorm/tok-s`，按 `ckpt_every` 存 checkpoint，返回逐步 loss 轨迹。
+
+## 验证方法（验收）
+
+GPU 测试全部 `#[cfg(not(no_cuda))]` 门控，在 dash5 实跑 capture：
+
+1. **AdamW 对拍 PyTorch**（严格正确性）：同一 tiny 模型 + 相同 init，Rust AdamW 与 `torch.optim.AdamW`（lr/wd/betas/eps 全对齐）各跑 N 步固定 batch → **loss 轨迹**与**终参**逐项 rtol 内一致。
+   - fixture：`cargo test -p xtrain-train --test adamw_parity_dump -- --ignored --nocapture`
+   - 对比：`python3 crates/xtrain-train/tests/adamw_parity.py /tmp/xtrain_adamw`
+2. **checkpoint round-trip**：训几步 → save → 载入**全新 init 的模型** → 固定输入 logits/loss 逐位一致（且证明载入前新模型确实不同）。
+   - `cargo test -p xtrain-train --test checkpoint_roundtrip`
+3. **真训练**（端到端学习信号）：TinyStories 上有界训练（几百~几千步）→ loss 大幅下降 + greedy 采样显出英文结构（非乱码）。
+   - `cargo test -p xtrain-train --release --test real_training -- --ignored --nocapture`
+   - 或 `cargo run -p xtrain-train --release --bin train -- <tokenizer.json> <corpus.txt> [steps] [ckpt]`
+4. **host 单测**（本地即跑）：AdamW 数学对独立参考递推、LR schedule 形状、grad-norm/clip 数学。
+   - `cargo test -p xtrain-optim -p xtrain-train`