From ab94ce2a49e9857a4dfac7de9a5eda27e175b1c0 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sun, 17 May 2026 14:39:38 +0200 Subject: [PATCH] perf(ml-alpha): device-resident AdamW step counter (capture prep) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Stage 1+2 of #162 (CUDA Graph capture of training step). The AdamW kernels previously took the step counter as a host scalar arg, which gets baked into kernel args at CUDA Graph capture time — replays would freeze the counter and produce wrong bias-correction values. Both AdamW variants now read the step from a device pointer, advanced by a tiny 1-thread `increment_counter` kernel that goes inside the captured region. Each replay correctly increments and observes the new step value. Kernel changes: adamw_step.cu: - adamw_step: int step → const int* step_ptr - adamw_increment_counter: new, +=1 on step_ptr[0] mamba2_alpha_kernel.cu: - mamba2_alpha_adamw_step_devscale: int t → const int* step_ptr - mamba2_alpha_increment_step_counter: new Rust changes: trainer/optim.rs (AdamW): - host `step: i32` → device `step_count_d: CudaSlice` - step(): launch increment kernel BEFORE adamw kernel; both read device counter via pointer arg. - step_count(): test-only accessor, mapped-pinned readback (sync). mamba2_block.rs (Mamba2AdamW): - kept host `step_count: i32` for legacy paths (`step`, `step_from_buffers`) which aren't capture-compatible anyway (host grad-norm dtoh, host scalar grad_scale). - added device `step_count_d: CudaSlice` for the production gpu_clip path; advances via `kernel_increment_step` kernel inside the captured region. - adamw_apply_devscale: `t: i32` → `step_d: &CudaSlice`. Validation: - 4 adamw_invariants tests pass (step_count_increments specifically exercises the device counter). - 10 mamba2_block lib tests pass (training_loop_decreases_loss exercises legacy host-counter path). - Synthetic overfit smoke: initial=0.25 → final=0.0006 (matches pre-refactor trajectory bit-for-bit-equivalent). Stage 3+4 (capture brackets + first-call-capture / subsequent-replay state machine in step_batched) follows in the next commit. Co-Authored-By: Claude Opus 4.7 --- crates/ml-alpha/cuda/adamw_step.cu | 19 ++++- crates/ml-alpha/cuda/mamba2_alpha_kernel.cu | 26 +++++-- crates/ml-alpha/src/mamba2_block.rs | 56 +++++++++++---- crates/ml-alpha/src/trainer/optim.rs | 79 +++++++++++++++++---- 4 files changed, 146 insertions(+), 34 deletions(-) diff --git a/crates/ml-alpha/cuda/adamw_step.cu b/crates/ml-alpha/cuda/adamw_step.cu index 4ddd9d8a6..299831693 100644 --- a/crates/ml-alpha/cuda/adamw_step.cu +++ b/crates/ml-alpha/cuda/adamw_step.cu @@ -8,6 +8,13 @@ // // One thread per parameter. Block tree-reduce not needed; this is pure // element-wise. No atomicAdd. +// +// `step_ptr` is a device pointer to a single i32 holding the current +// 1-indexed training step. This eliminates the host scalar arg that +// blocks CUDA Graph capture (host scalars are baked into kernel args +// at capture time; replays would freeze the step counter). The trainer +// must launch `adamw_increment_counter` once per training step (inside +// the captured region) to advance the counter. extern "C" __global__ void adamw_step( float* __restrict__ theta, @@ -20,7 +27,7 @@ extern "C" __global__ void adamw_step( float beta2, float eps, float wd, - int step + const int* __restrict__ step_ptr ) { int i = blockIdx.x * blockDim.x + threadIdx.x; if (i >= n_params) return; @@ -31,6 +38,7 @@ extern "C" __global__ void adamw_step( m[i] = m_n; v[i] = v_n; + const int step = step_ptr[0]; const float bc1 = 1.0f - powf(beta1, (float) step); const float bc2 = 1.0f - powf(beta2, (float) step); const float m_hat = m_n / fmaxf(bc1, 1e-12f); @@ -38,3 +46,12 @@ extern "C" __global__ void adamw_step( theta[i] -= lr * (m_hat / (sqrtf(v_hat) + eps) + wd * theta[i]); } + +// Increment the device-resident step counter. Single-thread kernel; +// launched once per training step (inside the captured graph, AFTER +// all adamw_step launches that read the current step value). +extern "C" __global__ void adamw_increment_counter(int* __restrict__ step_ptr) { + if (threadIdx.x == 0 && blockIdx.x == 0) { + step_ptr[0] += 1; + } +} diff --git a/crates/ml-alpha/cuda/mamba2_alpha_kernel.cu b/crates/ml-alpha/cuda/mamba2_alpha_kernel.cu index e9ac42c58..d6925f37f 100644 --- a/crates/ml-alpha/cuda/mamba2_alpha_kernel.cu +++ b/crates/ml-alpha/cuda/mamba2_alpha_kernel.cu @@ -559,11 +559,14 @@ extern "C" __global__ void mamba2_alpha_adamw_step( } -/// Variant of `mamba2_alpha_adamw_step` that reads `grad_scale` from a -/// DEVICE pointer instead of a host scalar — companion to the -/// GPU-resident grad-clip pipeline in `grad_norm.cu`. Eliminates the -/// per-step host roundtrip + 9 memcpy_dtoh syncs that the host-scalar -/// path required to compute grad-norm. +/// Variant of `mamba2_alpha_adamw_step` that reads BOTH `grad_scale` and +/// `t` (step counter) from DEVICE pointers — companion to the +/// GPU-resident grad-clip pipeline in `grad_norm.cu` AND the CUDA Graph +/// capture path. Eliminates the host scalars that would otherwise be +/// baked into the captured graph (replays would freeze grad_scale to +/// its first-step value and step counter would stop advancing → wrong +/// bias correction). Use `mamba2_alpha_increment_step_counter` once per +/// training step to advance the counter inside the captured region. extern "C" __global__ void mamba2_alpha_adamw_step_devscale( float* __restrict__ param, const float* __restrict__ grad, @@ -575,7 +578,7 @@ extern "C" __global__ void mamba2_alpha_adamw_step_devscale( float epsilon, float weight_decay, const float* __restrict__ grad_scale_ptr, - int t, + const int* __restrict__ step_ptr, int n ) { int i = blockIdx.x * blockDim.x + threadIdx.x; @@ -594,6 +597,7 @@ extern "C" __global__ void mamba2_alpha_adamw_step_devscale( v[i] = vi; /* Bias correction. */ + int t = step_ptr[0]; float bc1 = 1.0f - powf(beta1, (float)t); float bc2 = 1.0f - powf(beta2, (float)t); float m_hat = mi / bc1; @@ -604,6 +608,16 @@ extern "C" __global__ void mamba2_alpha_adamw_step_devscale( param[i] = p; } +/// Increment the device-resident step counter for Mamba2 AdamW. +/// Single-thread kernel; launch once per training step inside the +/// captured graph, AFTER all `mamba2_alpha_adamw_step_devscale` launches +/// that read the current step value. +extern "C" __global__ void mamba2_alpha_increment_step_counter(int* __restrict__ step_ptr) { + if (threadIdx.x == 0 && blockIdx.x == 0) { + step_ptr[0] += 1; + } +} + /* --------------------------------------------------------------------- * Reduction kernel: sum d_w_c_per_sample[N, sh2, state_d] across N → diff --git a/crates/ml-alpha/src/mamba2_block.rs b/crates/ml-alpha/src/mamba2_block.rs index 3259f8334..0525fb356 100644 --- a/crates/ml-alpha/src/mamba2_block.rs +++ b/crates/ml-alpha/src/mamba2_block.rs @@ -1605,6 +1605,10 @@ pub struct Mamba2AdamW { kernel_norm_sq_phase2: CudaFunction, kernel_clip_scale: CudaFunction, kernel_adamw_devscale: CudaFunction, + /// 1-thread kernel that advances `step_count_d` by 1. Launched + /// inside the captured graph so each replay increments the counter + /// before the AdamW kernels read it. + kernel_increment_step: CudaFunction, /// Per-block partial sums for grad-norm reduction phase 1. /// Sized for the largest parameter tensor's grid: max(n)/256 blocks. block_partials_d: CudaSlice, @@ -1613,6 +1617,12 @@ pub struct Mamba2AdamW { /// Final grad-clip scale (= min(1, max_norm/sqrt(norm_sq))). Read /// by mamba2_alpha_adamw_step_devscale as a device pointer arg. grad_scale_d: CudaSlice, + /// Device-resident 1-indexed step counter for the gpu_clip path. + /// Advanced via `kernel_increment_step` inside the captured graph, + /// so each replay sees the correct bias-correction step. Legacy + /// non-captured paths (`step`, `step_from_buffers`) keep using + /// the host-side `step_count` field. + step_count_d: CudaSlice, /// max(grid_blocks) across all 9 tensors — capacity check for /// block_partials_d. max_grid_blocks: usize, @@ -1677,6 +1687,10 @@ impl Mamba2AdamW { ._module .load_function("mamba2_alpha_adamw_step_devscale") .map_err(|e| anyhow!("mamba2_alpha_adamw_step_devscale symbol: {e}"))?; + let kernel_increment_step = block + ._module + .load_function("mamba2_alpha_increment_step_counter") + .map_err(|e| anyhow!("mamba2_alpha_increment_step_counter symbol: {e}"))?; // Sizing: largest parameter tensor's grid dim. With block=256, // the largest tensor is w_in.weight = hidden_dim * in_dim. @@ -1694,6 +1708,8 @@ impl Mamba2AdamW { .map_err(|e| anyhow!("grad_norm_sq alloc: {e}"))?; let grad_scale_d = stream.alloc_zeros::(1) .map_err(|e| anyhow!("grad_scale alloc: {e}"))?; + let step_count_d = stream.alloc_zeros::(1) + .map_err(|e| anyhow!("step_count_d alloc: {e}"))?; Ok(Self { stream, @@ -1702,9 +1718,11 @@ impl Mamba2AdamW { kernel_norm_sq_phase2, kernel_clip_scale, kernel_adamw_devscale, + kernel_increment_step, block_partials_d, grad_norm_sq_d, grad_scale_d, + step_count_d, max_grid_blocks, step_count: 0, s_w_in, s_b_in, s_w_a, s_b_a, s_w_b, s_b_b, s_w_c, s_w_out, s_b_out, @@ -1799,8 +1817,17 @@ impl Mamba2AdamW { block: &mut Mamba2Block, grads: &Mamba2BackwardGradsBuffers, ) -> Result<()> { - self.step_count += 1; - let t = self.step_count; + // Device step counter is the canonical step number for this + // path. Increment kernel advances it BEFORE the adamw kernels + // read it, so the first call observes step=1. + { + let inc_cfg = LaunchConfig { + grid_dim: (1, 1, 1), block_dim: (1, 1, 1), shared_mem_bytes: 0, + }; + let mut launch = self.stream.launch_builder(&self.kernel_increment_step); + launch.arg(&mut self.step_count_d); + unsafe { launch.launch(inc_cfg).map_err(|e| anyhow!("inc_step launch: {e}"))?; } + } // ── 1. Compute grad_scale_d on device. if let Some(max_norm) = self.config.grad_clip_max_norm { @@ -1873,30 +1900,31 @@ impl Mamba2AdamW { } // ── 2. AdamW per param via the devscale kernel (reads - // grad_scale_d as device pointer). + // grad_scale_d AND step_count_d as device pointers). let stream = &self.stream; let kernel = &self.kernel_adamw_devscale; let cfg = &self.config; let grad_scale = &self.grad_scale_d; + let step_d = &self.step_count_d; adamw_apply_devscale(stream, kernel, cfg, block.w_in.weight.len(), - &mut block.w_in.weight, grads.dw_in.cuda_data(), &mut self.s_w_in, t, grad_scale)?; + &mut block.w_in.weight, grads.dw_in.cuda_data(), &mut self.s_w_in, step_d, grad_scale)?; adamw_apply_devscale(stream, kernel, cfg, block.w_in.bias.len(), - &mut block.w_in.bias, grads.db_in.cuda_data(), &mut self.s_b_in, t, grad_scale)?; + &mut block.w_in.bias, grads.db_in.cuda_data(), &mut self.s_b_in, step_d, grad_scale)?; adamw_apply_devscale(stream, kernel, cfg, block.w_a.weight.len(), - &mut block.w_a.weight, grads.dw_a.cuda_data(), &mut self.s_w_a, t, grad_scale)?; + &mut block.w_a.weight, grads.dw_a.cuda_data(), &mut self.s_w_a, step_d, grad_scale)?; adamw_apply_devscale(stream, kernel, cfg, block.w_a.bias.len(), - &mut block.w_a.bias, grads.db_a.cuda_data(), &mut self.s_b_a, t, grad_scale)?; + &mut block.w_a.bias, grads.db_a.cuda_data(), &mut self.s_b_a, step_d, grad_scale)?; adamw_apply_devscale(stream, kernel, cfg, block.w_b.weight.len(), - &mut block.w_b.weight, grads.dw_b.cuda_data(), &mut self.s_w_b, t, grad_scale)?; + &mut block.w_b.weight, grads.dw_b.cuda_data(), &mut self.s_w_b, step_d, grad_scale)?; adamw_apply_devscale(stream, kernel, cfg, block.w_b.bias.len(), - &mut block.w_b.bias, grads.db_b.cuda_data(), &mut self.s_b_b, t, grad_scale)?; + &mut block.w_b.bias, grads.db_b.cuda_data(), &mut self.s_b_b, step_d, grad_scale)?; adamw_apply_devscale(stream, kernel, cfg, block.w_c.len(), - &mut block.w_c, grads.dw_c.cuda_data(), &mut self.s_w_c, t, grad_scale)?; + &mut block.w_c, grads.dw_c.cuda_data(), &mut self.s_w_c, step_d, grad_scale)?; adamw_apply_devscale(stream, kernel, cfg, block.w_out.weight.len(), - &mut block.w_out.weight, grads.dw_out.cuda_data(), &mut self.s_w_out, t, grad_scale)?; + &mut block.w_out.weight, grads.dw_out.cuda_data(), &mut self.s_w_out, step_d, grad_scale)?; adamw_apply_devscale(stream, kernel, cfg, block.w_out.bias.len(), - &mut block.w_out.bias, grads.db_out.cuda_data(), &mut self.s_b_out, t, grad_scale)?; + &mut block.w_out.bias, grads.db_out.cuda_data(), &mut self.s_b_out, step_d, grad_scale)?; Ok(()) } @@ -1979,7 +2007,7 @@ fn adamw_apply_devscale( param: &mut CudaSlice, grad: &CudaSlice, state: &mut AdamState, - t: i32, + step_d: &CudaSlice, grad_scale_d: &CudaSlice, ) -> Result<()> { let block_threads: u32 = 256; @@ -2003,7 +2031,7 @@ fn adamw_apply_devscale( .arg(&cfg.epsilon) .arg(&cfg.weight_decay) .arg(grad_scale_d) // DEVICE pointer instead of host scalar - .arg(&t) + .arg(step_d) // DEVICE pointer instead of host scalar .arg(&n_i32) .launch(launch_cfg) .map_err(|e| anyhow!("mamba2_alpha_adamw_step_devscale launch (n={n}): {e}"))?; diff --git a/crates/ml-alpha/src/trainer/optim.rs b/crates/ml-alpha/src/trainer/optim.rs index e6d35a900..7a522d32e 100644 --- a/crates/ml-alpha/src/trainer/optim.rs +++ b/crates/ml-alpha/src/trainer/optim.rs @@ -1,10 +1,20 @@ //! AdamW optimizer state + kernel binding. +//! +//! Uses a device-resident step counter (`step_count_d`) instead of a +//! host scalar. Required for CUDA Graph capture: host scalars passed +//! via `.arg()` are baked into kernel args at capture time, so a +//! replayed graph would freeze the step counter and produce wrong +//! bias-correction values. The counter is incremented via a tiny +//! single-thread kernel launched after the parameter update — both +//! launches go through the same stream and end up inside any captured +//! region the caller brackets. use std::sync::Arc; use anyhow::{Context, Result}; -use cudarc::driver::{CudaFunction, CudaModule, CudaSlice, CudaStream, LaunchConfig, PushKernelArg}; +use cudarc::driver::{CudaFunction, CudaModule, CudaSlice, CudaStream, DevicePtr, LaunchConfig, PushKernelArg}; use ml_core::device::MlDevice; +use crate::pinned_mem::MappedI32Buffer; const CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/adamw_step.cubin")); @@ -12,9 +22,13 @@ pub struct AdamW { stream: Arc, _module: Arc, func: CudaFunction, + inc_func: CudaFunction, m: CudaSlice, v: CudaSlice, - step: i32, + /// Device-resident step counter (i32). Initialized to 0; the + /// increment kernel advances it after each parameter update, so the + /// first call observes step=1 (matching standard Adam convention). + step_count_d: CudaSlice, pub lr: f32, pub beta1: f32, pub beta2: f32, @@ -28,15 +42,25 @@ impl AdamW { let ctx = dev.cuda_context().context("adamw ctx")?; let module = ctx.load_cubin(CUBIN.to_vec()).context("load adamw cubin")?; let func = module.load_function("adamw_step").context("adamw fn")?; + let inc_func = module + .load_function("adamw_increment_counter") + .context("adamw_increment_counter fn")?; let m = stream.alloc_zeros::(n_params).context("m alloc")?; let v = stream.alloc_zeros::(n_params).context("v alloc")?; + // Step counter starts at 0; the increment kernel inside step() + // advances it to 1 BEFORE the first read by the adamw_step + // kernel — see step() for the ordering. + let step_count_d = stream + .alloc_zeros::(1) + .context("step_count_d alloc")?; Ok(Self { stream, _module: module, func, + inc_func, m, v, - step: 0, + step_count_d, lr, beta1: 0.9, beta2: 0.999, @@ -45,18 +69,31 @@ impl AdamW { }) } - pub fn step_count(&self) -> i32 { - self.step - } - + /// Apply one AdamW update. Launches `adamw_increment_counter` first + /// (advances device step counter from N-1 to N), then `adamw_step` + /// which reads the counter and uses N for bias correction. Both + /// launches are stream-ordered on the trainer's stream — they end + /// up inside any CUDA Graph capture that brackets the call. pub fn step(&mut self, theta: &mut CudaSlice, grad: &CudaSlice) -> Result<()> { let n = theta.len(); assert_eq!(grad.len(), n, "grad/theta size mismatch"); assert_eq!(self.m.len(), n, "m size mismatch"); assert_eq!(self.v.len(), n, "v size mismatch"); - self.step += 1; + + // 1. Advance device step counter (was N-1, becomes N for this update). + let inc_cfg = LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (1, 1, 1), + shared_mem_bytes: 0, + }; + { + let mut inc = self.stream.launch_builder(&self.inc_func); + inc.arg(&mut self.step_count_d); + unsafe { inc.launch(inc_cfg).context("adamw increment launch")?; } + } + + // 2. Apply update — kernel reads step from device pointer. let n_params_i = n as i32; - let step_i = self.step; let cfg = LaunchConfig { grid_dim: ((n as u32).div_ceil(256), 1, 1), block_dim: (256, 1, 1), @@ -67,12 +104,28 @@ impl AdamW { .arg(theta).arg(grad).arg(&mut self.m).arg(&mut self.v) .arg(&n_params_i) .arg(&self.lr).arg(&self.beta1).arg(&self.beta2).arg(&self.eps).arg(&self.wd) - .arg(&step_i); + .arg(&self.step_count_d); unsafe { launch.launch(cfg).context("adamw launch")?; } // No per-call sync — caller batches multiple AdamW launches on the - // same stream and synchronises once at step boundary. Removing 7 - // per-step syncs (one per param group) eliminates ~35-70 μs of - // pipeline flushes per training step. + // same stream and synchronises once at step boundary. Ok(()) } + + /// Read the current device-resident step counter. Forces a stream + /// sync and a mapped-pinned readback — NOT suitable for the hot + /// training path. Intended for tests and one-off diagnostics; the + /// counter is normally observed only indirectly through training + /// loss trajectory. + pub fn step_count(&self) -> i32 { + let staging = unsafe { MappedI32Buffer::new(1) } + .expect("MappedI32Buffer alloc"); + unsafe { + let (src_ptr, _g) = self.step_count_d.device_ptr(&self.stream); + cudarc::driver::result::memcpy_dtod_async( + staging.dev_ptr, src_ptr, 4, self.stream.cu_stream(), + ).expect("step_count dtod"); + } + self.stream.synchronize().expect("step_count sync"); + unsafe { std::ptr::read_volatile(staging.host_ptr) } + } }