Neural network training

Deep dive

How it works

Builds on: AdamW, Learning-rate warmup and cosine decay, Mixed-precision training, ZeRO and FSDP sharding, Ring all-reduce, Data loading pipelines.

Parts:

• data loader — shuffles, shards per worker, decodes and prefetches fixed-size batches
• model — forward pass in low precision, loss computed in fp32
• loss scaler — multiplies the loss by S before backward, divides gradients by S after; skips the step on inf/NaN
• gradient sync — averages gradients across workers (all-reduce, or reduce-scatter when the optimizer is sharded)
• clipper — rescales gradients if their global L2 norm exceeds a limit
• optimizer — AdamW with first and second moments in fp32
• schedule — learning rate as a function of the step: linear warmup then cosine decay
• checkpointer — saves weights, optimizer moments, step, schedule and data position
• evaluator — fixed held-out set, no gradients, run every K steps

Operations: load, forward, backward, sync, clip, update, evaluate, checkpoint (steps below).

1. Load. Draw a batch; each of W workers takes its own shard so no example repeats within an epoch.
2. Forward and loss. Run in fp16 or bf16 autocast; reduce and compute the loss in fp32. Multiply by the loss scale S (fp16 only; bf16 usually needs none).
3. Backward. Gradients flow in low precision into fp32 buffers. Divide by S. If any value is inf or NaN, skip the update, halve S and continue; after N clean steps double S.
4. Sync. All-reduce the mean of the gradients across workers so every replica applies the same update (data parallelism).
5. Clip. Compute the global L2 norm of all gradients; if it exceeds the limit (1.0 is a common default), multiply all gradients by limit / norm (Pascanu et al., 2013).
6. Update. lr = schedule(step); AdamW on the fp32 master weights: moments m, v, bias correction, step, decoupled decay.
7. Evaluate. Every K steps run the held-out set with the model in eval mode (no dropout, no gradient), log loss and task metrics; keep the best checkpoint.
8. Checkpoint and resume. Save weights, both moments, the step counter, loss scale and data position; write to a temporary file then rename. Resuming from this state continues the same trajectory.

data ─► shard/prefetch ─► forward (fp16) ─► loss ×S ─► backward ─► ÷S ─► all-reduce ─► clip ─► AdamW(lr(step)) ─► weights
                                                                                                   │
                                   eval on held-out set ◄── every K steps ──── checkpoint ◄───────┘

Code

/// MLP in -> h (tanh) -> o, parameters in one flat vector: [w1 (in*h), b1 (h), w2 (h*o), b2 (o)].
pub struct Mlp { pub i: usize, pub h: usize, pub o: usize }

impl Mlp {
    pub fn n_params(&self) -> usize { self.i * self.h + self.h + self.h * self.o + self.o }

    /// Mean-squared-error loss over a batch (x: [n,i], y: [n,o]) and its gradient, by hand.
    pub fn loss_grad(&self, p: &[f32], x: &[f32], y: &[f32], g: &mut [f32]) -> f32 {
        let (i, h, o) = (self.i, self.h, self.o);
        let n = x.len() / i;
        let (w1, rest) = p.split_at(i * h);
        let (b1, rest) = rest.split_at(h);
        let (w2, b2) = rest.split_at(h * o);
        g.iter_mut().for_each(|v| *v = 0.0);
        let (gw1, grest) = g.split_at_mut(i * h);
        let (gb1, grest) = grest.split_at_mut(h);
        let (gw2, gb2) = grest.split_at_mut(h * o);
        let mut loss = 0.0;
        for r in 0..n {
            let xr = &x[r * i..(r + 1) * i];
            let a: Vec<f32> = (0..h).map(|j| (b1[j] + (0..i).map(|k| xr[k] * w1[k * h + j]).sum::<f32>()).tanh()).collect();
            let dout: Vec<f32> = (0..o).map(|c| {              // dL/dout for this row
                let out = b2[c] + (0..h).map(|j| a[j] * w2[j * o + c]).sum::<f32>();
                let d = out - y[r * o + c];
                loss += d * d / (n * o) as f32;
                2.0 * d / (n * o) as f32
            }).collect();
            for c in 0..o {
                gb2[c] += dout[c];
                for j in 0..h { gw2[j * o + c] += dout[c] * a[j]; }
            }
            for j in 0..h {                                     // backprop through tanh
                let da: f32 = (0..o).map(|c| dout[c] * w2[j * o + c]).sum();
                let dz = da * (1.0 - a[j] * a[j]);
                gb1[j] += dz;
                for k in 0..i { gw1[k * h + j] += dz * xr[k]; }
            }
        }
        loss
    }
}

