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>
248 lines
9.6 KiB
Rust
248 lines
9.6 KiB
Rust
//! Hand-written AdamW optimizer (Phase T6).
|
||
//!
|
||
//! AdamW = Adam with **decoupled** weight decay (Loshchilov & Hutter, 2019): the
|
||
//! weight-decay term is applied directly to the parameter, NOT folded into the
|
||
//! gradient (so it does not interact with the adaptive `v` denominator). This
|
||
//! matches `torch.optim.AdamW`.
|
||
//!
|
||
//! Update for parameter `θ` at step `t` (1-indexed), with gradient `g`:
|
||
//! ```text
|
||
//! m ← β1·m + (1−β1)·g
|
||
//! v ← β2·v + (1−β2)·g²
|
||
//! m̂ ← m / (1 − β1ᵗ) (bias correction)
|
||
//! v̂ ← v / (1 − β2ᵗ)
|
||
//! θ ← θ − lr·( m̂ / (√v̂ + ε) + wd·θ )
|
||
//! ```
|
||
//! The `lr·wd·θ` term is the decoupled decay. Note PyTorch applies decay as
|
||
//! `θ ← θ·(1 − lr·wd)` then the Adam step; both are algebraically the same
|
||
//! first-order update — we fold decay into the single subtraction above, which
|
||
//! is what PyTorch's default (`maximize=False`, no `amsgrad`) computes.
|
||
//!
|
||
//! The math operates on flat host `f32` buffers ([`AdamW::step_host`]) so it is
|
||
//! unit-testable on a GPU-less host; [`AdamW::step`] is a thin wrapper that
|
||
//! round-trips each parameter's value/grad through the GPU tensor and is gated
|
||
//! behind `not(no_cuda)`.
|
||
|
||
/// Per-parameter optimizer state: the first (`m`) and second (`v`) moment
|
||
/// estimates, one f32 per element, kept flat (matching the parameter layout).
|
||
struct ParamState {
|
||
m: Vec<f32>,
|
||
v: Vec<f32>,
|
||
}
|
||
|
||
/// Decoupled-weight-decay Adam. One instance owns the moment state for a fixed
|
||
/// list of parameters, keyed by their index in the slice passed to `step`
|
||
/// (the model's stable `params()` order).
|
||
pub struct AdamW {
|
||
pub lr: f32,
|
||
beta1: f32,
|
||
beta2: f32,
|
||
eps: f32,
|
||
weight_decay: f32,
|
||
/// Global step count (shared across all params for bias correction).
|
||
t: u64,
|
||
/// Lazily sized to the parameter list on the first `step`.
|
||
state: Vec<ParamState>,
|
||
}
|
||
|
||
impl AdamW {
|
||
/// PyTorch-default hyperparameters except `lr`/`weight_decay`, which you set
|
||
/// (β1=0.9, β2=0.999, ε=1e-8).
|
||
pub fn new(lr: f32, weight_decay: f32) -> Self {
|
||
Self::with_betas(lr, weight_decay, 0.9, 0.999, 1e-8)
|
||
}
|
||
|
||
pub fn with_betas(lr: f32, weight_decay: f32, beta1: f32, beta2: f32, eps: f32) -> Self {
|
||
Self {
|
||
lr,
|
||
beta1,
|
||
beta2,
|
||
eps,
|
||
weight_decay,
|
||
t: 0,
|
||
state: Vec::new(),
|
||
}
|
||
}
|
||
|
||
/// Current global step (number of `step` calls so far).
|
||
pub fn step_count(&self) -> u64 {
|
||
self.t
|
||
}
|
||
|
||
/// Pure-host AdamW step over flat parameter/gradient buffers. `params[i]` is
|
||
/// updated in place using `grads[i]`; both are the i-th parameter's elements
|
||
/// in the model's stable order. Lazily allocates moment state on first call.
|
||
///
|
||
/// This is the testable core — no GPU, no autograd. `lr` is passed per call
|
||
/// so a schedule can vary it each step.
|
||
pub fn step_host(&mut self, lr: f32, params: &mut [Vec<f32>], grads: &[Vec<f32>]) {
|
||
assert_eq!(params.len(), grads.len(), "param/grad count mismatch");
|
||
if self.state.is_empty() {
|
||
self.state = params
|
||
.iter()
|
||
.map(|p| ParamState {
|
||
m: vec![0.0; p.len()],
|
||
v: vec![0.0; p.len()],
|
||
})
|
||
.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 (i, (p, g)) in params.iter_mut().zip(grads).enumerate() {
|
||
assert_eq!(p.len(), g.len(), "param/grad len mismatch at {i}");
|
||
let st = &mut self.state[i];
|
||
for j in 0..p.len() {
|
||
let gj = g[j];
|
||
st.m[j] = self.beta1 * st.m[j] + (1.0 - self.beta1) * gj;
|
||
st.v[j] = self.beta2 * st.v[j] + (1.0 - self.beta2) * gj * gj;
|
||
let mhat = st.m[j] / bc1;
|
||
let vhat = st.v[j] / bc2;
|
||
// Decoupled weight decay: decay term uses the *current* param,
|
||
// matching PyTorch's `p ← p − lr·wd·p` applied alongside the step.
|
||
p[j] -= lr * (mhat / (vhat.sqrt() + self.eps) + self.weight_decay * p[j]);
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
#[cfg(not(no_cuda))]
|
||
mod gpu {
|
||
use super::AdamW;
|
||
use xtrain_autodiff::tape::Var;
|
||
use xtrain_tensor::{DType, Device, Tensor};
|
||
|
||
impl AdamW {
|
||
/// Apply one AdamW step to every parameter `Var`, using `lr` for this step
|
||
/// (so an LR schedule can vary it). Pulls each param's value and `.grad()`
|
||
/// to the host, runs [`AdamW::step_host`], and writes the updated value
|
||
/// back with `set_value`. A param with no grad is fed a zero grad, so the
|
||
/// Adam term vanishes and only decoupled weight decay applies (the model's
|
||
/// params all receive grads each step, so this is just a safety default).
|
||
///
|
||
/// 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>> =
|
||
params.iter().map(|p| p.value().shape().to_vec()).collect();
|
||
|
||
let mut host_params: Vec<Vec<f32>> = params
|
||
.iter()
|
||
.map(|p| p.value().to_device(Device::Cpu).as_slice::<f32>().to_vec())
|
||
.collect();
|
||
let host_grads: Vec<Vec<f32>> = params
|
||
.iter()
|
||
.zip(&host_params)
|
||
.map(|(p, hp)| match p.grad() {
|
||
Some(g) => g.to_device(Device::Cpu).as_slice::<f32>().to_vec(),
|
||
None => vec![0.0; hp.len()], // no grad → no update this step
|
||
})
|
||
.collect();
|
||
|
||
self.step_host(lr, &mut host_params, &host_grads);
|
||
|
||
for ((p, data), shape) in params.iter().zip(&host_params).zip(&shapes) {
|
||
let t = Tensor::from_slice(data, shape).to_device(device);
|
||
p.set_value(t);
|
||
}
|
||
}
|
||
}
|
||
|
||
/// 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;
|