refactor(ml-alpha): move Mamba2 stack 2 from PerceptionTrainer to CfcTrunk (X5)
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:
@@ -157,10 +157,8 @@ pub struct PerceptionTrainer {
|
|||||||
// Access via `self.trunk.mamba2_l1()` / `mamba2_l1_mut()`. Gradient
|
// Access via `self.trunk.mamba2_l1()` / `mamba2_l1_mut()`. Gradient
|
||||||
// buffers + AdamW state remain at trainer level.
|
// buffers + AdamW state remain at trainer level.
|
||||||
pub mamba2_adamw: Mamba2AdamW,
|
pub mamba2_adamw: Mamba2AdamW,
|
||||||
/// Phase 2B: SECOND Mamba2 stack. Sits between the (first-stack-output
|
// X5: Mamba2 stack 2 weights moved to `self.trunk.mamba2_stack_2`.
|
||||||
/// + LN_a) and the existing LN_b → CfC. Same hidden_dim as stack 1
|
// Access via `self.trunk.mamba2_l2()` / `mamba2_l2_mut()`.
|
||||||
/// (128), in_dim=HIDDEN_DIM (consumes the LN_a output stream).
|
|
||||||
pub mamba2_l2: Mamba2Block,
|
|
||||||
pub mamba2_l2_adamw: Mamba2AdamW,
|
pub mamba2_l2_adamw: Mamba2AdamW,
|
||||||
mamba2_l2_fwd_scratch: Mamba2BlockForwardScratch,
|
mamba2_l2_fwd_scratch: Mamba2BlockForwardScratch,
|
||||||
mamba2_l2_bwd_scratch: Mamba2BackwardScratch,
|
mamba2_l2_bwd_scratch: Mamba2BackwardScratch,
|
||||||
@@ -672,6 +670,8 @@ impl PerceptionTrainer {
|
|||||||
let mamba2_l2_grads_buffers = Mamba2BackwardGradsBuffers::new(
|
let mamba2_l2_grads_buffers = Mamba2BackwardGradsBuffers::new(
|
||||||
&stream, cfg.n_batch, cfg.seq_len, HIDDEN_DIM, HIDDEN_DIM, cfg.mamba2_state_dim,
|
&stream, cfg.n_batch, cfg.seq_len, HIDDEN_DIM, HIDDEN_DIM, cfg.mamba2_state_dim,
|
||||||
).context("Mamba2BackwardGradsBuffers::new (l2)")?;
|
).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)
|
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}"))?;
|
.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)
|
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_fwd_scratch,
|
||||||
mamba2_bwd_scratch,
|
mamba2_bwd_scratch,
|
||||||
mamba2_grads_buffers,
|
mamba2_grads_buffers,
|
||||||
mamba2_l2,
|
|
||||||
mamba2_l2_adamw,
|
mamba2_l2_adamw,
|
||||||
mamba2_l2_fwd_scratch,
|
mamba2_l2_fwd_scratch,
|
||||||
mamba2_l2_bwd_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.
|
// ── 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)
|
.forward_train_seq_into(&self.ln_a_out_d, &mut self.mamba2_l2_fwd_scratch)
|
||||||
.context("mamba2 (l2) forward_train_seq_into")?;
|
.context("mamba2 (l2) forward_train_seq_into")?;
|
||||||
|
|
||||||
@@ -2048,7 +2048,8 @@ impl PerceptionTrainer {
|
|||||||
// Writes:
|
// Writes:
|
||||||
// mamba2_l2_grads_buffers.d_x_from_in = grad w.r.t. m2's
|
// 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).
|
// 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(
|
.backward_from_h_enriched_seq_full_into(
|
||||||
&self.ln_a_out_d,
|
&self.ln_a_out_d,
|
||||||
&self.mamba2_l2_fwd_scratch,
|
&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)
|
.step_from_buffers_gpu_clip(self.trunk.mamba2_l1_mut(), &self.mamba2_grads_buffers)
|
||||||
.context("mamba2 (l1) AdamW step_from_buffers_gpu_clip")?;
|
.context("mamba2 (l1) AdamW step_from_buffers_gpu_clip")?;
|
||||||
self.mamba2_l2_adamw
|
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")?;
|
.context("mamba2 (l2) AdamW step_from_buffers_gpu_clip")?;
|
||||||
|
|
||||||
// Queue loss_d → loss_host_d (mapped-pinned). Captured inside
|
// Queue loss_d → loss_host_d (mapped-pinned). Captured inside
|
||||||
@@ -2504,7 +2505,8 @@ impl PerceptionTrainer {
|
|||||||
}
|
}
|
||||||
|
|
||||||
// Mamba2 stack-2 fwd → m2.h_enriched_seq.
|
// 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)
|
.forward_train_seq_into(&self.ln_a_out_d, &mut self.mamba2_l2_fwd_scratch)
|
||||||
.context("eval mamba2 (l2) fwd_into")?;
|
.context("eval mamba2 (l2) fwd_into")?;
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user