pub struct AdamW { pub lr: f32, pub b1: f32, pub b2: f32, pub eps: f32, pub wd: f32, m: Vec<f32>, v: Vec<f32>, pub t: u32 }

impl AdamW {
    pub fn new(n: usize) -> Self { AdamW { lr: 1e-3, b1: 0.9, b2: 0.999, eps: 1e-8, wd: 0.01, m: vec![0.0; n], v: vec![0.0; n], t: 0 } }
    /// Decoupled weight decay: applied to the weights directly, not through the gradient.
    pub fn step(&mut self, p: &mut [f32], g: &[f32]) {
        self.t += 1;
        let (c1, c2) = (1.0 - self.b1.powi(self.t as i32), 1.0 - self.b2.powi(self.t as i32));
        for k in 0..p.len() {
            self.m[k] = self.b1 * self.m[k] + (1.0 - self.b1) * g[k];
            self.v[k] = self.b2 * self.v[k] + (1.0 - self.b2) * g[k] * g[k];
            p[k] -= self.lr * ((self.m[k] / c1) / ((self.v[k] / c2).sqrt() + self.eps) + self.wd * p[k]);
        }
    }
    /// Checkpoint = everything needed to resume bit-exactly: params, moments, step.
    pub fn save(&self, p: &[f32]) -> Vec<u8> {
        let mut b = self.t.to_le_bytes().to_vec();
        for v in p.iter().chain(&self.m).chain(&self.v) { b.extend(v.to_le_bytes()); }
        b
    }
    pub fn load(&mut self, bytes: &[u8], p: &mut [f32]) {
        self.t = u32::from_le_bytes(bytes[..4].try_into().unwrap());
        let f: Vec<f32> = bytes[4..].chunks(4).map(|c| f32::from_le_bytes(c.try_into().unwrap())).collect();
        let n = p.len();
        p.copy_from_slice(&f[..n]); self.m.copy_from_slice(&f[n..2 * n]); self.v.copy_from_slice(&f[2 * n..]);
    }
}

/// Linear warmup to `peak`, then cosine decay to `floor`.
pub fn lr_at(step: u32, warmup: u32, total: u32, peak: f32, floor: f32) -> f32 {
    if step < warmup { return peak * (step + 1) as f32 / warmup as f32; }
    let t = (step - warmup) as f32 / (total - warmup).max(1) as f32;
    floor + 0.5 * (peak - floor) * (1.0 + (std::f32::consts::PI * t).cos())
}

/// Clip the global gradient norm; returns the norm before clipping.
pub fn clip_grad_norm(g: &mut [f32], max: f32) -> f32 {
    let norm = g.iter().map(|v| v * v).sum::<f32>().sqrt();
    if norm > max { g.iter_mut().for_each(|v| *v *= max / norm); }
    norm
}

