diff --git a/crates/ml-alpha/src/cfc/trunk.rs b/crates/ml-alpha/src/cfc/trunk.rs index 0d5ba002b..9bc3d2940 100644 --- a/crates/ml-alpha/src/cfc/trunk.rs +++ b/crates/ml-alpha/src/cfc/trunk.rs @@ -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::(cfg.n_in * cfg.n_in).context("vsn_w alloc")?, vsn_b_d: stream.alloc_zeros::(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::(cfg.n_hid).context("ln_a_gain alloc")?, ln_a_bias_d: stream.alloc_zeros::(cfg.n_hid).context("ln_a_bias alloc")?, ln_b_gain_d: stream.alloc_zeros::(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| -> Vec { @@ -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)); + } } diff --git a/crates/ml-alpha/src/trainer/perception.rs b/crates/ml-alpha/src/trainer/perception.rs index 2a598e4ad..339255bea 100644 --- a/crates/ml-alpha/src/trainer/perception.rs +++ b/crates/ml-alpha/src/trainer/perception.rs @@ -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)