perf: GPU AdamW + grad-norm
Eliminate the per-step GPU↔host roundtrip of every parameter/gradient. - optim.cu: adamw_step (m/v on device, in-place param update), sumsq_accum (block-reduced global grad sum-of-squares), scale_inplace. - GpuAdamW: device m/v state per param; step launches the kernel reading each param's .grad() and rewriting the param buffer in place — no host roundtrip. Host AdamW kept as the torch-parity reference. - clip_grad_norm_gpu: device sum-of-squares reduction (only the scalar norm comes back), in-place rescale of grads by pre_scale·clip_factor. - train_loop: use GpuAdamW + clip_grad_norm_gpu. - test: GPU AdamW vs host reference parity (max abs err < 1e-6). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
@@ -34,6 +34,7 @@ fn main() {
|
||||
.file("../../csrc/ops/gemm.cu")
|
||||
.file("../../csrc/ops/nn.cu")
|
||||
.file("../../csrc/ops/model.cu")
|
||||
.file("../../csrc/ops/optim.cu")
|
||||
.compile("xtrain_cuda_kernels");
|
||||
}
|
||||
|
||||
|
||||
@@ -212,6 +212,34 @@ unsafe extern "C" {
|
||||
);
|
||||
}
|
||||
|
||||
// GPU-side optimizer kernels (csrc/ops/optim.cu): AdamW step (m/v on device) and
|
||||
// the global grad-norm reduction + in-place rescale (Phase T7).
|
||||
#[cfg(not(no_cuda))]
|
||||
unsafe extern "C" {
|
||||
// One in-place AdamW step over a parameter tensor of `n` elements. `bc1`/`bc2`
|
||||
// are the bias-correction denominators 1-beta^t.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn launch_adamw_step_f32(
|
||||
p: *mut f32,
|
||||
g: *const f32,
|
||||
m: *mut f32,
|
||||
v: *mut f32,
|
||||
lr: f32,
|
||||
b1: f32,
|
||||
b2: f32,
|
||||
eps: f32,
|
||||
wd: f32,
|
||||
bc1: f32,
|
||||
bc2: f32,
|
||||
n: i32,
|
||||
s: CudaStream,
|
||||
);
|
||||
// acc += sum_i g[i]^2 (acc is one f32 on device, pre-zeroed). atomicAdd.
|
||||
pub fn launch_sumsq_accum_f32(g: *const f32, acc: *mut f32, n: i32, s: CudaStream);
|
||||
// In-place scalar scale: x[i] *= factor.
|
||||
pub fn launch_scale_inplace_f32(x: *mut f32, factor: f32, n: i32, s: CudaStream);
|
||||
}
|
||||
|
||||
// cuBLAS — the production GEMM backend (Phase T7) and the correctness oracle the
|
||||
// T3 GEMM tests still compare against. Declared (and linked, see build.rs) only
|
||||
// when CUDA is compiled in.
|
||||
|
||||
Reference in New Issue
Block a user