feat(ml-alpha): CheckpointV2 envelope + save/load for full v2 trunk (X12)

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).
This commit is contained in:
jgrusewski
2026-05-19 08:39:28 +02:00
parent 716e0d0781
commit 47a605e4c5

View File

@@ -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<f32>,
vsn_b: Vec<f32>,
// Mamba2 stack 1 weights
m1_w_in_w: Vec<f32>, m1_w_in_b: Vec<f32>,
m1_w_a_w: Vec<f32>, m1_w_a_b: Vec<f32>,
m1_w_b_w: Vec<f32>, m1_w_b_b: Vec<f32>,
m1_w_c: Vec<f32>,
m1_w_out_w: Vec<f32>, m1_w_out_b: Vec<f32>,
// LN_a
ln_a_gain: Vec<f32>, ln_a_bias: Vec<f32>,
// Mamba2 stack 2 weights
m2_w_in_w: Vec<f32>, m2_w_in_b: Vec<f32>,
m2_w_a_w: Vec<f32>, m2_w_a_b: Vec<f32>,
m2_w_b_w: Vec<f32>, m2_w_b_b: Vec<f32>,
m2_w_c: Vec<f32>,
m2_w_out_w: Vec<f32>, m2_w_out_b: Vec<f32>,
// LN_b
ln_b_gain: Vec<f32>, ln_b_bias: Vec<f32>,
// Attention pool
attn_q: Vec<f32>,
// CfC
w_in: Vec<f32>, w_rec: Vec<f32>, b: Vec<f32>, tau: Vec<f32>,
// Heads (GRN structure × 5 horizons)
heads_w1: Vec<f32>, heads_b1: Vec<f32>,
heads_w2: Vec<f32>, heads_b2: Vec<f32>,
heads_w_gate: Vec<f32>, heads_b_gate: Vec<f32>,
heads_w_main: Vec<f32>, heads_b_main: Vec<f32>,
heads_w_skip: Vec<f32>, heads_b_skip: Vec<f32>,
}
/// 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<f32>,
w_rec: Vec<f32>,
b: Vec<f32>,
tau: Vec<f32>,
heads_w: Vec<f32>,
heads_b: Vec<f32>,
proj_w: Vec<f32>,
proj_b: Vec<f32>,
proj_g: Vec<f32>,
proj_n: Vec<f32>,
}
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<f32>` into a host Vec
// sized to the slice's length.
let download = |slice: &CudaSlice<f32>| -> Result<Vec<f32>> {
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<Self> {
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<f32>, 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);
}
}