/// Data parallel: every worker computes grads on its shard, then all hold the mean.
pub fn allreduce_mean(grads: &mut [Vec<f32>]) {
    let w = grads.len() as f32;
    for k in 0..grads[0].len() {
        let s: f32 = grads.iter().map(|g| g[k]).sum::<f32>() / w;
        grads.iter_mut().for_each(|g| g[k] = s);
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    fn lcg(s: &mut u64) -> f32 { *s = s.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407); ((*s >> 40) as f32 / (1u64 << 24) as f32) * 2.0 - 1.0 }

    #[test]
    fn gradient_check() {
        let net = Mlp { i: 3, h: 4, o: 2 };
        let mut s = 1;
        let p: Vec<f32> = (0..net.n_params()).map(|_| lcg(&mut s) * 0.5).collect();
        let x: Vec<f32> = (0..15).map(|_| lcg(&mut s)).collect();
        let y: Vec<f32> = (0..10).map(|_| lcg(&mut s)).collect();
        let mut g = vec![0.0; p.len()];
        net.loss_grad(&p, &x, &y, &mut g);
        let eps = 1e-3;
        for k in 0..p.len() { // central differences, all parameters
            let (mut a, mut b, mut tmp) = (p.clone(), p.clone(), vec![0.0; p.len()]);
            a[k] += eps; b[k] -= eps;
            let num = (net.loss_grad(&a, &x, &y, &mut tmp) - net.loss_grad(&b, &x, &y, &mut tmp)) / (2.0 * eps);
            assert!((num - g[k]).abs() < 2e-3, "param {k}: numeric {num} vs analytic {}", g[k]);
        }
    }

    #[test]
    fn training_reduces_loss_and_resumes() {
        let net = Mlp { i: 2, h: 16, o: 1 };
        let mut s = 3;
        let x: Vec<f32> = (0..256).map(|_| lcg(&mut s)).collect();
        let y: Vec<f32> = x.chunks(2).map(|v| v[0] * v[1]).collect(); // target: product
        let mut p: Vec<f32> = (0..net.n_params()).map(|_| lcg(&mut s) * 0.5).collect();
        let (mut opt, mut g) = (AdamW::new(p.len()), vec![0.0; p.len()]);
        let total = 400;
        let first = net.loss_grad(&p, &x, &y, &mut g);
        let mut ckpt = None;
        for step in 0..total {
            let loss = net.loss_grad(&p, &x, &y, &mut g);
            clip_grad_norm(&mut g, 1.0);
            opt.lr = lr_at(step, 20, total, 1e-2, 1e-4);
            opt.step(&mut p, &g);
            if step == 199 { ckpt = Some((opt.save(&p), loss)); }
        }
        let last = net.loss_grad(&p, &x, &y, &mut g);
        assert!(last < 0.2 * first, "loss {first} -> {last}");
        // resume from the step-200 checkpoint, rerun the rest: same final parameters
        let (bytes, _) = ckpt.unwrap();
        let (mut p2, mut opt2) = (vec![0.0; p.len()], AdamW::new(p.len()));
        opt2.load(&bytes, &mut p2);
        for step in 200..total {
            net.loss_grad(&p2, &x, &y, &mut g);
            clip_grad_norm(&mut g, 1.0);
            opt2.lr = lr_at(step, 20, total, 1e-2, 1e-4);
            opt2.step(&mut p2, &g);
        }
        assert_eq!(p, p2);
    }
}

The first test compares every analytic gradient with central finite differences; the second checks the loss falls to under a fifth and that resuming from a step-200 checkpoint reproduces the final weights exactly. Mixed precision and the real all-reduce are left out of the code; allreduce_mean shows what each worker ends up holding.

Complexity and constants

• Forward plus backward costs about 3× the forward FLOPs (backward is about twice the forward); for dense transformers a common rule of thumb is about 6 FLOPs per parameter per token (order of).
• Optimizer step is O(params) and memory-bound: it reads and writes about 16 bytes per parameter (fp32 AdamW, derived).
• Memory per parameter, derived: fp32 weights, gradients and two moments = 16 bytes; mixed precision with fp32 master copy = 2 (weights) + 2 (gradients) + 4 (master) + 8 (moments) = 16 bytes; activations come on top and often exceed this.
• All-reduce: a ring moves about 2·(W−1)/W of the gradient bytes per worker, nearly independent of W; see ring all-reduce.

