test: AdamW PyTorch parity + checkpoint round-trip + real training

Acceptance tests (GPU-gated not(no_cuda), run on dash5):
- adamw_parity_dump.rs + adamw_parity.py: build the tiny model with fixed init,
  run N AdamW steps on a fixed batch, dump the loss trajectory + final params;
  the Python side rebuilds the identical model and runs torch.optim.AdamW with
  matched lr/wd/betas/eps, comparing trajectory + final params within rtol.
- checkpoint_roundtrip.rs: train a few steps, save, load into a fresh model with
  a DIFFERENT init, assert identical logits/loss on a fixed input.
- real_training.rs (#[ignore], --release): train on TinyStories for a bounded
  budget; assert loss drops substantially and print greedy samples.

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

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}");
+}