From c24cf7423b57e0e6acad7b91b8753f97d91e5406 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sun, 17 May 2026 21:38:15 +0200 Subject: [PATCH] feat(ml-alpha): wire LayerNorm into PerceptionTrainer (Phase 1.6) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- crates/ml-alpha/src/trainer/perception.rs | 206 +++++++++++++++++++++- 1 file changed, 199 insertions(+), 7 deletions(-) diff --git a/crates/ml-alpha/src/trainer/perception.rs b/crates/ml-alpha/src/trainer/perception.rs index 73abe904f..f317620d1 100644 --- a/crates/ml-alpha/src/trainer/perception.rs +++ b/crates/ml-alpha/src/trainer/perception.rs @@ -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, + /// 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, + /// LayerNorm bias (`[HIDDEN_DIM]`, initialised to 0.0). + ln_bias_d: CudaSlice, + /// Per-row LN stats `[B*K, 2]` (mean, inv_std) saved by fwd, used + /// by bwd. + ln_stats_d: CudaSlice, + /// 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, + /// Per-row gain-grad scratch `[B*K, HIDDEN]`, reduced into + /// `grad_ln_gain_d`. + grad_ln_gain_per_row_d: CudaSlice, + /// Per-row bias-grad scratch `[B*K, HIDDEN]`, reduced into + /// `grad_ln_bias_d`. + grad_ln_bias_per_row_d: CudaSlice, + /// Final LN gain gradient `[HIDDEN]` — Adam consumes this. + grad_ln_gain_d: CudaSlice, + /// Final LN bias gradient `[HIDDEN]` — Adam consumes this. + grad_ln_bias_d: CudaSlice, + /// 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, + /// 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 = vec![1.0; HIDDEN_DIM]; + let ln_bias_init: Vec = 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::(1)?, loss_host_d: unsafe { MappedF32Buffer::new(1) }.map_err(|e| anyhow::anyhow!("loss_host_d: {e}"))?, loss_per_horizon_d: stream.alloc_zeros::(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::(cfg.n_batch * k * 2)?, + ln_out_d: stream.alloc_zeros::(cfg.n_batch * k * HIDDEN_DIM)?, + grad_ln_gain_per_row_d: stream.alloc_zeros::(cfg.n_batch * k * HIDDEN_DIM)?, + grad_ln_bias_per_row_d: stream.alloc_zeros::(cfg.n_batch * k * HIDDEN_DIM)?, + grad_ln_gain_d: stream.alloc_zeros::(HIDDEN_DIM)?, + grad_ln_bias_d: stream.alloc_zeros::(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::(N_HORIZONS)?, lambda_d: stream.alloc_zeros::(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")?; }