Gahow Wang
b0e397ca81
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>
77 lines
2.7 KiB
Rust
77 lines
2.7 KiB
Rust
// GPU AdamW parity (Phase T7): the device-side AdamW kernel (m/v on device, no
|
|
// host roundtrip) must produce the same update as the host reference
|
|
// `AdamW::step_host` given identical params + grads across several steps with a
|
|
// varying lr. This is the new correctness gate for the GPU optimizer; the host
|
|
// path itself is already pinned to PyTorch by xtrain-train's adamw_parity test.
|
|
//
|
|
// Gated #![cfg(not(no_cuda))] (runs on dash5; needs a GPU to link + launch).
|
|
#![cfg(not(no_cuda))]
|
|
|
|
use xtrain_autodiff::tape::Var;
|
|
use xtrain_cuda::device;
|
|
use xtrain_optim::{AdamW, GpuAdamW};
|
|
use xtrain_tensor::{Device, Tensor};
|
|
|
|
fn grad(step: usize, idx: usize, j: usize) -> f32 {
|
|
let s = (step * 13 + idx * 7 + j * 3) as f32;
|
|
(s * 0.123).sin() * 0.5
|
|
}
|
|
|
|
#[test]
|
|
fn gpu_adamw_matches_host() {
|
|
assert!(device::device_count().unwrap() > 0, "no CUDA device");
|
|
device::set_device(0).unwrap();
|
|
let dev = Device::Cuda(0);
|
|
|
|
let wd = 0.1f32;
|
|
// Two params of different sizes (exercises per-param device state).
|
|
let shapes: Vec<Vec<usize>> = vec![vec![2, 2], vec![3]];
|
|
let init: Vec<Vec<f32>> = vec![vec![0.5, -1.0, 2.0, 0.0], vec![1.5, -0.25, 0.75]];
|
|
|
|
// GPU side: leaf Vars on device.
|
|
let params: Vec<Var> = init
|
|
.iter()
|
|
.zip(&shapes)
|
|
.map(|(d, s)| Var::leaf(Tensor::from_slice(d, s).to_device(dev)))
|
|
.collect();
|
|
let mut gpu_opt = GpuAdamW::new(wd);
|
|
|
|
// Host reference.
|
|
let mut host_params = init.clone();
|
|
let mut host_opt = AdamW::new(0.0, wd);
|
|
|
|
for step in 0..15 {
|
|
let lr = 0.01 + 0.001 * step as f32; // varying lr
|
|
let grads: Vec<Vec<f32>> = shapes
|
|
.iter()
|
|
.enumerate()
|
|
.map(|(idx, s)| {
|
|
let n: usize = s.iter().product();
|
|
(0..n).map(|j| grad(step, idx, j)).collect()
|
|
})
|
|
.collect();
|
|
|
|
// Push grads onto the GPU Vars, run the device step, then clear.
|
|
for (p, (g, s)) in params.iter().zip(grads.iter().zip(&shapes)) {
|
|
p.zero_grad();
|
|
Var::push_grad(p, Tensor::from_slice(g, s).to_device(dev));
|
|
}
|
|
gpu_opt.step(lr, ¶ms);
|
|
for p in ¶ms {
|
|
p.zero_grad();
|
|
}
|
|
|
|
host_opt.step_host(lr, &mut host_params, &grads);
|
|
}
|
|
|
|
let mut max_err = 0.0f32;
|
|
for (p, hp) in params.iter().zip(&host_params) {
|
|
let got = p.value().to_device(Device::Cpu).as_slice::<f32>().to_vec();
|
|
for (a, b) in got.iter().zip(hp) {
|
|
max_err = max_err.max((a - b).abs());
|
|
}
|
|
}
|
|
println!("gpu vs host AdamW: max abs err = {max_err:.3e}");
|
|
assert!(max_err < 1e-6, "GPU AdamW diverged from host: {max_err:e}");
|
|
}
|