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

feat(ml-alpha): wire LayerNorm into PerceptionTrainer (Phase 1.6)

Browse Source 
LN sits between Mamba2 encoder output and the K-loop CfC consumer:
  Mamba2.h_enriched_seq [B,K,H] → LN fwd → ln_out_d [B,K,H]
  → transpose → h_enriched_seq_t_d [K,B,H] → CfC + heads

Forward path: LN consumes raw Mamba2 output, writes ln_out_d + per-row
stats (mean, inv_std). Transpose now reads from ln_out_d.

Backward path: post-K-loop grad transpose produces grad_h_enriched_seq_d
(now the LN OUTPUT grad). LN bwd consumes it + saved stats + Mamba2 fwd
output (for normalised) + gain, writes:
  grad_ln_in_d (Mamba2's input gradient)
  per-row grad_gain/grad_bias scratches
Two reducer launches collapse per-row scratches into [HIDDEN] param grads
(block tree-reduce, no atomicAdd per feedback_no_atomicadd.md).

AdamW state for ln_gain + ln_bias added (wd=0, lr=lr_cfc). Mamba2 bwd
now consumes grad_ln_in_d, NOT grad_h_enriched_seq_d.

Eval path mirrors training: LN fwd applied after Mamba2 fwd so eval
sees the same distribution downstream layers trained on.

Synthetic overfit smoke passes: BCE 0.37 → 0.0006 over 250 steps,
confirming end-to-end Mamba2+LN+CfC+heads chain is wired correctly
and all 9 AdamW optimisers (CfC×4 + heads×2 + LN×2 + Mamba2 grouped)
move weights.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-17 21:38:15 +02:00
parent 59e236b4e8
commit c24cf7423b
1 changed files with 199 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -59,6 +59,7 @@ const STEP_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/cfc_step.cub
const HEADS_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/multi_horizon_heads.cubin"));
const BCE_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/bce_loss_multi_horizon.cubin"));
const HORIZON_LAMBDA_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/horizon_lambda.cubin"));
const LAYER_NORM_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/layer_norm.cubin"));
#[derive(Clone, Debug)]
pub struct PerceptionTrainerConfig {
@@ -241,6 +242,41 @@ pub struct PerceptionTrainer {
/// tracks to compute dynamic per-horizon weights. Refreshed by
/// the BCE kernel on every training step. Shape: [N_HORIZONS].
loss_per_horizon_d: CudaSlice<f32>,
/// LayerNorm gain (`[HIDDEN_DIM]`, initialised to 1.0). Applied
/// to Mamba2's `h_enriched_seq` output before the K-loop CfC
/// consumes it. Phase 1: stabilises the trunk-to-CfC distribution.
ln_gain_d: CudaSlice<f32>,
/// LayerNorm bias (`[HIDDEN_DIM]`, initialised to 0.0).
ln_bias_d: CudaSlice<f32>,
/// Per-row LN stats `[B*K, 2]` (mean, inv_std) saved by fwd, used
/// by bwd.
ln_stats_d: CudaSlice<f32>,
/// LN forward output `[B*K, HIDDEN]` — fed to the transpose that
/// produces `h_enriched_seq_t_d`. Replaces direct reads of
/// Mamba2's `h_enriched_seq` in the K-loop path.
ln_out_d: CudaSlice<f32>,
/// Per-row gain-grad scratch `[B*K, HIDDEN]`, reduced into
/// `grad_ln_gain_d`.
grad_ln_gain_per_row_d: CudaSlice<f32>,
/// Per-row bias-grad scratch `[B*K, HIDDEN]`, reduced into
/// `grad_ln_bias_d`.
grad_ln_bias_per_row_d: CudaSlice<f32>,
/// Final LN gain gradient `[HIDDEN]` — Adam consumes this.
grad_ln_gain_d: CudaSlice<f32>,
/// Final LN bias gradient `[HIDDEN]` — Adam consumes this.
grad_ln_bias_d: CudaSlice<f32>,
/// LN backward grad-x: the gradient w.r.t. Mamba2's `h_enriched_seq`
/// output. Replaces direct use of `grad_h_enriched_seq_d` when
/// feeding Mamba2 backward.
grad_ln_in_d: GpuTensor,
/// Cached LN kernel handles + module lifetime owner.
ln_fwd_fn: CudaFunction,
ln_bwd_fn: CudaFunction,
ln_reduce_fn: CudaFunction,
_ln_module: Arc<CudaModule>,
/// AdamW state for the two LN params (gain + bias).
opt_ln_gain: AdamW,
opt_ln_bias: AdamW,
/// Per-horizon EMA of `loss_per_horizon_d`. Zero-initialised
/// (sentinel); the horizon_ema_and_lambda kernel detects the
/// sentinel on step 1 and replaces directly per
@@ -327,6 +363,18 @@ impl PerceptionTrainer {
let horizon_lambda_fn = horizon_lambda_module
.load_function("horizon_ema_and_lambda")
.context("horizon_ema_and_lambda symbol")?;
let ln_module = ctx
.load_cubin(LAYER_NORM_CUBIN.to_vec())
.context("layer_norm cubin")?;
let ln_fwd_fn = ln_module
.load_function("layer_norm_fwd")
.context("layer_norm_fwd symbol")?;
let ln_bwd_fn = ln_module
.load_function("layer_norm_bwd")
.context("layer_norm_bwd symbol")?;
let ln_reduce_fn = ln_module
.load_function("layer_norm_reduce_param_grads")
.context("layer_norm_reduce_param_grads symbol")?;
let snap_batched_fn = snap_module.load_function("snap_feature_assemble_batched")?;
let bce_fn = bce_module.load_function("bce_multi_horizon_forward_backward")?;
let step_batched_fn = step_module.load_function("cfc_step_batched")?;
@@ -411,6 +459,17 @@ impl PerceptionTrainer {
let mut opt_heads_b = AdamW::new(dev, N_HORIZONS, cfg.lr_cfc)?;
opt_heads_b.wd = 0.0;
// LayerNorm parameters: gain initialised to 1.0, bias to 0.0.
// Adam wd=0 — LN params are scale/shift, never penalised.
let ln_gain_init: Vec<f32> = vec![1.0; HIDDEN_DIM];
let ln_bias_init: Vec<f32> = vec![0.0; HIDDEN_DIM];
let ln_gain_d = upload(&stream, &ln_gain_init)?;
let ln_bias_d = upload(&stream, &ln_bias_init)?;
let mut opt_ln_gain = AdamW::new(dev, HIDDEN_DIM, cfg.lr_cfc)?;
opt_ln_gain.wd = 0.0;
let mut opt_ln_bias = AdamW::new(dev, HIDDEN_DIM, cfg.lr_cfc)?;
opt_ln_bias.wd = 0.0;
let k = cfg.seq_len;
Ok(Self {
cfg: cfg.clone(),
@@ -422,6 +481,27 @@ impl PerceptionTrainer {
loss_d: stream.alloc_zeros::<f32>(1)?,
loss_host_d: unsafe { MappedF32Buffer::new(1) }.map_err(|e| anyhow::anyhow!("loss_host_d: {e}"))?,
loss_per_horizon_d: stream.alloc_zeros::<f32>(N_HORIZONS)?,
// LayerNorm state. `ln_gain_d` / `ln_bias_d` constructed
// above via `upload` (gain=1, bias=0). LN_HIDDEN matches
// HIDDEN_DIM (128) by construction — checked at kernel
// launch time. Per-row scratch is [B*K, HIDDEN]; the
// single param-grad reducer collapses it to [HIDDEN].
ln_gain_d,
ln_bias_d,
ln_stats_d: stream.alloc_zeros::<f32>(cfg.n_batch * k * 2)?,
ln_out_d: stream.alloc_zeros::<f32>(cfg.n_batch * k * HIDDEN_DIM)?,
grad_ln_gain_per_row_d: stream.alloc_zeros::<f32>(cfg.n_batch * k * HIDDEN_DIM)?,
grad_ln_bias_per_row_d: stream.alloc_zeros::<f32>(cfg.n_batch * k * HIDDEN_DIM)?,
grad_ln_gain_d: stream.alloc_zeros::<f32>(HIDDEN_DIM)?,
grad_ln_bias_d: stream.alloc_zeros::<f32>(HIDDEN_DIM)?,
grad_ln_in_d: GpuTensor::zeros(&[cfg.n_batch, k, HIDDEN_DIM], &stream)
.map_err(|e| anyhow::anyhow!("grad_ln_in_d alloc: {e}"))?,
ln_fwd_fn,
ln_bwd_fn,
ln_reduce_fn,
_ln_module: ln_module,
opt_ln_gain,
opt_ln_bias,
loss_ema_d: stream.alloc_zeros::<f32>(N_HORIZONS)?,
lambda_d: stream.alloc_zeros::<f32>(N_HORIZONS)?,
horizon_lambda_fn,
@@ -812,7 +892,31 @@ impl PerceptionTrainer {
.forward_train_seq_into(&self.window_tensor_d, &mut self.mamba2_fwd_scratch)
.context("mamba2 forward_train_seq_into")?;
// ── 2b. Transpose [B, K, H] → [K, B, H] into pre-allocated buffer.
// ── 2a. LayerNorm forward over h_enriched_seq [B*K, H].
// Stabilises the trunk-to-CfC distribution. One block per
// row, block-wide tree-reduce for mean/var. Stats saved
// for backward in self.ln_stats_d.
{
let n_rows_ln: i32 = (b_sz * k_seq) as i32;
let cfg_ln = LaunchConfig {
grid_dim: (n_rows_ln as u32, 1, 1),
block_dim: (128, 1, 1),
shared_mem_bytes: 0,
};
let mut launch = self.stream.launch_builder(&self.ln_fwd_fn);
launch
.arg(self.mamba2_fwd_scratch.h_enriched_seq.cuda_data())
.arg(&self.ln_gain_d)
.arg(&self.ln_bias_d)
.arg(&n_rows_ln)
.arg(&mut self.ln_out_d)
.arg(&mut self.ln_stats_d);
unsafe { launch.launch(cfg_ln).context("layer_norm_fwd")?; }
}
// ── 2b. Transpose LN-normalised h_enriched [B, K, H] → [K, B, H]
// into pre-allocated buffer (source = ln_out_d, NOT raw
// Mamba2 output).
{
let block_n3: u32 = 32;
let grid_z = (HIDDEN_DIM as u32).div_ceil(block_n3);
@@ -826,7 +930,7 @@ impl PerceptionTrainer {
let n3 = HIDDEN_DIM as i32;
let mut launch = self.stream.launch_builder(&self.transpose_3d_fn);
launch
.arg(self.mamba2_fwd_scratch.h_enriched_seq.cuda_data())
.arg(&self.ln_out_d)
.arg(self.h_enriched_seq_t_d.data_mut())
.arg(&n1).arg(&n2).arg(&n3);
unsafe { launch.launch(cfg_tx).context("transpose h_enriched fwd")?; }
@@ -1069,7 +1173,8 @@ impl PerceptionTrainer {
// to flush before reading the mapped-pinned loss shadow.
// ── 7b. Transpose grad_h_enriched_seq_t_d [K, B, H] → [B, K, H]
// (pre-allocated grad_h_enriched_seq_d).
// (pre-allocated grad_h_enriched_seq_d). This is the
// gradient w.r.t. the LN OUTPUT (= grad_y for LN bwd).
{
let block_n3: u32 = 32;
let grid_z = (HIDDEN_DIM as u32).div_ceil(block_n3);
@@ -1089,25 +1194,91 @@ impl PerceptionTrainer {
unsafe { launch.launch(cfg_tx).context("transpose grad bwd")?; }
}
// ── 7c. LayerNorm backward. Consumes:
// x = Mamba2's h_enriched_seq [B, K, H]
// gain = self.ln_gain_d [H]
// stats = self.ln_stats_d (from fwd) [B*K, 2]
// grad_y = self.grad_h_enriched_seq_d [B, K, H]
// Produces:
// grad_x = self.grad_ln_in_d [B, K, H]
// grad_gain/row = self.grad_ln_gain_per_row_d [B*K, H]
// grad_bias/row = self.grad_ln_bias_per_row_d [B*K, H]
// Then two reducer launches collapse the per-row
// scratches into the final [H] param grad buffers
// (no atomicAdd per feedback_no_atomicadd.md).
{
let n_rows_ln: i32 = (b_sz * k_seq) as i32;
let cfg_ln = LaunchConfig {
grid_dim: (n_rows_ln as u32, 1, 1),
block_dim: (128, 1, 1),
shared_mem_bytes: 0,
};
let mut launch = self.stream.launch_builder(&self.ln_bwd_fn);
launch
.arg(self.mamba2_fwd_scratch.h_enriched_seq.cuda_data())
.arg(&self.ln_gain_d)
.arg(&self.ln_stats_d)
.arg(self.grad_h_enriched_seq_d.cuda_data())
.arg(&n_rows_ln)
.arg(self.grad_ln_in_d.data_mut())
.arg(&mut self.grad_ln_gain_per_row_d)
.arg(&mut self.grad_ln_bias_per_row_d);
unsafe { launch.launch(cfg_ln).context("layer_norm_bwd")?; }
}
// 7c.i — reduce per-row grad_gain → [H].
{
let cfg_red = LaunchConfig {
grid_dim: (HIDDEN_DIM as u32, 1, 1),
block_dim: (128, 1, 1),
shared_mem_bytes: 0,
};
let n_rows_ln: i32 = (b_sz * k_seq) as i32;
let mut launch = self.stream.launch_builder(&self.ln_reduce_fn);
launch
.arg(&self.grad_ln_gain_per_row_d)
.arg(&n_rows_ln)
.arg(&mut self.grad_ln_gain_d);
unsafe { launch.launch(cfg_red).context("layer_norm_reduce gain")?; }
}
// 7c.ii — reduce per-row grad_bias → [H].
{
let cfg_red = LaunchConfig {
grid_dim: (HIDDEN_DIM as u32, 1, 1),
block_dim: (128, 1, 1),
shared_mem_bytes: 0,
};
let n_rows_ln: i32 = (b_sz * k_seq) as i32;
let mut launch = self.stream.launch_builder(&self.ln_reduce_fn);
launch
.arg(&self.grad_ln_bias_per_row_d)
.arg(&n_rows_ln)
.arg(&mut self.grad_ln_bias_d);
unsafe { launch.launch(cfg_red).context("layer_norm_reduce bias")?; }
}
// ── 8. Mamba2 backward — fully pre-allocated path. Writes all
// grads into self.mamba2_grads_buffers; no allocation.
// Now consumes grad_ln_in_d (LN bwd output), NOT the raw
// grad_h_enriched_seq_d — LN sits between Mamba2 and CfC.
self.mamba2
.backward_from_h_enriched_seq_full_into(
&self.window_tensor_d,
&self.mamba2_fwd_scratch,
&self.grad_h_enriched_seq_d,
&self.grad_ln_in_d,
&mut self.mamba2_bwd_scratch,
&mut self.mamba2_grads_buffers,
)
.context("mamba2 backward_from_h_enriched_seq_full_into")?;
// ── 9. Apply AdamW updates on all 7 param groups (added tau).
// ── 9. Apply AdamW updates on all 9 param groups (CfC×4 + heads×2 + LN×2 + Mamba2 grouped).
self.opt_w_in.step(&mut self.w_in_d, &self.grad_w_in_d)?;
self.opt_w_rec.step(&mut self.w_rec_d, &self.grad_w_rec_d)?;
self.opt_b.step(&mut self.b_d, &self.grad_b_d)?;
self.opt_tau.step(&mut self.tau_d, &self.grad_tau_d)?;
self.opt_heads_w.step(&mut self.heads_w_d, &self.grad_heads_w_d)?;
self.opt_heads_b.step(&mut self.heads_b_d, &self.grad_heads_b_d)?;
self.opt_ln_gain.step(&mut self.ln_gain_d, &self.grad_ln_gain_d)?;
self.opt_ln_bias.step(&mut self.ln_bias_d, &self.grad_ln_bias_d)?;
// GPU-resident grad-clip + AdamW (zero host roundtrips).
// Replaces step_from_buffers which did 9× memcpy_dtoh per step.
self.mamba2_adamw
@@ -1270,7 +1441,28 @@ impl PerceptionTrainer {
.forward_train_seq_into(&self.window_tensor_d, &mut self.mamba2_fwd_scratch)
.context("eval mamba2 fwd_into")?;
// Transpose [B, K, H] → [K, B, H] into pre-allocated buffer.
// LN forward over h_enriched_seq [B*K, H] — same as training path
// so eval reads the same distribution downstream layers were
// trained on.
{
let n_rows_ln: i32 = (b_sz * k_seq) as i32;
let cfg_ln = LaunchConfig {
grid_dim: (n_rows_ln as u32, 1, 1),
block_dim: (128, 1, 1),
shared_mem_bytes: 0,
};
let mut launch = self.stream.launch_builder(&self.ln_fwd_fn);
launch
.arg(self.mamba2_fwd_scratch.h_enriched_seq.cuda_data())
.arg(&self.ln_gain_d)
.arg(&self.ln_bias_d)
.arg(&n_rows_ln)
.arg(&mut self.ln_out_d)
.arg(&mut self.ln_stats_d);
unsafe { launch.launch(cfg_ln).context("eval layer_norm_fwd")?; }
}
// Transpose LN-normalised h_enriched [B, K, H] → [K, B, H].
{
let block_n3: u32 = 32;
let grid_z = (HIDDEN_DIM as u32).div_ceil(block_n3);
@@ -1284,7 +1476,7 @@ impl PerceptionTrainer {
let n3 = HIDDEN_DIM as i32;
let mut launch = self.stream.launch_builder(&self.transpose_3d_fn);
launch
.arg(self.mamba2_fwd_scratch.h_enriched_seq.cuda_data())
.arg(&self.ln_out_d)
.arg(self.h_enriched_seq_t_d.data_mut())
.arg(&n1).arg(&n2).arg(&n3);
unsafe { launch.launch(cfg_tx).context("eval transpose h_enriched")?; }