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>

crates/xtrain-optim/src/lib.rs

@@ -113,7 +113,7 @@ impl AdamW {
 mod gpu {
     use super::AdamW;
     use xtrain_autodiff::tape::Var;
-    use xtrain_tensor::{Device, Tensor};
+    use xtrain_tensor::{DType, Device, Tensor};
 
     impl AdamW {
         /// Apply one AdamW step to every parameter `Var`, using `lr` for this step
@@ -125,6 +125,10 @@ mod gpu {
         ///
         /// Does NOT zero grads — the caller does that (matching the GD-step
         /// template in the T5 overfit test).
+        ///
+        /// This is the host-roundtrip reference path; training uses
+        /// [`GpuAdamW`] (kernel, m/v on device). Both are checked against the
+        /// torch parity in tests.
         pub fn step(&mut self, lr: f32, params: &[Var]) {
             let device = params[0].value().device();
             let shapes: Vec<Vec<usize>> =
@@ -151,4 +155,93 @@ mod gpu {
             }
         }
     }
+
+    /// GPU AdamW (Phase T7): the optimizer state (m/v moments) lives on the device
+    /// as one tensor pair per parameter, and the update runs as a CUDA kernel that
+    /// reads each param's `.grad()` and rewrites the param buffer in place — no
+    /// per-step GPU↔host roundtrip of params/grads. Same math as
+    /// [`AdamW::step_host`] (the parity reference).
+    pub struct GpuAdamW {
+        beta1: f32,
+        beta2: f32,
+        eps: f32,
+        weight_decay: f32,
+        t: u64,
+        /// Per-parameter (m, v) device buffers, sized lazily on first step.
+        state: Vec<(Tensor, Tensor)>,
+    }
+
+    impl GpuAdamW {
+        /// PyTorch-default betas/eps; you set lr (per-step) + weight decay.
+        pub fn new(weight_decay: f32) -> Self {
+            Self {
+                beta1: 0.9,
+                beta2: 0.999,
+                eps: 1e-8,
+                weight_decay,
+                t: 0,
+                state: Vec::new(),
+            }
+        }
+
+        pub fn step_count(&self) -> u64 {
+            self.t
+        }
+
+        /// One in-place AdamW step over every parameter `Var` at learning rate
+        /// `lr`. Updates the param value buffer and the device m/v state via the
+        /// `adamw_step_f32` kernel. Params are mutated in place, so the leaf `Var`
+        /// identities stay stable across steps (no `set_value`). Does NOT zero
+        /// grads — the caller does. A param without a grad is skipped this step.
+        pub fn step(&mut self, lr: f32, params: &[Var]) {
+            let device = params[0].value().device();
+            if self.state.is_empty() {
+                self.state = params
+                    .iter()
+                    .map(|p| {
+                        let shape = p.value().shape().to_vec();
+                        (
+                            Tensor::zeros(&shape, DType::F32, device),
+                            Tensor::zeros(&shape, DType::F32, device),
+                        )
+                    })
+                    .collect();
+            }
+            assert_eq!(self.state.len(), params.len(), "param count changed");
+
+            self.t += 1;
+            let bc1 = 1.0 - self.beta1.powi(self.t as i32);
+            let bc2 = 1.0 - self.beta2.powi(self.t as i32);
+
+            for (p, (m, v)) in params.iter().zip(&self.state) {
+                let g = match p.grad() {
+                    Some(g) => g,
+                    None => continue,
+                };
+                let pv = p.value();
+                let n = pv.numel() as i32;
+                unsafe {
+                    xtrain_cuda::ffi::launch_adamw_step_f32(
+                        pv.data_ptr() as *mut f32,
+                        g.data_ptr() as *const f32,
+                        m.data_ptr() as *mut f32,
+                        v.data_ptr() as *mut f32,
+                        lr,
+                        self.beta1,
+                        self.beta2,
+                        self.eps,
+                        self.weight_decay,
+                        bc1,
+                        bc2,
+                        n,
+                        std::ptr::null_mut(),
+                    );
+                }
+            }
+            xtrain_cuda::device::synchronize().expect("adamw step sync failed");
+        }
+    }
 }
+
+#[cfg(not(no_cuda))]
+pub use gpu::GpuAdamW;

crates/xtrain-optim/tests/adamw_gpu.rs

@@ -0,0 +1,76 @@
+// 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, &params);
+        for p in &params {
+            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}");
+}