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