Efficiency

• Cost: order of 6 FLOPs per parameter per token for a dense transformer; 16 bytes of state per parameter. Bounds: compute on GPU for large models, input pipeline for small models, interconnect for many workers.
• Levers:
  • mixed precision (bf16 or fp16 with loss scaling) — halves activation and weight traffic and uses tensor cores — any GPU generation that supports it
  • prefetch and parallel decode in the loader — hides I/O behind compute — whenever GPU utilisation drops between steps
  • gradient accumulation — larger effective batch without more memory — memory-limited, small batches
  • ZeRO or FSDP sharding — divides optimizer state (and optionally gradients, weights) across workers — models whose 16 bytes per parameter do not fit
  • activation checkpointing — trades about one extra forward for much less activation memory — memory-bound runs
  • overlap communication with backward — bucketed all-reduce while earlier layers still compute — many workers
• Naive → standard → state of the art:
  • one process, fp32, constant learning rate, no checkpoints · slow and fragile
  • data-parallel, AdamW, warmup-cosine, mixed precision, clipping, periodic checkpoints · the default recipe
  • sharded optimizer state, overlapped communication, resumable data loader, automatic restart · large-scale runs (Goyal et al. 2017 for the large-batch recipe)

Implementation notes

• Apply weight decay to matrices only; leave biases and normalisation scales undecayed.
• Seed everything and store the RNG and data position in the checkpoint; otherwise a resumed run repeats or skips examples.
• Warmup length of order of 1-5 % of the steps; Post-LayerNorm transformers especially need it (Xiong et al., 2020).
• Compute the clip norm after the all-reduce and after unscaling, over all parameters together, not per tensor.
• Atomic checkpoint writes: write ckpt.tmp, fsync, rename; keep the last few and the best.
• Evaluate with a fixed set and fixed seeds so curves are comparable across runs.
• Log the gradient norm every step: spikes show up there before the loss diverges.

Pitfalls

• Resuming without the optimizer moments or step counter: the first steps use a cold Adam and the schedule restarts, usually causing a loss spike.
• Clipping before unscaling: the norm is S times too large, so everything is clipped.
• A skipped overflow step that still advances the schedule: the learning rate drifts from the intended curve.
• Evaluating in train mode, or on data that overlaps training (in games, positions from the same game in both sets): metrics look better than reality; split by game.
• Changing the batch size or worker count without rescaling the learning rate.
• Averaging per-worker losses of unequal batch sizes instead of weighting by count.

Variants and alternatives

In production

• PyTorch with DistributedDataParallel and automatic mixed precision for most open training code; Megatron-LM for large transformer runs; Hugging Face Accelerate and PyTorch Lightning wrap the same loop; nanoGPT is a short, readable reference.

Implementations

Sources

• Ilya Loshchilov, Frank Hutter, Decoupled Weight Decay Regularization, ICLR 2019, arXiv:1711.05101. paper
• Ilya Loshchilov, Frank Hutter, SGDR: Stochastic Gradient Descent with Warm Restarts, ICLR 2017, arXiv:1608.03983. paper
• Paulius Micikevicius et al., Mixed Precision Training, ICLR 2018, arXiv:1710.03740. paper
• Priya Goyal et al., Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour, arXiv:1706.02677, 2017. paper
• Razvan Pascanu, Tomas Mikolov, Yoshua Bengio, On the difficulty of training Recurrent Neural Networks, ICML 2013, arXiv:1211.5063. paper
• Diederik P. Kingma, Jimmy Ba, Adam: A Method for Stochastic Optimization, ICLR 2015, arXiv:1412.6980. paper
• Ruibin Xiong et al., On Layer Normalization in the Transformer Architecture, ICML 2020, arXiv:2002.04745. paper
• PyTorch documentation, Automatic Mixed Precision examples and Distributed Data Parallel. AMP · DDP