xtrain/crates/xtrain-train/tests/checkpoint_roundtrip.rs

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