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>
108 lines
3.7 KiB
Rust
108 lines
3.7 KiB
Rust
//! The training loop: sample sequences → forward `loss` → backward → grad clip
|
||
//! (with batch averaging) → AdamW step → zero grads; with an LR schedule,
|
||
//! periodic loss logging, and periodic checkpointing.
|
||
//!
|
||
//! The T5 model is single-sequence, so a "batch" of `batch_size` sequences is
|
||
//! handled by running forward+backward on each and letting the tape SUM their
|
||
//! grads (its fan-out rule); the clip pass then multiplies by `1/batch_size` to
|
||
//! recover the batch-mean gradient before clipping + the optimizer step.
|
||
|
||
#![cfg(not(no_cuda))]
|
||
|
||
use std::path::PathBuf;
|
||
use std::time::Instant;
|
||
|
||
use xtrain_model::{TinyTransformer, ids_tensor};
|
||
use xtrain_optim::GpuAdamW;
|
||
use xtrain_tensor::Device;
|
||
|
||
use crate::checkpoint;
|
||
use crate::clip::clip_grad_norm_gpu;
|
||
use crate::data::Corpus;
|
||
use crate::schedule::LrSchedule;
|
||
|
||
/// Knobs for a training run.
|
||
pub struct TrainConfig {
|
||
pub seq_len: usize,
|
||
pub batch_size: usize,
|
||
pub steps: usize,
|
||
pub schedule: LrSchedule,
|
||
pub weight_decay: f32,
|
||
pub max_grad_norm: f32,
|
||
pub log_every: usize,
|
||
/// Optional checkpoint path written every `ckpt_every` steps (and at the end).
|
||
pub ckpt_path: Option<PathBuf>,
|
||
pub ckpt_every: usize,
|
||
/// Seed for reproducible sequence sampling.
|
||
pub seed: u64,
|
||
}
|
||
|
||
/// Train `model` on `corpus` for `cfg.steps` AdamW steps. Returns the per-step
|
||
/// loss trace (one mean loss per step, read from the first sequence of the
|
||
/// batch — cheap and representative). Logs progress and checkpoints as configured.
|
||
pub fn train(
|
||
model: &TinyTransformer,
|
||
device: Device,
|
||
corpus: &Corpus,
|
||
cfg: &TrainConfig,
|
||
) -> Vec<f32> {
|
||
let params = model.params();
|
||
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;
|
||
let start = Instant::now();
|
||
let mut tokens_seen: u64 = 0;
|
||
|
||
for step in 0..cfg.steps {
|
||
let lr = cfg.schedule.lr(step);
|
||
|
||
// Accumulate grads over `batch_size` sequences (tape SUMs them).
|
||
let mut step_loss = 0.0f32;
|
||
for _ in 0..cfg.batch_size {
|
||
let (input, target) = corpus.sample(cfg.seq_len, &mut rng);
|
||
let ids = ids_tensor(&input, device);
|
||
let targets = ids_tensor(&target, device);
|
||
let loss = model.loss(&ids, &targets);
|
||
step_loss += read_scalar(&loss);
|
||
loss.backward();
|
||
tokens_seen += cfg.seq_len as u64;
|
||
}
|
||
step_loss *= inv_batch;
|
||
losses.push(step_loss);
|
||
|
||
// Average the summed grads (×1/batch) and clip to the global norm.
|
||
let gnorm = clip_grad_norm_gpu(¶ms, cfg.max_grad_norm, inv_batch);
|
||
opt.step(lr, ¶ms);
|
||
for p in ¶ms {
|
||
p.zero_grad();
|
||
}
|
||
|
||
if step % cfg.log_every == 0 || step == cfg.steps - 1 {
|
||
let elapsed = start.elapsed().as_secs_f32();
|
||
let tps = tokens_seen as f32 / elapsed.max(1e-6);
|
||
println!(
|
||
"step {step:5}/{}: loss {step_loss:.4} lr {lr:.2e} gnorm {gnorm:.3} \
|
||
({tps:.0} tok/s)",
|
||
cfg.steps
|
||
);
|
||
}
|
||
|
||
if let Some(path) = &cfg.ckpt_path {
|
||
if cfg.ckpt_every > 0 && (step + 1) % cfg.ckpt_every == 0 {
|
||
checkpoint::save(path, ¶ms).expect("checkpoint save");
|
||
}
|
||
}
|
||
}
|
||
|
||
if let Some(path) = &cfg.ckpt_path {
|
||
checkpoint::save(path, ¶ms).expect("final checkpoint save");
|
||
println!("saved checkpoint → {}", path.display());
|
||
}
|
||
losses
|
||
}
|
||
|
||
fn read_scalar(v: &xtrain_autodiff::tape::Var) -> f32 {
|
||
v.value().to_device(Device::Cpu).as_slice::<f32>()[0]
|
||
}
|