gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor(ml-alpha): move Mamba2 stack 2 from PerceptionTrainer to CfcTrunk (X5)

Browse Source 
Same pattern as X3: trainer constructs mamba2_l2 + Mamba2AdamW (against
&mamba2_l2), then moves the block into trunk.mamba2_stack_2. All
forward/backward/AdamW access sites redirect to self.trunk.mamba2_l2_mut().
Gradient buffers + AdamW state remain at trainer level.

Verification: golden bit-exact; ml-alpha lib green.

Per spec §2.2 (X5).
This commit is contained in:
jgrusewski
2026-05-19 01:43:00 +02:00
parent 868021e818
commit 5a534e9972
1 changed files with 11 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
crates/ml-alpha/src/trainer/perception.rs
View File

@@ -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")?;