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-train/src/clip.rs

@@ -73,8 +73,61 @@ mod gpu {
     }
 }
 
+#[cfg(not(no_cuda))]
+mod gpu_norm {
+    use super::clip_scale;
+    use xtrain_autodiff::tape::Var;
+    use xtrain_tensor::{DType, Device, Tensor};
+
+    /// GPU-side global-norm grad clip (Phase T7): compute the joint L2 norm of all
+    /// `pre_scale`-applied grads with a device reduction, then rescale every grad
+    /// in place by `pre_scale·clip_factor` — no per-step grad roundtrip to host
+    /// (only the single scalar norm comes back). Returns the post-pre_scale total
+    /// norm. Params without a grad contribute 0 and are skipped on rescale.
+    pub fn clip_grad_norm_gpu(params: &[Var], max_norm: f32, pre_scale: f32) -> f32 {
+        let device = params[0].value().device();
+        // sum-of-squares of the RAW grads accumulated on device.
+        let acc = Tensor::zeros(&[1], DType::F32, device);
+        for p in params {
+            if let Some(g) = p.grad() {
+                unsafe {
+                    xtrain_cuda::ffi::launch_sumsq_accum_f32(
+                        g.data_ptr() as *const f32,
+                        acc.data_ptr() as *mut f32,
+                        g.numel() as i32,
+                        std::ptr::null_mut(),
+                    );
+                }
+            }
+        }
+        xtrain_cuda::device::synchronize().expect("grad-norm reduce sync failed");
+        let raw_sumsq = acc.to_device(Device::Cpu).as_slice::<f32>()[0];
+        // Norm of the pre_scale-applied grads = pre_scale · sqrt(raw_sumsq).
+        let total = pre_scale * raw_sumsq.max(0.0).sqrt();
+        let factor = pre_scale * clip_scale(total, max_norm);
+        if (factor - 1.0).abs() >= f32::EPSILON {
+            for p in params {
+                if let Some(g) = p.grad() {
+                    unsafe {
+                        xtrain_cuda::ffi::launch_scale_inplace_f32(
+                            g.data_ptr() as *mut f32,
+                            factor,
+                            g.numel() as i32,
+                            std::ptr::null_mut(),
+                        );
+                    }
+                }
+            }
+            xtrain_cuda::device::synchronize().expect("grad rescale sync failed");
+        }
+        total
+    }
+}
+
 #[cfg(not(no_cuda))]
 pub use gpu::clip_grad_norm;
+#[cfg(not(no_cuda))]
+pub use gpu_norm::clip_grad_norm_gpu;
 
 #[cfg(test)]
 mod tests {

crates/xtrain-train/src/train_loop.rs

@@ -13,11 +13,11 @@ use std::path::PathBuf;
 use std::time::Instant;
 
 use xtrain_model::{TinyTransformer, ids_tensor};
-use xtrain_optim::AdamW;
+use xtrain_optim::GpuAdamW;
 use xtrain_tensor::Device;
 
 use crate::checkpoint;
-use crate::clip::clip_grad_norm;
+use crate::clip::clip_grad_norm_gpu;
 use crate::data::Corpus;
 use crate::schedule::LrSchedule;
 
@@ -47,7 +47,7 @@ pub fn train(
     cfg: &TrainConfig,
 ) -> Vec<f32> {
     let params = model.params();
-    let mut opt = AdamW::new(cfg.schedule.max_lr, cfg.weight_decay);
+    let mut opt = GpuAdamW::new(cfg.weight_decay);
     let mut rng = cfg.seed;
     let mut losses = Vec::with_capacity(cfg.steps);
     let inv_batch = 1.0 / cfg.batch_size as f32;
@@ -72,7 +72,7 @@ pub fn train(
         losses.push(step_loss);
 
         // Average the summed grads (×1/batch) and clip to the global norm.
-        let gnorm = clip_grad_norm(&params, cfg.max_grad_norm, inv_batch);
+        let gnorm = clip_grad_norm_gpu(&params, cfg.max_grad_norm, inv_batch);
         opt.step(lr, &params);
         for p in &params {
             p.zero_grad();