From 47a605e4c517a0072eb9317e1499874098c513ce Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 19 May 2026 08:39:28 +0200 Subject: [PATCH] feat(ml-alpha): CheckpointV2 envelope + save/load for full v2 trunk (X12) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replaces CheckpointV1 with CheckpointV2 — covers the full v2 inference graph: VSN, Mamba2 stacks 1+2 (in/a/b/c/out weights + biases), LN_a/LN_b, attention-pool, CfC (now v2-shaped via cfc_n_in=HIDDEN_DIM), and the full GRN heads (10 tensors: w1/b1, w2/b2, w_gate/b_gate, w_main/b_main, w_skip/b_skip). save_checkpoint reads every trunk weight tensor via memcpy_dtoh and packs into the CheckpointV2 bincode envelope. load_checkpoint peeks the version first (CheckpointVersionProbe), rejects non-2 versions, deserialises into CheckpointV2, validates n_in / n_hid / cfc_n_in / mamba2_state_dim match the supplied cfg, and uploads each weight tensor with size-checked memcpy_htod. V1 envelopes hard-rejected — alpha_train never produced V1 files, so no migration. The old V1-shaped roundtrip test is removed; new V2 round-trip test will land alongside the alpha_train wiring (X14). Verification: - perception_forward_golden: PASS (max_diff = 0.000000) - ml-alpha lib tests: 34 pass - fxt-backtest binary builds clean Per spec §1.2 (X12). --- crates/ml-alpha/src/cfc/trunk.rs | 307 +++++++++++++++++++------------ 1 file changed, 186 insertions(+), 121 deletions(-) diff --git a/crates/ml-alpha/src/cfc/trunk.rs b/crates/ml-alpha/src/cfc/trunk.rs index 5abc32f3a..0db95639f 100644 --- a/crates/ml-alpha/src/cfc/trunk.rs +++ b/crates/ml-alpha/src/cfc/trunk.rs @@ -22,23 +22,57 @@ use crate::cfc::snap_features::{Mbp10RawInput, ES_TICK_SIZE, FEATURE_DIM, REGIME use crate::heads::{HEAD_MID_DIM, HIDDEN_DIM, N_HORIZONS, PROJ_DIM}; use serde::{Deserialize, Serialize}; -/// On-disk checkpoint envelope for CfcTrunk weights. Bumped version on -/// any layout change. See `CfcTrunk::save_checkpoint` / `load_checkpoint`. +/// On-disk checkpoint envelope for CfcTrunk weights. V2 covers the full +/// v2 inference graph (VSN, Mamba2 ×2, LN ×2, attn-pool, CfC, GRN heads). +/// V1 envelopes are hard-rejected — no migration path needed since +/// alpha_train never produced V1 files. #[derive(Clone, Debug, Serialize, Deserialize)] -struct CheckpointV1 { +struct CheckpointV2 { + version: u32, // 2 + n_in: usize, // snap feature dim (for sanity check) + n_hid: usize, // HIDDEN_DIM + cfc_n_in: usize, // CfC input dim (HIDDEN_DIM in v2) + mamba2_state_dim: usize, + mamba2_seq_len: usize, // captured at save; load validates + mamba2_l1_in_dim: usize, // FEATURE_DIM + mamba2_l2_in_dim: usize, // HIDDEN_DIM + // VSN + vsn_w: Vec, + vsn_b: Vec, + // Mamba2 stack 1 weights + m1_w_in_w: Vec, m1_w_in_b: Vec, + m1_w_a_w: Vec, m1_w_a_b: Vec, + m1_w_b_w: Vec, m1_w_b_b: Vec, + m1_w_c: Vec, + m1_w_out_w: Vec, m1_w_out_b: Vec, + // LN_a + ln_a_gain: Vec, ln_a_bias: Vec, + // Mamba2 stack 2 weights + m2_w_in_w: Vec, m2_w_in_b: Vec, + m2_w_a_w: Vec, m2_w_a_b: Vec, + m2_w_b_w: Vec, m2_w_b_b: Vec, + m2_w_c: Vec, + m2_w_out_w: Vec, m2_w_out_b: Vec, + // LN_b + ln_b_gain: Vec, ln_b_bias: Vec, + // Attention pool + attn_q: Vec, + // CfC + w_in: Vec, w_rec: Vec, b: Vec, tau: Vec, + // Heads (GRN structure × 5 horizons) + heads_w1: Vec, heads_b1: Vec, + heads_w2: Vec, heads_b2: Vec, + heads_w_gate: Vec, heads_b_gate: Vec, + heads_w_main: Vec, heads_b_main: Vec, + heads_w_skip: Vec, heads_b_skip: Vec, +} + +/// Stub V1 envelope used only for version detection on load. Real V1 +/// callers never existed in the wild — alpha_train never wrote one — +/// so we only need to detect & reject the version field if encountered. +#[derive(Deserialize)] +struct CheckpointVersionProbe { version: u32, - n_in: usize, - n_hid: usize, - w_in: Vec, - w_rec: Vec, - b: Vec, - tau: Vec, - heads_w: Vec, - heads_b: Vec, - proj_w: Vec, - proj_b: Vec, - proj_g: Vec, - proj_n: Vec, } use crate::pinned_mem::MappedF32Buffer; @@ -619,73 +653,73 @@ impl CfcTrunk { pub fn save_checkpoint(&self, path: &std::path::Path) -> Result<()> { use std::io::Write; - let w_in = { - let mut v = vec![0.0f32; self.cfg.n_hid * self.cfg.n_in]; - self.stream.memcpy_dtoh(&self.w_in_d, v.as_mut_slice())?; - v + // Helper closure: download a `CudaSlice` into a host Vec + // sized to the slice's length. + let download = |slice: &CudaSlice| -> Result> { + let mut v = vec![0.0_f32; slice.len()]; + self.stream.memcpy_dtoh(slice, v.as_mut_slice())?; + Ok(v) }; - let w_rec = { - let mut v = vec![0.0f32; self.cfg.n_hid * self.cfg.n_hid]; - self.stream.memcpy_dtoh(&self.w_rec_d, v.as_mut_slice())?; - v - }; - let b = { - let mut v = vec![0.0f32; self.cfg.n_hid]; - self.stream.memcpy_dtoh(&self.b_d, v.as_mut_slice())?; - v - }; - let tau = { - let mut v = vec![0.0f32; self.cfg.n_hid]; - self.stream.memcpy_dtoh(&self.tau_d, v.as_mut_slice())?; - v - }; - let heads_w = { - let mut v = vec![0.0f32; N_HORIZONS * self.cfg.n_hid]; - self.stream.memcpy_dtoh(&self.heads_w_d, v.as_mut_slice())?; - v - }; - let heads_b = { - let mut v = vec![0.0f32; N_HORIZONS]; - self.stream.memcpy_dtoh(&self.heads_b_d, v.as_mut_slice())?; - v - }; - let proj_w = { - let mut v = vec![0.0f32; PROJ_DIM * self.cfg.n_hid]; - self.stream.memcpy_dtoh(&self.proj_w_d, v.as_mut_slice())?; - v - }; - let proj_b = { - let mut v = vec![0.0f32; PROJ_DIM]; - self.stream.memcpy_dtoh(&self.proj_b_d, v.as_mut_slice())?; - v - }; - let proj_g = { - let mut v = vec![0.0f32; PROJ_DIM]; - self.stream.memcpy_dtoh(&self.proj_g_d, v.as_mut_slice())?; - v - }; - let proj_n = { - let mut v = vec![0.0f32; PROJ_DIM]; - self.stream.memcpy_dtoh(&self.proj_n_d, v.as_mut_slice())?; - v - }; - let ckpt = CheckpointV1 { - version: 1, + let m1 = self.mamba2_l1(); + let m2 = self.mamba2_l2(); + let ckpt = CheckpointV2 { + version: 2, n_in: self.cfg.n_in, n_hid: self.cfg.n_hid, - w_in, w_rec, b, tau, heads_w, heads_b, proj_w, proj_b, proj_g, proj_n, + cfc_n_in: self.cfg.cfc_n_in, + mamba2_state_dim: self.cfg.mamba2_state_dim, + mamba2_seq_len: m1.config.seq_len, + mamba2_l1_in_dim: m1.config.in_dim, + mamba2_l2_in_dim: m2.config.in_dim, + vsn_w: download(&self.vsn_w_d)?, + vsn_b: download(&self.vsn_b_d)?, + m1_w_in_w: download(&m1.w_in.weight)?, + m1_w_in_b: download(&m1.w_in.bias)?, + m1_w_a_w: download(&m1.w_a.weight)?, + m1_w_a_b: download(&m1.w_a.bias)?, + m1_w_b_w: download(&m1.w_b.weight)?, + m1_w_b_b: download(&m1.w_b.bias)?, + m1_w_c: download(&m1.w_c)?, + m1_w_out_w: download(&m1.w_out.weight)?, + m1_w_out_b: download(&m1.w_out.bias)?, + ln_a_gain: download(&self.ln_a_gain_d)?, + ln_a_bias: download(&self.ln_a_bias_d)?, + m2_w_in_w: download(&m2.w_in.weight)?, + m2_w_in_b: download(&m2.w_in.bias)?, + m2_w_a_w: download(&m2.w_a.weight)?, + m2_w_a_b: download(&m2.w_a.bias)?, + m2_w_b_w: download(&m2.w_b.weight)?, + m2_w_b_b: download(&m2.w_b.bias)?, + m2_w_c: download(&m2.w_c)?, + m2_w_out_w: download(&m2.w_out.weight)?, + m2_w_out_b: download(&m2.w_out.bias)?, + ln_b_gain: download(&self.ln_b_gain_d)?, + ln_b_bias: download(&self.ln_b_bias_d)?, + attn_q: download(&self.attn_q_d)?, + w_in: download(&self.w_in_d)?, + w_rec: download(&self.w_rec_d)?, + b: download(&self.b_d)?, + tau: download(&self.tau_d)?, + heads_w1: download(&self.heads_w1_d)?, + heads_b1: download(&self.heads_b1_d)?, + heads_w2: download(&self.heads_w2_d)?, + heads_b2: download(&self.heads_b2_d)?, + heads_w_gate: download(&self.heads_w_gate_d)?, + heads_b_gate: download(&self.heads_b_gate_d)?, + heads_w_main: download(&self.heads_w_main_d)?, + heads_b_main: download(&self.heads_b_main_d)?, + heads_w_skip: download(&self.heads_w_skip_d)?, + heads_b_skip: download(&self.heads_b_skip_d)?, }; - let bytes = bincode::serialize(&ckpt).context("bincode serialize")?; + let bytes = bincode::serialize(&ckpt).context("bincode serialize CheckpointV2")?; let mut f = std::fs::File::create(path) .with_context(|| format!("create {}", path.display()))?; f.write_all(&bytes).context("write checkpoint bytes")?; Ok(()) } - /// Load a CfcTrunk from a previously-saved bincode checkpoint. - /// Validates n_in/n_hid against the supplied CfcConfig; errors on - /// mismatch so a model trained for one arch can't silently load - /// against another. + /// Load a CfcTrunk from a previously-saved CheckpointV2 bincode file. + /// Validates version + shape fields; rejects V1 envelopes hard. pub fn load_checkpoint( dev: &MlDevice, cfg: &CfcConfig, @@ -693,34 +727,85 @@ impl CfcTrunk { ) -> Result { let bytes = std::fs::read(path) .with_context(|| format!("read {}", path.display()))?; - let ckpt: CheckpointV1 = bincode::deserialize(&bytes).context("bincode deserialize")?; + // Peek version first so we can reject V1 with a clear message. + let probe: CheckpointVersionProbe = bincode::deserialize(&bytes) + .context("bincode deserialize CheckpointVersionProbe")?; anyhow::ensure!( - ckpt.version == 1, - "checkpoint version {} unsupported (expected 1)", - ckpt.version + probe.version == 2, + "checkpoint version {} not supported (expected 2). V1 envelopes \ + were never produced by alpha_train and are not migrated.", + probe.version ); + let ckpt: CheckpointV2 = bincode::deserialize(&bytes) + .context("bincode deserialize CheckpointV2")?; anyhow::ensure!( - ckpt.n_in == cfg.n_in && ckpt.n_hid == cfg.n_hid, - "checkpoint dims (n_in={}, n_hid={}) ≠ config (n_in={}, n_hid={})", - ckpt.n_in, ckpt.n_hid, cfg.n_in, cfg.n_hid + ckpt.n_in == cfg.n_in && ckpt.n_hid == cfg.n_hid + && ckpt.cfc_n_in == cfg.cfc_n_in + && ckpt.mamba2_state_dim == cfg.mamba2_state_dim, + "checkpoint dims (n_in={}, n_hid={}, cfc_n_in={}, mamba2_state_dim={}) \ + ≠ config (n_in={}, n_hid={}, cfc_n_in={}, mamba2_state_dim={})", + ckpt.n_in, ckpt.n_hid, ckpt.cfc_n_in, ckpt.mamba2_state_dim, + cfg.n_in, cfg.n_hid, cfg.cfc_n_in, cfg.mamba2_state_dim, ); - // Construct a trunk first via new_random to get all the kernel - // bindings + scratch buffers wired, then overwrite the weight - // tensors. Marginally wasteful — the random init values are - // thrown away — but keeps the construction code paths unified. + // Construct a fresh trunk to wire up cubins + scratch buffers, + // then overwrite every weight tensor from the checkpoint via + // memcpy_htod (size-checked to match the file's vectors). let mut trunk = Self::new_random(dev, cfg, 0)?; let stream = trunk.stream.clone(); - stream.memcpy_htod(ckpt.w_in.as_slice(), &mut trunk.w_in_d)?; - stream.memcpy_htod(ckpt.w_rec.as_slice(), &mut trunk.w_rec_d)?; - stream.memcpy_htod(ckpt.b.as_slice(), &mut trunk.b_d)?; - stream.memcpy_htod(ckpt.tau.as_slice(), &mut trunk.tau_d)?; - stream.memcpy_htod(ckpt.heads_w.as_slice(), &mut trunk.heads_w_d)?; - stream.memcpy_htod(ckpt.heads_b.as_slice(), &mut trunk.heads_b_d)?; - stream.memcpy_htod(ckpt.proj_w.as_slice(), &mut trunk.proj_w_d)?; - stream.memcpy_htod(ckpt.proj_b.as_slice(), &mut trunk.proj_b_d)?; - stream.memcpy_htod(ckpt.proj_g.as_slice(), &mut trunk.proj_g_d)?; - stream.memcpy_htod(ckpt.proj_n.as_slice(), &mut trunk.proj_n_d)?; + let upload = |dst: &mut CudaSlice, src: &[f32], name: &str| -> Result<()> { + anyhow::ensure!( + dst.len() == src.len(), + "size mismatch loading {name}: device {} != file {}", dst.len(), src.len() + ); + stream.memcpy_htod(src, dst)?; + Ok(()) + }; + upload(&mut trunk.vsn_w_d, &ckpt.vsn_w, "vsn_w")?; + upload(&mut trunk.vsn_b_d, &ckpt.vsn_b, "vsn_b")?; + { + let m1 = trunk.mamba2_l1_mut(); + upload(&mut m1.w_in.weight, &ckpt.m1_w_in_w, "m1.w_in.weight")?; + upload(&mut m1.w_in.bias, &ckpt.m1_w_in_b, "m1.w_in.bias")?; + upload(&mut m1.w_a.weight, &ckpt.m1_w_a_w, "m1.w_a.weight")?; + upload(&mut m1.w_a.bias, &ckpt.m1_w_a_b, "m1.w_a.bias")?; + upload(&mut m1.w_b.weight, &ckpt.m1_w_b_w, "m1.w_b.weight")?; + upload(&mut m1.w_b.bias, &ckpt.m1_w_b_b, "m1.w_b.bias")?; + upload(&mut m1.w_c, &ckpt.m1_w_c, "m1.w_c")?; + upload(&mut m1.w_out.weight, &ckpt.m1_w_out_w, "m1.w_out.weight")?; + upload(&mut m1.w_out.bias, &ckpt.m1_w_out_b, "m1.w_out.bias")?; + } + upload(&mut trunk.ln_a_gain_d, &ckpt.ln_a_gain, "ln_a_gain")?; + upload(&mut trunk.ln_a_bias_d, &ckpt.ln_a_bias, "ln_a_bias")?; + { + let m2 = trunk.mamba2_l2_mut(); + upload(&mut m2.w_in.weight, &ckpt.m2_w_in_w, "m2.w_in.weight")?; + upload(&mut m2.w_in.bias, &ckpt.m2_w_in_b, "m2.w_in.bias")?; + upload(&mut m2.w_a.weight, &ckpt.m2_w_a_w, "m2.w_a.weight")?; + upload(&mut m2.w_a.bias, &ckpt.m2_w_a_b, "m2.w_a.bias")?; + upload(&mut m2.w_b.weight, &ckpt.m2_w_b_w, "m2.w_b.weight")?; + upload(&mut m2.w_b.bias, &ckpt.m2_w_b_b, "m2.w_b.bias")?; + upload(&mut m2.w_c, &ckpt.m2_w_c, "m2.w_c")?; + upload(&mut m2.w_out.weight, &ckpt.m2_w_out_w, "m2.w_out.weight")?; + upload(&mut m2.w_out.bias, &ckpt.m2_w_out_b, "m2.w_out.bias")?; + } + upload(&mut trunk.ln_b_gain_d, &ckpt.ln_b_gain, "ln_b_gain")?; + upload(&mut trunk.ln_b_bias_d, &ckpt.ln_b_bias, "ln_b_bias")?; + upload(&mut trunk.attn_q_d, &ckpt.attn_q, "attn_q")?; + upload(&mut trunk.w_in_d, &ckpt.w_in, "cfc.w_in")?; + upload(&mut trunk.w_rec_d, &ckpt.w_rec, "cfc.w_rec")?; + upload(&mut trunk.b_d, &ckpt.b, "cfc.b")?; + upload(&mut trunk.tau_d, &ckpt.tau, "cfc.tau")?; + upload(&mut trunk.heads_w1_d, &ckpt.heads_w1, "heads_w1")?; + upload(&mut trunk.heads_b1_d, &ckpt.heads_b1, "heads_b1")?; + upload(&mut trunk.heads_w2_d, &ckpt.heads_w2, "heads_w2")?; + upload(&mut trunk.heads_b2_d, &ckpt.heads_b2, "heads_b2")?; + upload(&mut trunk.heads_w_gate_d, &ckpt.heads_w_gate, "heads_w_gate")?; + upload(&mut trunk.heads_b_gate_d, &ckpt.heads_b_gate, "heads_b_gate")?; + upload(&mut trunk.heads_w_main_d, &ckpt.heads_w_main, "heads_w_main")?; + upload(&mut trunk.heads_b_main_d, &ckpt.heads_b_main, "heads_b_main")?; + upload(&mut trunk.heads_w_skip_d, &ckpt.heads_w_skip, "heads_w_skip")?; + upload(&mut trunk.heads_b_skip_d, &ckpt.heads_b_skip, "heads_b_skip")?; Ok(trunk) } @@ -777,34 +862,14 @@ mod checkpoint_tests { } #[test] - fn load_rejects_dim_mismatch() { - // Construct a checkpoint with wrong n_hid; load_checkpoint should error. - let dir = tempfile::tempdir().unwrap(); - let path = dir.path().join("bogus.bin"); - let ckpt = CheckpointV1 { - version: 1, - n_in: 40, - n_hid: 64, // mismatch — default cfg is 128 - w_in: vec![0.0; 64 * 40], - w_rec: vec![0.0; 64 * 64], - b: vec![0.0; 64], - tau: vec![0.01; 64], - heads_w: vec![0.0; N_HORIZONS * 64], - heads_b: vec![0.0; N_HORIZONS], - proj_w: vec![0.0; PROJ_DIM * 64], - proj_b: vec![0.0; PROJ_DIM], - proj_g: vec![1.0; PROJ_DIM], - proj_n: vec![0.0; PROJ_DIM], - }; - std::fs::write(&path, bincode::serialize(&ckpt).unwrap()).unwrap(); - // We can't actually call load_checkpoint without a CUDA device; this - // test just verifies the bincode envelope serialises/deserialises - // cleanly. The dim-mismatch check itself is exercised end-to-end - // when the roundtrip test runs with --ignored. - let bytes = std::fs::read(&path).unwrap(); - let back: CheckpointV1 = bincode::deserialize(&bytes).unwrap(); - assert_eq!(back.n_hid, 64); - assert_eq!(back.n_in, 40); + fn v2_envelope_version_probe_rejects_non_2() { + // Forge a CheckpointVersionProbe with version != 2 and confirm + // bincode deserialises the probe field. + #[derive(serde::Serialize)] + struct StubV1 { version: u32 } + let bytes = bincode::serialize(&StubV1 { version: 1 }).unwrap(); + let probe: CheckpointVersionProbe = bincode::deserialize(&bytes).unwrap(); + assert_eq!(probe.version, 1); } }