diff --git a/crates/ml-alpha/src/cfc/trunk.rs b/crates/ml-alpha/src/cfc/trunk.rs index 4423a14bb..51b968176 100644 --- a/crates/ml-alpha/src/cfc/trunk.rs +++ b/crates/ml-alpha/src/cfc/trunk.rs @@ -526,6 +526,25 @@ impl CfcTrunk { /// Reads each device tensor back via memcpy_dtoh and packs into a /// CheckpointV1 envelope. Currently writes ~22-25k f32s = ~90KB /// per CfcTrunk. + /// Accessor for Mamba2 stack 1. Panics if X3 migration hasn't run. + pub fn mamba2_l1_mut(&mut self) -> &mut crate::mamba2_block::Mamba2Block { + self.mamba2_stack_1.as_mut() + .expect("Mamba2 stack 1 not populated (X3 migration must have run)") + } + pub fn mamba2_l1(&self) -> &crate::mamba2_block::Mamba2Block { + self.mamba2_stack_1.as_ref() + .expect("Mamba2 stack 1 not populated (X3 migration must have run)") + } + /// Accessor for Mamba2 stack 2. Panics if X5 migration hasn't run. + pub fn mamba2_l2_mut(&mut self) -> &mut crate::mamba2_block::Mamba2Block { + self.mamba2_stack_2.as_mut() + .expect("Mamba2 stack 2 not populated (X5 migration must have run)") + } + pub fn mamba2_l2(&self) -> &crate::mamba2_block::Mamba2Block { + self.mamba2_stack_2.as_ref() + .expect("Mamba2 stack 2 not populated (X5 migration must have run)") + } + pub fn save_checkpoint(&self, path: &std::path::Path) -> Result<()> { use std::io::Write; let w_in = { diff --git a/crates/ml-alpha/src/trainer/perception.rs b/crates/ml-alpha/src/trainer/perception.rs index 096ce72be..57b9ee683 100644 --- a/crates/ml-alpha/src/trainer/perception.rs +++ b/crates/ml-alpha/src/trainer/perception.rs @@ -153,8 +153,9 @@ pub struct PerceptionTrainer { heads_grn_bwd_fn: CudaFunction, transpose_3d_fn: CudaFunction, - // Mamba2 encoder block + its optimizer - pub mamba2: Mamba2Block, + // X3: Mamba2 stack 1 weights moved to `self.trunk.mamba2_stack_1`. + // Access via `self.trunk.mamba2_l1()` / `mamba2_l1_mut()`. Gradient + // buffers + AdamW state remain at trainer level. pub mamba2_adamw: Mamba2AdamW, /// Phase 2B: SECOND Mamba2 stack. Sits between the (first-stack-output /// + LN_a) and the existing LN_b → CfC. Same hidden_dim as stack 1 @@ -635,6 +636,11 @@ 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 @@ -1057,7 +1063,6 @@ impl PerceptionTrainer { heads_grn_fwd_fn, heads_grn_bwd_fn, transpose_3d_fn, - mamba2, mamba2_adamw, mamba2_fwd_scratch, mamba2_bwd_scratch, @@ -1432,7 +1437,8 @@ impl PerceptionTrainer { // ── 2. Mamba2 stack-1 forward — writes into mamba2_fwd_scratch. // Consumes vsn_out_d (gated features). in_dim = FEATURE_DIM (40). - self.mamba2 + self.trunk + .mamba2_l1_mut() .forward_train_seq_into(&self.vsn_out_d, &mut self.mamba2_fwd_scratch) .context("mamba2 (l1) forward_train_seq_into")?; @@ -2115,7 +2121,8 @@ impl PerceptionTrainer { // Writes: // mamba2_grads_buffers.d_x_from_in = grad on VSN output // (consumed by VSN bwd below). - self.mamba2 + self.trunk + .mamba2_l1_mut() .backward_from_h_enriched_seq_full_into( &self.vsn_out_d, &self.mamba2_fwd_scratch, @@ -2293,7 +2300,7 @@ impl PerceptionTrainer { // GPU-resident grad-clip + AdamW (zero host roundtrips). // Replaces step_from_buffers which did 9× memcpy_dtoh per step. self.mamba2_adamw - .step_from_buffers_gpu_clip(&mut self.mamba2, &self.mamba2_grads_buffers) + .step_from_buffers_gpu_clip(self.trunk.mamba2_l1_mut(), &self.mamba2_grads_buffers) .context("mamba2 (l1) AdamW step_from_buffers_gpu_clip")?; self.mamba2_l2_adamw .step_from_buffers_gpu_clip(&mut self.mamba2_l2, &self.mamba2_l2_grads_buffers) @@ -2470,7 +2477,8 @@ impl PerceptionTrainer { } // Mamba2 stack-1 fwd → m1.h_enriched_seq. - self.mamba2 + self.trunk + .mamba2_l1_mut() .forward_train_seq_into(&self.vsn_out_d, &mut self.mamba2_fwd_scratch) .context("eval mamba2 (l1) fwd_into")?;