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>
Gahow Wang
b0e397ca81