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/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();