b6fb720acd
GPU-resident grad-norm + clip-scale; mapped-pinned loss readback.
Replaces 9× memcpy_dtoh per Mamba2 AdamW step (grad-norm host roundtrip)
+ 1× per-step download() (loss). Saves ~10 stream-sync barriers/step.
New kernels (cuda/grad_norm.cu):
- grad_norm_sq_phase1: per-block tree-reduce of x[i]^2 (no atomicAdd)
- grad_norm_sq_phase2: cross-tensor accumulator (sequential stream-ordered)
- grad_clip_scale: writes min(1, max_norm/sqrt(norm_sq)) to device ptr
- mamba2_alpha_adamw_step_devscale: reads grad_scale from device pointer
instead of host scalar, allowing AdamW kernels to launch async without
waiting for a CPU-side norm computation.
Trainer changes (perception.rs):
- loss_d kept device-side; mapped-pinned MappedF32Buffer shadow.
- Single stream.synchronize() at end of step (was 2: post-bwd + download).
- DtoD copy loss_d → loss_host_d queued, then sync flushes both kernels
+ copy in one barrier. Loss read via host_ptr (no dtoh).
Mamba2 AdamW (mamba2_block.rs):
- step_from_buffers_gpu_clip(): all grad-norm tensors processed via
phase1+phase2 chain, scale computed on-device, AdamW launches with
devscale variant. Zero host roundtrips.
- Pre-allocated block_partials_d, grad_norm_sq_d, grad_scale_d.
Optimizer (optim.rs): removed redundant stream.synchronize() per AdamW step.
Each per-tensor AdamW kernel is stream-ordered; sync only needed before
host reads, which the trainer handles centrally.
Synthetic overfit smoke: initial=0.30 → final=0.0006 (matches pre-refactor
trajectory). Full ml-alpha test suite passes (45 tests across lib +
integration).
Honors:
- feedback_no_htod_htoh_only_mapped_pinned.md (MappedF32Buffer only)
- feedback_no_atomicadd.md (block tree-reduce only)
- feedback_no_legacy_aliases.md (step_from_buffers replaced, not aliased)
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
85 lines
3.1 KiB
Plaintext
85 lines
3.1 KiB
Plaintext
// grad_norm.cu — GPU-resident grad-norm + clip-scale computation.
|
||
//
|
||
// Replaces the host-side memcpy_dtoh + sum-of-squares loop in
|
||
// `Mamba2AdamW::step_from_buffers`. The previous path did 9× dtoh per
|
||
// training step — each a stream sync barrier costing ~1ms+ on the hot
|
||
// path. This kernel suite keeps the entire computation on GPU.
|
||
//
|
||
// Workflow (per training step):
|
||
// 1. Zero `grad_norm_sq_d` once via `memset_zeros`.
|
||
// 2. For each of N gradient tensors:
|
||
// - launch `grad_norm_sq_phase1` → writes block-tree partials to scratch
|
||
// - launch `grad_norm_sq_phase2` → reduces partials, ACCUMULATES into grad_norm_sq_d
|
||
// 3. launch `grad_clip_scale` → reads grad_norm_sq_d, writes
|
||
// `grad_scale_d = min(1, max_norm / sqrt(norm_sq))`
|
||
// 4. AdamW kernels read `grad_scale_d` as a device pointer (not host scalar).
|
||
//
|
||
// All per-block reductions use shared-memory tree reduce per
|
||
// `feedback_no_atomicadd.md`. Phase 2 uses a single-thread `+=` for
|
||
// the cross-tensor accumulation, which is race-free because phase-2
|
||
// launches are sequential in the stream and each launch writes exactly
|
||
// one float to the accumulator slot.
|
||
|
||
#include <cuda_runtime.h>
|
||
|
||
#define GRAD_NORM_BLOCK 256
|
||
|
||
/// Phase 1: per-block tree-reduce of `x[]*x[]` → block_partials[blockIdx.x].
|
||
extern "C" __global__ void grad_norm_sq_phase1(
|
||
const float* __restrict__ x,
|
||
int n,
|
||
float* __restrict__ block_partials
|
||
) {
|
||
__shared__ float smem[GRAD_NORM_BLOCK];
|
||
int tid = threadIdx.x;
|
||
int gid = blockIdx.x * blockDim.x + tid;
|
||
smem[tid] = (gid < n) ? x[gid] * x[gid] : 0.0f;
|
||
__syncthreads();
|
||
for (int s = blockDim.x / 2; s > 0; s >>= 1) {
|
||
if (tid < s) smem[tid] += smem[tid + s];
|
||
__syncthreads();
|
||
}
|
||
if (tid == 0) block_partials[blockIdx.x] = smem[0];
|
||
}
|
||
|
||
/// Phase 2: single-block tree-reduce of `block_partials[0..n_blocks]`
|
||
/// → accumulates into `accum[0]`. `accumulate=1` adds; `accumulate=0`
|
||
/// overwrites (use for the first call in a step).
|
||
extern "C" __global__ void grad_norm_sq_phase2(
|
||
const float* __restrict__ block_partials,
|
||
int n_blocks,
|
||
float* __restrict__ accum,
|
||
int accumulate
|
||
) {
|
||
__shared__ float smem[GRAD_NORM_BLOCK];
|
||
int tid = threadIdx.x;
|
||
float v = 0.0f;
|
||
for (int i = tid; i < n_blocks; i += blockDim.x) {
|
||
v += block_partials[i];
|
||
}
|
||
smem[tid] = v;
|
||
__syncthreads();
|
||
for (int s = blockDim.x / 2; s > 0; s >>= 1) {
|
||
if (tid < s) smem[tid] += smem[tid + s];
|
||
__syncthreads();
|
||
}
|
||
if (tid == 0) {
|
||
float prev = (accumulate != 0) ? accum[0] : 0.0f;
|
||
accum[0] = prev + smem[0];
|
||
}
|
||
}
|
||
|
||
/// Final clamp: grad_scale = min(1, max_norm / sqrt(norm_sq + eps)).
|
||
/// One-thread kernel — writes a single float to `grad_scale_out[0]`.
|
||
/// AdamW kernels read this device pointer instead of a host scalar.
|
||
extern "C" __global__ void grad_clip_scale(
|
||
const float* __restrict__ norm_sq,
|
||
float max_norm,
|
||
float* __restrict__ grad_scale_out
|
||
) {
|
||
if (threadIdx.x != 0 || blockIdx.x != 0) return;
|
||
const float eps = 1e-12f;
|
||
float norm = sqrtf(norm_sq[0] + eps);
|
||
grad_scale_out[0] = (norm > max_norm) ? (max_norm / norm) : 1.0f;
|
||
}
|