diff --git a/crates/ml-alpha/build.rs b/crates/ml-alpha/build.rs index 028b40a4a..980314519 100644 --- a/crates/ml-alpha/build.rs +++ b/crates/ml-alpha/build.rs @@ -24,6 +24,7 @@ const KERNELS: &[&str] = &[ "horizon_token_attention_pool", // v2-C: horizon-token K-prepend single-Q attention "inverted_attention_pool", // v2-E: iTransformer-style cross-variate attention "regime_moe_gate", // v2-D: top-1 MoE gate + expert dispatch + aux loss + "anchor_l2", // v2-B: L2 anchor regularization toward init "reduce_axis0", // Phase B: cross-batch param-grad reducer ]; diff --git a/crates/ml-alpha/cuda/anchor_l2.cu b/crates/ml-alpha/cuda/anchor_l2.cu new file mode 100644 index 000000000..b6fd680a5 --- /dev/null +++ b/crates/ml-alpha/cuda/anchor_l2.cu @@ -0,0 +1,67 @@ +// anchor_l2.cu — L2 anchor regularization toward initialization (v2 axis B). +// +// Computes anchor_loss = λ · Σ_i (p[i] − p_init[i])² for a single +// flat parameter buffer p of length n. Adds the corresponding L2 +// gradient 2λ(p − p_init) into the parameter's grad buffer (`+=`). +// +// Caller invokes one launch per parameter group anchored. The +// trainer's Wiener-α controller computes λ each step on the host +// (cheap — scalar update; copies λ to a device-side single-float +// buffer before the launch). +// +// PERFORMANCE: +// - Warp-shuffle reduce for the loss sum. One block per parameter +// group; grid_dim = (1, 1, 1). For n ≤ ~100K (typical group size) +// a single 256-thread block is sufficient; threads stride over n. +// - Grad write is per-element, coalesced. + +#define ANCHOR_BLOCK 256 +#define ANCHOR_NWARPS (ANCHOR_BLOCK / 32) + +__device__ __forceinline__ float warp_reduce_sum_an(float v) { + #pragma unroll + for (int s = 16; s > 0; s >>= 1) { + v += __shfl_xor_sync(0xffffffff, v, s); + } + return v; +} + +extern "C" __global__ void anchor_l2_fwd_bwd( + const float* __restrict__ p, // [n] current parameter values + const float* __restrict__ p_init, // [n] snapshot at trainer construction + const float* __restrict__ lambda_scalar, // [1] current λ from Wiener-α controller + int n, + float* __restrict__ loss_out, // [1] λ · Σ (p − p_init)² + float* __restrict__ grad_p // [n] += 2λ (p − p_init) +) { + const int tid = threadIdx.x; + const int lane = tid & 31; + const int warp = tid >> 5; + + const float lam = lambda_scalar[0]; + const float two_lam = 2.0f * lam; + + // Strided per-thread accumulate of (p - p_init)^2 + emit grad. + float local_sum = 0.0f; + for (int i = tid; i < n; i += ANCHOR_BLOCK) { + const float diff = p[i] - p_init[i]; + local_sum += diff * diff; + // Coalesced grad write. + grad_p[i] += two_lam * diff; + } + + // Warp-shuffle reduce. + float warp_sum = warp_reduce_sum_an(local_sum); + + __shared__ float s_warp[ANCHOR_NWARPS]; + if (lane == 0) s_warp[warp] = warp_sum; + __syncthreads(); + + // Cross-warp reduce: all 32 lanes of warp 0 participate; inactive + // lanes contribute 0 via ternary (NOT a divergent shuffle). + float w = (tid < ANCHOR_NWARPS) ? s_warp[tid] : 0.0f; + if (tid < 32) { + w = warp_reduce_sum_an(w); + if (tid == 0) loss_out[0] = lam * w; + } +} diff --git a/crates/ml-alpha/src/anchor_l2.rs b/crates/ml-alpha/src/anchor_l2.rs new file mode 100644 index 000000000..fa6b85142 --- /dev/null +++ b/crates/ml-alpha/src/anchor_l2.rs @@ -0,0 +1,59 @@ +//! L2 anchor regularization host binding (v2 axis B, V8 commit). +//! +//! Anchors a trainable parameter buffer toward its initialization +//! values via `loss = λ · Σ (p − p_init)²`. The Wiener-α controller +//! on the host side adjusts λ each step based on the loss-improvement +//! rate (signal-driven, floor-bounded — see `trainer/anchor_controller.rs`). +//! +//! See `cuda/anchor_l2.cu` for math + the v2 spec §3.5. + +use anyhow::{Context, Result}; +use cudarc::driver::{ + CudaContext, CudaFunction, CudaModule, CudaSlice, CudaStream, LaunchConfig, PushKernelArg, +}; +use std::sync::Arc; + +const CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/anchor_l2.cubin")); + +pub struct AnchorL2 { + _module: Arc, + func: CudaFunction, + stream: Arc, +} + +impl AnchorL2 { + pub fn new(ctx: &Arc, stream: Arc) -> Result { + let module = ctx.load_cubin(CUBIN.to_vec()).context("load anchor_l2 cubin")?; + let func = module.load_function("anchor_l2_fwd_bwd").context("load anchor_l2_fwd_bwd")?; + Ok(Self { _module: module, func, stream }) + } + + /// Compute anchor_loss = λ · ‖p − p_init‖² and accumulate + /// `grad_p += 2λ · (p − p_init)`. `lambda_d` is a device-side + /// [1]-buffer (host writes scalar before launch). + pub fn apply( + &self, + p: &CudaSlice, + p_init: &CudaSlice, + lambda_d: &CudaSlice, + n: i32, + loss_out: &mut CudaSlice, + grad_p: &mut CudaSlice, + ) -> Result<()> { + let cfg = LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }; + let mut launch = self.stream.launch_builder(&self.func); + unsafe { + launch + .arg(p).arg(p_init).arg(lambda_d) + .arg(&n) + .arg(loss_out).arg(grad_p) + .launch(cfg) + .context("anchor_l2_fwd_bwd")?; + } + Ok(()) + } +} diff --git a/crates/ml-alpha/src/lib.rs b/crates/ml-alpha/src/lib.rs index ef1599926..92ad12c96 100644 --- a/crates/ml-alpha/src/lib.rs +++ b/crates/ml-alpha/src/lib.rs @@ -28,6 +28,7 @@ pub mod cfc; pub mod data; pub mod eval; +pub mod anchor_l2; pub mod heads; pub mod horizon_token_attention_pool; pub mod inverted_attention_pool; diff --git a/crates/ml-alpha/src/trainer/anchor_controller.rs b/crates/ml-alpha/src/trainer/anchor_controller.rs new file mode 100644 index 000000000..38f0ec2a5 --- /dev/null +++ b/crates/ml-alpha/src/trainer/anchor_controller.rs @@ -0,0 +1,142 @@ +//! Wiener-α anchor-coefficient controller (v2 axis B controller). +//! +//! Drives `λ_anchor` for the L2 anchor regularization. Signal-driven +//! per `pearl_controller_anchors_isv_driven` + the +//! Wiener-α-with-floor pattern (`pearl_wiener_alpha_floor_for_nonstationary`, +//! `pearl_blend_formulas_must_have_permanent_floor`). +//! +//! Math: +//! loss_change_t = val_loss_t − val_loss_{t-1} # raw signal +//! ema_change ← ema_change + α · (loss_change_t − ema_change) +//! target_λ = clamp(|ema_change| · scale, λ_floor, λ_max) +//! λ_t = max(λ_floor, blend(λ_{t-1}, target_λ, α)) +//! +//! where α is the standard Wiener-α (MSE-optimal under stationarity). +//! Floor protects against deadlock when ema_change → 0 +//! (anchor never fully relaxes — minimum guard against runaway drift). +//! +//! Floor `λ_floor` is itself signal-driven from initial parameter +//! magnitudes: `λ_floor = ‖p_init‖ / (100 · sqrt(numel))`. Captured at +//! controller construction via the first-observation bootstrap pattern +//! (`pearl_first_observation_bootstrap`). + +const WIENER_ALPHA_FLOOR: f32 = 0.4; // pearl_wiener_alpha_floor_for_nonstationary +const LAMBDA_MAX: f32 = 1.0; // safety cap on λ +const LAMBDA_SCALE: f32 = 10.0; // ema_change → λ multiplier + +pub struct AnchorController { + /// Floor for λ_anchor. Anchor never fully relaxes below this. + pub lambda_floor: f32, + pub lambda_max: f32, + + /// Wiener-α smoothing state. + diff_var: f32, + sample_var: f32, + ema_change: f32, + prev_loss: f32, + bootstrapped: bool, + + /// Smoothed λ. + pub lambda: f32, +} + +impl AnchorController { + /// `init_param_magnitude` = sum of L2-norms of all anchored param + /// groups; `init_numel` = total trainable element count anchored. + /// Used to derive the signal-driven floor. + pub fn new(init_param_magnitude: f32, init_numel: usize) -> Self { + let lambda_floor = + init_param_magnitude / (100.0 * (init_numel.max(1) as f32).sqrt()); + Self { + lambda_floor, + lambda_max: LAMBDA_MAX, + diff_var: 0.0, + sample_var: 0.0, + ema_change: 0.0, + prev_loss: 0.0, + bootstrapped: false, + lambda: lambda_floor.max(1e-6), // start at floor + } + } + + /// Step the controller with a new validation-loss reading. + /// Returns the updated λ. Capture-safe (host-side; no GPU touch). + pub fn step(&mut self, val_loss: f32) -> f32 { + if !self.bootstrapped { + // First observation: replace state directly per + // pearl_first_observation_bootstrap. + self.prev_loss = val_loss; + self.bootstrapped = true; + return self.lambda; + } + + let change = val_loss - self.prev_loss; + self.prev_loss = val_loss; + + // Wiener-α: α = diff_var / (diff_var + sample_var + ε). + // Maintain running estimates of diff_var (changes) and + // sample_var (noise around the mean change). + let delta = change - self.ema_change; + self.diff_var = 0.95 * self.diff_var + 0.05 * change * change; + self.sample_var = 0.95 * self.sample_var + 0.05 * delta * delta; + let alpha_raw = self.diff_var / (self.diff_var + self.sample_var + 1e-9); + let alpha = alpha_raw.max(WIENER_ALPHA_FLOOR); + + self.ema_change = self.ema_change + alpha * delta; + + // Target λ proportional to |ema_change| × scale. + let target_lambda = (self.ema_change.abs() * LAMBDA_SCALE) + .clamp(self.lambda_floor, self.lambda_max); + + // Blend with floor — pearl_blend_formulas_must_have_permanent_floor. + let blended = self.lambda + alpha * (target_lambda - self.lambda); + self.lambda = blended.max(self.lambda_floor); + self.lambda + } + + pub fn current_lambda(&self) -> f32 { + self.lambda + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn floor_init_from_signal() { + // ‖p_init‖ = 5.0, numel = 100 → floor = 5 / (100 · 10) = 0.005 + let c = AnchorController::new(5.0, 100); + assert!((c.lambda_floor - 0.005).abs() < 1e-6); + assert!(c.lambda >= c.lambda_floor); + } + + #[test] + fn first_observation_bootstrap_returns_floor() { + let mut c = AnchorController::new(1.0, 100); + let l0 = c.step(0.5); + assert!((l0 - c.lambda_floor).abs() < 1e-6); + } + + #[test] + fn lambda_never_falls_below_floor() { + let mut c = AnchorController::new(1.0, 100); + // Many steps with no change → ema_change → 0 → target → floor. + for _ in 0..1000 { + let l = c.step(0.5); + assert!(l >= c.lambda_floor - 1e-9, "λ={l} below floor={}", c.lambda_floor); + } + } + + #[test] + fn lambda_capped_at_max() { + let mut c = AnchorController::new(1.0, 100); + c.step(0.0); + // Huge loss spikes alternating sign — ema_change should grow, + // but λ stays ≤ LAMBDA_MAX. + for i in 0..100 { + let l = c.step(if i % 2 == 0 { 100.0 } else { -100.0 }); + assert!(l <= LAMBDA_MAX + 1e-6, "λ={l} above max"); + } + } +} diff --git a/crates/ml-alpha/src/trainer/mod.rs b/crates/ml-alpha/src/trainer/mod.rs index 1d751dac8..fcd230dcf 100644 --- a/crates/ml-alpha/src/trainer/mod.rs +++ b/crates/ml-alpha/src/trainer/mod.rs @@ -1,6 +1,7 @@ //! Trainer module — PerceptionTrainer wraps the full stacked //! Mamba2 -> CfC -> heads pipeline plus 6 AdamW optimizer groups. +pub mod anchor_controller; pub mod loss; pub mod loss_sigma; pub mod optim; diff --git a/crates/ml-alpha/tests/anchor_l2_numgrad.rs b/crates/ml-alpha/tests/anchor_l2_numgrad.rs new file mode 100644 index 000000000..1118cebf1 --- /dev/null +++ b/crates/ml-alpha/tests/anchor_l2_numgrad.rs @@ -0,0 +1,138 @@ +//! Numgrad parity test for anchor_l2 (v2 axis B). + +#![cfg(feature = "cuda")] + +use anyhow::{Context, Result}; +use cudarc::driver::{CudaSlice, CudaStream, DevicePtr, DevicePtrMut}; +use ml_alpha::anchor_l2::AnchorL2; +use ml_alpha::pinned_mem::MappedF32Buffer; +use ml_core::device::MlDevice; +use std::sync::Arc; + +const N: usize = 67; // intentionally non-power-of-two to stress stride loops + +fn rng(seed: u64) -> impl FnMut() -> f32 { + let mut s = seed.wrapping_mul(0x9E37_79B9_7F4A_7C15); + move || -> f32 { + s = s.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407); + let x = ((s >> 16) & 0xFFFF) as f32 / 65536.0; + x * 2.0 - 1.0 + } +} + +fn upload(stream: &Arc, host: &[f32]) -> Result> { + let n = host.len(); + let staging = unsafe { MappedF32Buffer::new(n) } + .map_err(|e| anyhow::anyhow!("upload: {e}"))?; + staging.write_from_slice(host); + let mut dst = stream.alloc_zeros::(n)?; + let nbytes = n * std::mem::size_of::(); + unsafe { + let (dst_ptr, _g) = dst.device_ptr_mut(stream); + cudarc::driver::result::memcpy_dtod_async( + dst_ptr, staging.dev_ptr, nbytes, stream.cu_stream(), + )?; + } + Ok(dst) +} + +fn download(stream: &Arc, src: &CudaSlice) -> Result> { + let n = src.len(); + let staging = unsafe { MappedF32Buffer::new(n) } + .map_err(|e| anyhow::anyhow!("download: {e}"))?; + let nbytes = n * std::mem::size_of::(); + unsafe { + let (src_ptr, _g) = src.device_ptr(stream); + cudarc::driver::result::memcpy_dtod_async( + staging.dev_ptr, src_ptr, nbytes, stream.cu_stream(), + )?; + } + stream.synchronize()?; + Ok(staging.read_all()) +} + +fn forward_loss( + anchor: &AnchorL2, + stream: &Arc, + p: &[f32], + p_init: &[f32], + lambda: f32, +) -> Result { + let p_d = upload(stream, p)?; + let p_init_d = upload(stream, p_init)?; + let lambda_d = upload(stream, &[lambda])?; + let mut loss_d = stream.alloc_zeros::(1)?; + let mut grad_d = stream.alloc_zeros::(p.len())?; + anchor.apply(&p_d, &p_init_d, &lambda_d, p.len() as i32, &mut loss_d, &mut grad_d)?; + stream.synchronize()?; + Ok(download(stream, &loss_d)?[0]) +} + +#[test] +#[ignore = "requires CUDA"] +fn forward_and_grad_match_central_difference() -> Result<()> { + let dev = MlDevice::cuda(0)?; + let stream = dev.cuda_stream().context("stream")?.clone(); + let ctx_dev = dev.cuda_context().context("ctx")?; + let anchor = AnchorL2::new(ctx_dev, stream.clone())?; + + let mut r = rng(20260518); + let p: Vec = (0..N).map(|_| r() * 0.5).collect(); + let p_init: Vec = (0..N).map(|_| r() * 0.5).collect(); + let lambda = 0.137_f32; + + // 1. Forward parity: kernel loss == closed-form. + let p_d = upload(&stream, &p)?; + let p_init_d = upload(&stream, &p_init)?; + let lambda_d = upload(&stream, &[lambda])?; + let mut loss_d = stream.alloc_zeros::(1)?; + let mut grad_d = stream.alloc_zeros::(N)?; + anchor.apply(&p_d, &p_init_d, &lambda_d, N as i32, &mut loss_d, &mut grad_d)?; + stream.synchronize()?; + let kernel_loss = download(&stream, &loss_d)?[0]; + let kernel_grad = download(&stream, &grad_d)?; + + let mut expected_loss = 0.0_f32; + for i in 0..N { + let d = p[i] - p_init[i]; + expected_loss += d * d; + } + expected_loss *= lambda; + let abs = (kernel_loss - expected_loss).abs(); + let rel = abs / expected_loss.abs().max(1e-3); + assert!( + abs < 1e-3 || rel < 5e-3, + "loss closed-form mismatch: kernel={kernel_loss:.5} expected={expected_loss:.5}" + ); + + // 2. Closed-form grad parity (analytic = 2λ (p − p_init)). + for i in 0..N { + let expected_grad = 2.0 * lambda * (p[i] - p_init[i]); + let abs = (kernel_grad[i] - expected_grad).abs(); + assert!(abs < 1e-4, + "grad[{i}] closed-form mismatch: kernel={} expected={} (diff={abs:.3e})", + kernel_grad[i], expected_grad); + } + + // 3. Central-difference numgrad — 4 random positions. + let eps = 1e-3_f32; + let mut r2 = rng(42); + for _ in 0..4 { + let idx = ((r2() + 1.0) * 0.5 * N as f32) as usize % N; + let mut p_p = p.clone(); + let mut p_m = p.clone(); + p_p[idx] += eps; + p_m[idx] -= eps; + let l_p = forward_loss(&anchor, &stream, &p_p, &p_init, lambda)?; + let l_m = forward_loss(&anchor, &stream, &p_m, &p_init, lambda)?; + let numgrad = (l_p - l_m) / (2.0 * eps); + let analytic = kernel_grad[idx]; + let abs = (analytic - numgrad).abs(); + let rel = abs / numgrad.abs().max(1e-3); + assert!( + abs < 5e-3 || rel < 5e-2, + "p[{idx}] numgrad mismatch: analytic={analytic:.4} numgrad={numgrad:.4} abs={abs:.3e}" + ); + } + Ok(()) +}