Merge pull request 'rtx-autograd: training smoke tests proving the tape learns' (#4) from tape-training-smoke into main

crates/core/rtx-autograd/tests/tape_train_smoke.rs

@@ -0,0 +1,232 @@
+//! End-to-end *training* smoke tests on the `Autodiff<CpuBackend>` tape:
+//! forward → backward → extract grads → SGD update → loss decreases. This is
+//! the capability the gradient-correctness fixes (PR: tape correct+sound on
+//! real backends) were for — it proves a tape-native network actually learns,
+//! the precondition for the SMT predictive-state teacher.
+//!
+//! Params are kept as host `Vec<f32>`; each step rebuilds grad-tracked leaves
+//! (define-by-run), reads each leaf's gradient out of the `GradientStorage`,
+//! and applies a plain SGD step on the host. Loss is sum-of-squared-error
+//! (the constant `1/N` of MSE is irrelevant to whether it decreases).
+
+use rtx_autograd::autodiff::{Autodiff, AutodiffDevice, GradTensor, backward_impl};
+use rtx_backend::{AutodiffBackend, Backend};
+use rtx_backend_cpu::CpuBackend;
+
+type Ad = Autodiff<CpuBackend>;
+
+fn ad_dev() -> AutodiffDevice<CpuBackend> {
+    AutodiffDevice::<CpuBackend>::default()
+}
+
+fn cpu_dev() -> <CpuBackend as Backend>::Device {
+    <CpuBackend as Backend>::Device::default()
+}
+
+/// Tiny deterministic init in [-scale, scale) via an LCG — avoids a uniform
+/// all-zeros start (which would give zero gradients through a linear layer).
+fn init(n: usize, scale: f32, seed: u64) -> Vec<f32> {
+    let mut s = seed.wrapping_add(0x9E37_79B9_7F4A_7C15);
+    (0..n)
+        .map(|_| {
+            s = s
+                .wrapping_mul(6364136223846793005)
+                .wrapping_add(1442695040888963407);
+            let u = ((s >> 33) as f32) / ((1u64 << 31) as f32); // [0,1)
+            (u * 2.0 - 1.0) * scale
+        })
+        .collect()
+}
+
+/// A trainable parameter: owns its data + shape, mints a fresh grad leaf each
+/// step, and applies an SGD update from the gradient it was given.
+struct Param {
+    data: Vec<f32>,
+    shape: [usize; 2],
+}
+
+impl Param {
+    fn new(rows: usize, cols: usize, scale: f32, seed: u64) -> Self {
+        Self {
+            data: init(rows * cols, scale, seed),
+            shape: [rows, cols],
+        }
+    }
+
+    /// A fresh grad-tracked leaf for this step; returns (leaf, leaf_id).
+    fn leaf(&self) -> (<Ad as Backend>::TensorPrimitive<2>, usize) {
+        let t = Ad::require_grad(Ad::from_data(&self.data, self.shape, &ad_dev()));
+        let id = t.id().0;
+        (t, id)
+    }
+
+    fn sgd(&mut self, grad: &[f32], lr: f32) {
+        for (w, g) in self.data.iter_mut().zip(grad.iter()) {
+            *w -= lr * g;
+        }
+    }
+}
+
+/// Read a 2-D leaf gradient out of the storage as a flat Vec.
+fn grad_of(storage: &rtx_autograd::autodiff::GradientStorage<CpuBackend>, id: usize) -> Vec<f32> {
+    match storage
+        .get(rtx_autograd::autodiff::TensorId(id))
+        .unwrap_or_else(|| panic!("no gradient for leaf {id}"))
+    {
+        GradTensor::D2(t) => t.to_vec(),
+        other => panic!("expected D2 gradient, got {other:?}"),
+    }
+}
+
+/// A non-grad constant 2-D tensor on the tape.
+fn constant(data: &[f32], shape: [usize; 2]) -> <Ad as Backend>::TensorPrimitive<2> {
+    Ad::from_data(data, shape, &ad_dev())
+}
+
+fn backward(
+    loss: &<Ad as Backend>::TensorPrimitive<1>,
+) -> rtx_autograd::autodiff::GradientStorage<CpuBackend> {
+    backward_impl(
+        loss,
+        Some(GradTensor::from_d1(CpuBackend::ones([1], &cpu_dev()))),
+    )
+    .expect("backward")
+}
+
+/// Train a 2-layer MLP (`x @ W1 -> gelu -> @ W2`) to fit a fixed synthetic
+/// target, and assert the loss falls substantially.
+#[test]
+fn mlp_trains_and_loss_decreases() {
+    const N: usize = 8; // samples
+    const DIN: usize = 4;
+    const H: usize = 8;
+    const DOUT: usize = 3;
+    const LR: f32 = 0.02;
+    const STEPS: usize = 200;
+
+    let x = init(N * DIN, 1.0, 1);
+    let y = init(N * DOUT, 1.0, 2); // arbitrary fixed target the net must fit
+
+    let mut w1 = Param::new(DIN, H, 0.3, 10);
+    let mut w2 = Param::new(H, DOUT, 0.3, 20);
+
+    let loss_at = |w1: &Param, w2: &Param| -> (f32, Vec<f32>, Vec<f32>) {
+        let (l1, id1) = w1.leaf();
+        let (l2, id2) = w2.leaf();
+        let xt = constant(&x, [N, DIN]);
+        let yt = constant(&y, [N, DOUT]);
+        let hidden = Ad::gelu(Ad::matmul(xt, l1));
+        let pred = Ad::matmul(hidden, l2);
+        let err = Ad::sub(pred, yt);
+        let loss = Ad::sum(Ad::mul(err.clone(), err));
+        let loss_val = Ad::to_data(&loss)[0];
+        let storage = backward(&loss);
+        (loss_val, grad_of(&storage, id1), grad_of(&storage, id2))
+    };
+
+    let (first_loss, _, _) = loss_at(&w1, &w2);
+    let mut last_loss = first_loss;
+    for step in 0..STEPS {
+        let (loss_val, g1, g2) = loss_at(&w1, &w2);
+        assert!(loss_val.is_finite(), "non-finite loss at step {step}");
+        w1.sgd(&g1, LR);
+        w2.sgd(&g2, LR);
+        last_loss = loss_val;
+    }
+
+    eprintln!("MLP loss: {first_loss:.4} -> {last_loss:.4}");
+    assert!(
+        last_loss < first_loss * 0.2,
+        "MLP did not learn: {first_loss:.4} -> {last_loss:.4}"
+    );
+}
+
+/// Train a minimal **attention set-encoder** — the SMT predictive-state teacher
+/// shape — to predict the per-dim mean of its input window. The embedding `e`
+/// fans out to Q, K, V *and* the residual (4 uses), so this also exercises the
+/// fan-out gradient-accumulation fix inside a real attention block. All ops
+/// used (matmul, transpose, softmax, add, mul, sum) are gradient-checked in
+/// `tape_cpu_gradcheck.rs`.
+#[test]
+fn attention_set_encoder_learns_window_mean() {
+    const L: usize = 6; // window length
+    const DIM: usize = 4; // token width
+    const DM: usize = 8; // model width
+    const LR: f32 = 0.01;
+    const STEPS: usize = 300;
+
+    // A fixed window; target = per-dim mean over the L tokens (needs memory).
+    let x = init(L * DIM, 1.0, 7);
+    let mut target = vec![0.0f32; DIM];
+    for d in 0..DIM {
+        let mut acc = 0.0;
+        for l in 0..L {
+            acc += x[l * DIM + d];
+        }
+        target[d] = acc / L as f32;
+    }
+
+    let mut w_embed = Param::new(DIM, DM, 0.3, 100);
+    let mut wq = Param::new(DM, DM, 0.3, 101);
+    let mut wk = Param::new(DM, DM, 0.3, 102);
+    let mut wv = Param::new(DM, DM, 0.3, 103);
+    let mut wdec = Param::new(DM, DIM, 0.3, 104);
+    let ones_row = vec![1.0f32; L];
+
+    let step = |w_embed: &Param,
+                wq: &Param,
+                wk: &Param,
+                wv: &Param,
+                wdec: &Param|
+     -> (f32, Vec<Vec<f32>>) {
+        let (we, we_id) = w_embed.leaf();
+        let (q_w, q_id) = wq.leaf();
+        let (k_w, k_id) = wk.leaf();
+        let (v_w, v_id) = wv.leaf();
+        let (dec_w, dec_id) = wdec.leaf();
+
+        let xt = constant(&x, [L, DIM]);
+        let e = Ad::matmul(xt, we); // [L, DM]
+        // e fans out 4 ways → exercises fan-out accumulation.
+        let q = Ad::matmul(e.clone(), q_w); // [L, DM]
+        let k = Ad::matmul(e.clone(), k_w);
+        let v = Ad::matmul(e.clone(), v_w);
+        let scores = Ad::matmul(q, Ad::transpose(k)); // [L, L]
+        let attn = Ad::softmax(scores, 1);
+        let ctx = Ad::matmul(attn, v); // [L, DM]
+        let h = Ad::add(e, ctx); // residual [L, DM]
+        let pool = Ad::matmul(constant(&ones_row, [1, L]), h); // sum-pool [1, DM]
+        let pred = Ad::matmul(pool, dec_w); // [1, DIM]
+        let err = Ad::sub(pred, constant(&target, [1, DIM]));
+        let loss = Ad::sum(Ad::mul(err.clone(), err));
+        let loss_val = Ad::to_data(&loss)[0];
+        let storage = backward(&loss);
+        let grads = vec![
+            grad_of(&storage, we_id),
+            grad_of(&storage, q_id),
+            grad_of(&storage, k_id),
+            grad_of(&storage, v_id),
+            grad_of(&storage, dec_id),
+        ];
+        (loss_val, grads)
+    };
+
+    let (first_loss, _) = step(&w_embed, &wq, &wk, &wv, &wdec);
+    let mut last_loss = first_loss;
+    for s in 0..STEPS {
+        let (loss_val, g) = step(&w_embed, &wq, &wk, &wv, &wdec);
+        assert!(loss_val.is_finite(), "non-finite loss at step {s}");
+        w_embed.sgd(&g[0], LR);
+        wq.sgd(&g[1], LR);
+        wk.sgd(&g[2], LR);
+        wv.sgd(&g[3], LR);
+        wdec.sgd(&g[4], LR);
+        last_loss = loss_val;
+    }
+
+    eprintln!("attention teacher loss: {first_loss:.5} -> {last_loss:.5}");
+    assert!(
+        last_loss < first_loss * 0.2,
+        "attention set-encoder did not learn: {first_loss:.5} -> {last_loss:.5}"
+    );
+}