xtrain/crates/xtrain-model/tests/recompute.rs

// T13 activation-recomputation correctness gate (the HARD gate).
//
// Gradient checkpointing is mathematically EXACT: the backward re-runs the same
// `segment_fn` from the same saved input and the same (unchanged) parameter
// values, so the recomputed activations equal the originals and the recovered
// grads equal the non-checkpointed grads — checkpointing trades compute for
// memory, never correctness. This test makes that a closed loop on-GPU:
//
//   build two identical models (same init), one with `--recompute` on, one off,
//   run the SAME batched loss + backward on both, and assert
//     1. the forward logits match (recompute doesn't touch forward output)
//     2. the loss matches
//     3. EVERY parameter's grad matches within a tight fp tolerance.
//
// Composition is covered by parameterising over fp32 AND bf16 (T12): the
// recompute path is the unchanged block forward, so it runs the same dtype path.
#![cfg(not(no_cuda))]

use xtrain_cuda::device;
use xtrain_model::{Config, TinyTransformer, batched_ids_tensor};
use xtrain_tensor::{DType, 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 build(cfg: Config, device: Device, dtype: DType, recompute: bool) -> TinyTransformer {
    let mut seed = 1u64;
    let m = 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)
        }
    });
    m.with_compute_dtype(dtype).with_recompute(recompute)
}

/// Upcast to fp32 then read to host — logits are bf16 in bf16 mode (grads are
/// always fp32 master, but this is uniform and harmless for fp32 tensors).
fn host(t: &xtrain_tensor::Tensor) -> Vec<f32> {
    t.to_dtype(DType::F32)
        .to_device(Device::Cpu)
        .as_slice::<f32>()
        .to_vec()
}

fn run(dtype: DType, logit_tol: f32, grad_tol: f32) {
    assert!(device::device_count().unwrap() > 0, "no CUDA device");
    device::set_device(0).unwrap();
    let device = Device::Cuda(0);

    // A few layers so checkpointing actually wraps multiple blocks.
    let mut cfg = Config::tiny();
    cfg.vocab = 16;
    cfg.n_layers = 4;
    let batch = 3usize;
    let seq = 6usize;
    let seqs: Vec<Vec<i32>> = (0..batch)
        .map(|b| {
            (0..seq)
                .map(|i| ((b * 7 + i * 3 + 1) % cfg.vocab) as i32)
                .collect()
        })
        .collect();
    let tgts: Vec<Vec<i32>> = (0..batch)
        .map(|b| {
            (0..seq)
                .map(|i| ((b * 5 + i * 2 + 2) % cfg.vocab) as i32)
                .collect()
        })
        .collect();
    let ids = batched_ids_tensor(&seqs, device);
    let tgt = batched_ids_tensor(&tgts, device);

    // --- recompute OFF (reference) ---
    let off = build(cfg, device, dtype, false);
    let off_logits = host(&off.forward_batched(&ids, batch).value());
    let off_loss = off.loss_batched(&ids, &tgt, batch);
    let off_loss_val = host(&off_loss.value())[0];
    off_loss.backward();
    let off_grads: Vec<Vec<f32>> = off
        .params()
        .iter()
        .map(|p| host(&p.grad().expect("off grad")))
        .collect();

    // --- recompute ON ---
    let on = build(cfg, device, dtype, true);
    let on_logits = host(&on.forward_batched(&ids, batch).value());
    let on_loss = on.loss_batched(&ids, &tgt, batch);
    let on_loss_val = host(&on_loss.value())[0];
    on_loss.backward();
    let on_grads: Vec<Vec<f32>> = on
        .params()
        .iter()
        .map(|p| host(&p.grad().expect("on grad")))
        .collect();

    // 1. Forward logits — recompute must not change the forward output.
    let logit_rel = off_logits
        .iter()
        .zip(&on_logits)
        .map(|(a, b)| (a - b).abs() / a.abs().max(1e-4))
        .fold(0.0f32, f32::max);
    // 2. Loss.
    let loss_rel = (off_loss_val - on_loss_val).abs() / off_loss_val.abs().max(1e-4);
    println!(
        "[{dtype:?}] recompute on/off: loss {off_loss_val:.6}/{on_loss_val:.6} (rel {loss_rel:.2e}), \
         logits max rel {logit_rel:.2e}"
    );
    assert!(
        logit_rel < logit_tol,
        "[{dtype:?}] logits diverged: {logit_rel:.2e}"
    );
    assert!(
        loss_rel < logit_tol,
        "[{dtype:?}] loss diverged: {loss_rel:.2e}"
    );

    // 3. Every parameter grad — the load-bearing gate.
    let mut max_grad_rel = 0.0f32;
    for (off_g, on_g) in off_grads.iter().zip(&on_grads) {
        for (a, b) in off_g.iter().zip(on_g) {
            let rel = (a - b).abs() / a.abs().max(1e-3);
            max_grad_rel = max_grad_rel.max(rel);
        }
    }
    println!("[{dtype:?}] recompute on/off: grad max rel err = {max_grad_rel:.3e}");
    assert!(
        max_grad_rel < grad_tol,
        "[{dtype:?}] recompute grads diverged from non-recompute: {max_grad_rel:.3e}"
    );
}

#[test]
fn recompute_matches_non_recompute_fp32() {
    // fp32: recompute runs the identical deterministic kernels → grads match to
    // (near) bit-exact; allow a hair for any nondeterministic GPU reduction.
    run(DType::F32, 1e-5, 1e-4);
}

#[test]
fn recompute_matches_non_recompute_bf16() {
    // bf16 (T12 composition): same bf16 path on recompute. The recompute is still
    // exact w.r.t. the bf16 forward, so on/off match tightly (looser tol only for
    // bf16 rounding, not for any recompute discrepancy).
    run(DType::BF16, 5e-3, 5e-3);
}