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

perf(ml-alpha): device-resident AdamW step counter (capture prep)

Browse Source 
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<i32>`
    - 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<i32>` for the production
      gpu_clip path; advances via `kernel_increment_step` kernel
      inside the captured region.
    - adamw_apply_devscale: `t: i32` → `step_d: &CudaSlice<i32>`.

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 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-17 14:39:38 +02:00
parent b6fb720acd
commit ab94ce2a49
4 changed files with 146 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
crates/ml-alpha/cuda/adamw_step.cu
View File

@@ -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;
}
}

26
crates/ml-alpha/cuda/mamba2_alpha_kernel.cu
View File

@@ -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 →

56
crates/ml-alpha/src/mamba2_block.rs
View File

@@ -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<f32>,
@@ -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<f32>,
/// 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<i32>,
/// 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::<f32>(1)
.map_err(|e| anyhow!("grad_scale alloc: {e}"))?;
let step_count_d = stream.alloc_zeros::<i32>(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<f32>,
grad: &CudaSlice<f32>,
state: &mut AdamState,
t: i32,
step_d: &CudaSlice<i32>,
grad_scale_d: &CudaSlice<f32>,
) -> 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}"))?;

79
crates/ml-alpha/src/trainer/optim.rs
View File

@@ -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<CudaStream>,
_module: Arc<CudaModule>,
func: CudaFunction,
inc_func: CudaFunction,
m: CudaSlice<f32>,
v: CudaSlice<f32>,
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<i32>,
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::<f32>(n_params).context("m alloc")?;
let v = stream.alloc_zeros::<f32>(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::<i32>(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<f32>, grad: &CudaSlice<f32>) -> 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) }
}
}