diff --git a/crates/ml-alpha/src/trainer/perception.rs b/crates/ml-alpha/src/trainer/perception.rs index af28441eb..6d18f1d73 100644 --- a/crates/ml-alpha/src/trainer/perception.rs +++ b/crates/ml-alpha/src/trainer/perception.rs @@ -157,10 +157,8 @@ pub struct PerceptionTrainer { // 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 - /// (128), in_dim=HIDDEN_DIM (consumes the LN_a output stream). - pub mamba2_l2: Mamba2Block, + // X5: Mamba2 stack 2 weights moved to `self.trunk.mamba2_stack_2`. + // Access via `self.trunk.mamba2_l2()` / `mamba2_l2_mut()`. pub mamba2_l2_adamw: Mamba2AdamW, mamba2_l2_fwd_scratch: Mamba2BlockForwardScratch, mamba2_l2_bwd_scratch: Mamba2BackwardScratch, @@ -672,6 +670,8 @@ 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) @@ -1071,7 +1071,6 @@ impl PerceptionTrainer { mamba2_fwd_scratch, mamba2_bwd_scratch, mamba2_grads_buffers, - mamba2_l2, mamba2_l2_adamw, mamba2_l2_fwd_scratch, mamba2_l2_bwd_scratch, @@ -1465,7 +1464,8 @@ impl PerceptionTrainer { } // ── 2b. Mamba2 stack-2 forward — consumes ln_a_out_d. in_dim = HIDDEN_DIM. - self.mamba2_l2 + self.trunk + .mamba2_l2_mut() .forward_train_seq_into(&self.ln_a_out_d, &mut self.mamba2_l2_fwd_scratch) .context("mamba2 (l2) forward_train_seq_into")?; @@ -2048,7 +2048,8 @@ impl PerceptionTrainer { // Writes: // mamba2_l2_grads_buffers.d_x_from_in = grad w.r.t. m2's // input = grad w.r.t. LN_a output (consumed by LN_a bwd below). - self.mamba2_l2 + self.trunk + .mamba2_l2_mut() .backward_from_h_enriched_seq_full_into( &self.ln_a_out_d, &self.mamba2_l2_fwd_scratch, @@ -2305,7 +2306,7 @@ impl PerceptionTrainer { .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) + .step_from_buffers_gpu_clip(self.trunk.mamba2_l2_mut(), &self.mamba2_l2_grads_buffers) .context("mamba2 (l2) AdamW step_from_buffers_gpu_clip")?; // Queue loss_d → loss_host_d (mapped-pinned). Captured inside @@ -2504,7 +2505,8 @@ impl PerceptionTrainer { } // Mamba2 stack-2 fwd → m2.h_enriched_seq. - self.mamba2_l2 + self.trunk + .mamba2_l2_mut() .forward_train_seq_into(&self.ln_a_out_d, &mut self.mamba2_l2_fwd_scratch) .context("eval mamba2 (l2) fwd_into")?;