fix(ml-alpha): construct Mamba2 stacks in CfcTrunk::new_random
Production smoke panicked at trunk.rs:mamba2_l1() during
PerceptionTrainer::from_checkpoint. Root cause: load_checkpoint
called new_random which left mamba2_stack_1/2 = None, then tried
to upload weights into None blocks.
PerceptionTrainer::new was the only caller that populated the
Mamba2 stacks (X3/X5 migrations stopped at the accessor methods
but never moved construction). Inference paths that didn't go
through PerceptionTrainer::new — exactly what fxt-backtest does
via from_checkpoint → load_checkpoint — hit the gap.
Move Mamba2Block::new + state allocation into CfcTrunk::new_random.
PerceptionTrainer::new now sets up its optimizer + scratches against
the trunk-owned blocks via existing accessors (mamba2_l1/mamba2_l2).
Extends the trunk roundtrip test to:
- assert mamba2 weight slices are non-empty (catches empty-buffer
regressions that would let bit-equivalence pass trivially)
- bit-equivalence-check Mamba2 stack 1 + stack 2 weights through
save -> load, reproducing the exact failure path that hit prod.
Verified locally on RTX 3050 Ti — full workspace cargo check passes,
test passes with non-empty weight slices.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -19,6 +19,7 @@ use rand_chacha::ChaCha8Rng;
|
||||
|
||||
use crate::cfc::snap_features::FEATURE_DIM;
|
||||
use crate::heads::{HEAD_MID_DIM, HIDDEN_DIM, N_HORIZONS};
|
||||
use crate::mamba2_block::{Mamba2Block, Mamba2BlockConfig};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// On-disk checkpoint envelope for CfcTrunk weights. Covers the full
|
||||
@@ -230,8 +231,36 @@ impl CfcTrunk {
|
||||
// tensors of the correct shape, just not trained.
|
||||
vsn_w_d: stream.alloc_zeros::<f32>(cfg.n_in * cfg.n_in).context("vsn_w alloc")?,
|
||||
vsn_b_d: stream.alloc_zeros::<f32>(cfg.n_in).context("vsn_b alloc")?,
|
||||
mamba2_stack_1: None,
|
||||
mamba2_stack_2: None,
|
||||
// Mamba2 stack 1: in=FEATURE_DIM, hidden=HIDDEN_DIM. Constructed
|
||||
// here so that load_checkpoint (which goes through new_random)
|
||||
// gets populated stacks ready for weight upload. Block's own
|
||||
// random init is overwritten by either PerceptionTrainer's seeded
|
||||
// re-init during training, or load_checkpoint's file weights.
|
||||
mamba2_stack_1: Some(
|
||||
Mamba2Block::new(
|
||||
Mamba2BlockConfig {
|
||||
in_dim: FEATURE_DIM,
|
||||
hidden_dim: HIDDEN_DIM,
|
||||
state_dim: cfg.mamba2_state_dim,
|
||||
seq_len: cfg.seq_len,
|
||||
},
|
||||
stream.clone(),
|
||||
)
|
||||
.context("mamba2_stack_1 construct")?,
|
||||
),
|
||||
// Mamba2 stack 2: same shape, consumes LN_a output (HIDDEN_DIM).
|
||||
mamba2_stack_2: Some(
|
||||
Mamba2Block::new(
|
||||
Mamba2BlockConfig {
|
||||
in_dim: HIDDEN_DIM,
|
||||
hidden_dim: HIDDEN_DIM,
|
||||
state_dim: cfg.mamba2_state_dim,
|
||||
seq_len: cfg.seq_len,
|
||||
},
|
||||
stream.clone(),
|
||||
)
|
||||
.context("mamba2_stack_2 construct")?,
|
||||
),
|
||||
ln_a_gain_d: stream.alloc_zeros::<f32>(cfg.n_hid).context("ln_a_gain alloc")?,
|
||||
ln_a_bias_d: stream.alloc_zeros::<f32>(cfg.n_hid).context("ln_a_bias alloc")?,
|
||||
ln_b_gain_d: stream.alloc_zeros::<f32>(cfg.n_hid).context("ln_b_gain alloc")?,
|
||||
@@ -447,9 +476,15 @@ mod checkpoint_tests {
|
||||
use super::*;
|
||||
use ml_core::device::MlDevice;
|
||||
|
||||
/// Regression test for the production panic at
|
||||
/// `trunk.rs:mamba2_l1()` during `PerceptionTrainer::from_checkpoint`.
|
||||
/// Reproduces save -> load -> assert all weights including Mamba2 ×2
|
||||
/// round-trip bit-identically. Originally `new_random` left
|
||||
/// `mamba2_stack_1/2 = None`, so `save_checkpoint` itself panicked
|
||||
/// before the file even got written.
|
||||
#[test]
|
||||
#[ignore = "requires CUDA"]
|
||||
fn save_load_roundtrip_preserves_cfc_weights() {
|
||||
fn save_load_roundtrip_preserves_all_weights() {
|
||||
let dev = match MlDevice::cuda(0) {
|
||||
Ok(d) => d,
|
||||
Err(e) => {
|
||||
@@ -461,6 +496,12 @@ mod checkpoint_tests {
|
||||
let trunk_a = CfcTrunk::new_random(&dev, &cfg, 0xDEAD_BEEF).expect("new_random");
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("trunk.bin");
|
||||
// Both stacks must allocate real weight buffers — zero-length slices
|
||||
// would let bit-equivalence pass trivially while shipping a useless
|
||||
// checkpoint.
|
||||
assert!(trunk_a.mamba2_l1().w_in.weight.len() > 0, "m1 w_in empty");
|
||||
assert!(trunk_a.mamba2_l1().w_a.weight.len() > 0, "m1 w_a empty");
|
||||
assert!(trunk_a.mamba2_l2().w_in.weight.len() > 0, "m2 w_in empty");
|
||||
trunk_a.save_checkpoint(&path).expect("save_checkpoint");
|
||||
let trunk_b = CfcTrunk::load_checkpoint(&dev, &cfg, &path).expect("load_checkpoint");
|
||||
let read = |trunk: &CfcTrunk, slice: &CudaSlice<f32>| -> Vec<f32> {
|
||||
@@ -476,7 +517,25 @@ mod checkpoint_tests {
|
||||
assert_eq!(read(&trunk_a, &trunk_a.ln_a_gain_d), read(&trunk_b, &trunk_b.ln_a_gain_d));
|
||||
assert_eq!(read(&trunk_a, &trunk_a.attn_q_d), read(&trunk_b, &trunk_b.attn_q_d));
|
||||
assert_eq!(read(&trunk_a, &trunk_a.heads_w_gate_d), read(&trunk_b, &trunk_b.heads_w_gate_d));
|
||||
}
|
||||
|
||||
// Mamba2 stack 1: the exact bug we shipped to production. Both
|
||||
// saved-side `mamba2_l1()` and load-side weight upload would panic
|
||||
// if mamba2_stack_1 weren't populated by new_random.
|
||||
let m1a = trunk_a.mamba2_l1();
|
||||
let m1b = trunk_b.mamba2_l1();
|
||||
assert_eq!(read(&trunk_a, &m1a.w_in.weight), read(&trunk_b, &m1b.w_in.weight));
|
||||
assert_eq!(read(&trunk_a, &m1a.w_in.bias), read(&trunk_b, &m1b.w_in.bias));
|
||||
assert_eq!(read(&trunk_a, &m1a.w_a.weight), read(&trunk_b, &m1b.w_a.weight));
|
||||
assert_eq!(read(&trunk_a, &m1a.w_b.weight), read(&trunk_b, &m1b.w_b.weight));
|
||||
assert_eq!(read(&trunk_a, &m1a.w_c), read(&trunk_b, &m1b.w_c));
|
||||
assert_eq!(read(&trunk_a, &m1a.w_out.weight), read(&trunk_b, &m1b.w_out.weight));
|
||||
|
||||
// Mamba2 stack 2: same.
|
||||
let m2a = trunk_a.mamba2_l2();
|
||||
let m2b = trunk_b.mamba2_l2();
|
||||
assert_eq!(read(&trunk_a, &m2a.w_in.weight), read(&trunk_b, &m2b.w_in.weight));
|
||||
assert_eq!(read(&trunk_a, &m2a.w_a.weight), read(&trunk_b, &m2b.w_a.weight));
|
||||
assert_eq!(read(&trunk_a, &m2a.w_out.weight), read(&trunk_b, &m2b.w_out.weight));
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -49,7 +49,7 @@ use crate::cfc::snap_features::{Mbp10RawInput, ES_TICK_SIZE, FEATURE_DIM, REGIME
|
||||
use crate::heads::{HEAD_MID_DIM, HIDDEN_DIM, N_HORIZONS};
|
||||
use crate::mamba2_block::{
|
||||
Mamba2AdamW, Mamba2AdamWConfig, Mamba2BackwardGradsBuffers, Mamba2BackwardScratch,
|
||||
Mamba2Block, Mamba2BlockConfig, Mamba2BlockForwardScratch,
|
||||
Mamba2BlockForwardScratch,
|
||||
};
|
||||
use crate::pinned_mem::{MappedF32Buffer, MappedI32Buffer, MappedI64Buffer};
|
||||
use crate::trainer::optim::AdamW;
|
||||
@@ -571,19 +571,10 @@ impl PerceptionTrainer {
|
||||
.load_function("multi_horizon_heads_grn_bwd_batched")
|
||||
.context("heads GRN bwd symbol")?;
|
||||
|
||||
// Mamba2 encoder: in_dim=FEATURE_DIM (32), hidden=HIDDEN_DIM (128).
|
||||
let mamba2 = Mamba2Block::new(
|
||||
Mamba2BlockConfig {
|
||||
in_dim: FEATURE_DIM,
|
||||
hidden_dim: HIDDEN_DIM,
|
||||
state_dim: cfg.mamba2_state_dim,
|
||||
seq_len: cfg.seq_len,
|
||||
},
|
||||
stream.clone(),
|
||||
)
|
||||
.context("Mamba2Block::new")?;
|
||||
// Mamba2 stacks live on the trunk from new_random; we wire optimizer
|
||||
// + training scratches against the trunk-owned blocks.
|
||||
let mamba2_adamw = Mamba2AdamW::new(
|
||||
&mamba2,
|
||||
trunk.mamba2_l1(),
|
||||
Mamba2AdamWConfig {
|
||||
lr: cfg.lr_mamba2,
|
||||
grad_clip_max_norm: Some(1.0),
|
||||
@@ -600,27 +591,10 @@ impl PerceptionTrainer {
|
||||
let mamba2_grads_buffers = Mamba2BackwardGradsBuffers::new(
|
||||
&stream, cfg.n_batch, cfg.seq_len, FEATURE_DIM, HIDDEN_DIM, cfg.mamba2_state_dim,
|
||||
).context("Mamba2BackwardGradsBuffers::new")?;
|
||||
// X3: move Mamba2 stack 1 into the trunk; subsequent access via
|
||||
// self.trunk.mamba2_l1() / mamba2_l1_mut(). Mamba2AdamW was built
|
||||
// above against &mamba2 (read-only borrow that ended before this
|
||||
// move), so the optimizer state is fully constructed.
|
||||
trunk.mamba2_stack_1 = Some(mamba2);
|
||||
|
||||
// ── Phase 2B: SECOND Mamba2 stack ──
|
||||
// Same hidden_dim + state_dim as stack 1. in_dim = HIDDEN_DIM
|
||||
// (consumes LN_a output which has the same shape as m1's output).
|
||||
let mamba2_l2 = Mamba2Block::new(
|
||||
Mamba2BlockConfig {
|
||||
in_dim: HIDDEN_DIM,
|
||||
hidden_dim: HIDDEN_DIM,
|
||||
state_dim: cfg.mamba2_state_dim,
|
||||
seq_len: cfg.seq_len,
|
||||
},
|
||||
stream.clone(),
|
||||
)
|
||||
.context("Mamba2Block::new (l2)")?;
|
||||
// ── Phase 2B: SECOND Mamba2 stack — same shape, on the trunk. ──
|
||||
let mamba2_l2_adamw = Mamba2AdamW::new(
|
||||
&mamba2_l2,
|
||||
trunk.mamba2_l2(),
|
||||
Mamba2AdamWConfig {
|
||||
lr: cfg.lr_mamba2,
|
||||
grad_clip_max_norm: Some(1.0),
|
||||
@@ -637,8 +611,6 @@ impl PerceptionTrainer {
|
||||
let mamba2_l2_grads_buffers = Mamba2BackwardGradsBuffers::new(
|
||||
&stream, cfg.n_batch, cfg.seq_len, HIDDEN_DIM, HIDDEN_DIM, cfg.mamba2_state_dim,
|
||||
).context("Mamba2BackwardGradsBuffers::new (l2)")?;
|
||||
// X5: move Mamba2 stack 2 into the trunk.
|
||||
trunk.mamba2_stack_2 = Some(mamba2_l2);
|
||||
let window_tensor_d = GpuTensor::zeros(&[cfg.n_batch, cfg.seq_len, FEATURE_DIM], &stream)
|
||||
.map_err(|e| anyhow::anyhow!("window_tensor_d alloc: {e}"))?;
|
||||
let h_enriched_seq_t_d = GpuTensor::zeros(&[cfg.seq_len, cfg.n_batch, HIDDEN_DIM], &stream)
|
||||
|
||||
Reference in New Issue
Block a user