AdamW

Deep dive

How it works

In: weights w, gradient g, learning rate lr, step count t · Out: updated weights · State: m and v per parameter (fp32) plus t · Cost: O(P) per step, 28 B/param of memory traffic in fp32.

1. Update the first moment m (momentum) and the second moment v (per-coordinate gradient scale).
2. Bias-correct both, since they start at 0.
3. Shrink the weights by lr·λ (decoupled decay).
4. Subtract lr·m̂/(√v̂+ε).

Per step t, with gradient g:

m ← β1·m + (1−β1)·g                 first moment (momentum)
v ← β2·v + (1−β2)·g²                second moment (per-coordinate scale)
m̂ = m / (1−β1ᵗ),  v̂ = v / (1−β2ᵗ)    bias correction: m and v start at 0
w ← w − lr·λ·w                      decoupled decay (PyTorch multiplies λ by lr)
w ← w − lr·m̂ / (√v̂ + ε)

• The update size is bounded by about lr per coordinate, since |m̂/√v̂| ≲ 1 for stationary gradients. That makes lr a step size in weight units.
• β2 sets the averaging window, about 1/(1−β2): 1000 steps at 0.999 and 20 at 0.95. A short window reacts quickly to gradient-scale jumps, which suppresses LLM loss spikes.
• The equilibrium weight norm is set by lr·λ. Changing lr also changes the effective regularization.

Code

pub struct AdamW {
    pub beta1: f32, pub beta2: f32, pub eps: f32, pub weight_decay: f32,
    t: i32, m: Vec<f32>, v: Vec<f32>,
}

impl AdamW {
    pub fn new(n: usize) -> Self {
        Self { beta1: 0.9, beta2: 0.95, eps: 1e-8, weight_decay: 0.1, t: 0, m: vec![0.0; n], v: vec![0.0; n] }
    }

    /// Fused update matching torch.optim.AdamW. `decay = false` for biases, norm scales, (often) embeddings.
    pub fn step(&mut self, w: &mut [f32], g: &[f32], lr: f32, decay: bool) {
        self.t += 1;
        let (b1, b2) = (self.beta1, self.beta2);
        let step = lr / (1.0 - b1.powi(self.t));                 // fold bias correction 1 into lr
        let bc2_sqrt = (1.0 - b2.powi(self.t)).sqrt();
        let shrink = if decay { 1.0 - lr * self.weight_decay } else { 1.0 };
        for i in 0..w.len() {
            let gi = g[i];
            self.m[i] = b1 * self.m[i] + (1.0 - b1) * gi;
            self.v[i] = b2 * self.v[i] + (1.0 - b2) * gi * gi;
            let denom = self.v[i].sqrt() / bc2_sqrt + self.eps;  // eps outside the sqrt, as in PyTorch
            w[i] = w[i] * shrink - step * self.m[i] / denom;
        }
    }
}

/// Global-norm clipping across ALL parameters (and all shards) before the step.
pub fn clip_global_norm(grads: &mut [&mut [f32]], max_norm: f32) -> f32 {
    let sq: f64 = grads.iter().flat_map(|g| g.iter()).map(|&x| (x as f64) * (x as f64)).sum();
    let norm = sq.sqrt() as f32;
    if norm > max_norm {
        let s = max_norm / (norm + 1e-6);
        grads.iter_mut().for_each(|g| g.iter_mut().for_each(|x| *x *= s));
    }
    norm
}

fn main() {
    // Two coordinates with gradients 1e4x apart still move at about the same speed.
    let mut w = vec![1.0f32, 1.0];
    let mut opt = AdamW::new(2);
    for _ in 0..100 {
        let mut g = vec![1e-3 * w[0], 10.0 * w[1]];
        clip_global_norm(&mut [&mut g[..]], 1e3);
        opt.step(&mut w, &g, 1e-2, true);
    }
    println!("w after 100 steps: {:?}", w); // both about 0.09: same speed despite 1e4x gradient ratio
    assert!(w[0] < 0.5 && w[1].abs() < 0.5);
}

Complexity and constants

• Traffic per param in fp32: read w, g, m, v and write w, m, v = 28 B. For 7B params that is 196 GB, about 60 ms at 3.35 TB/s if unsharded. Under ZeRO with N ranks it divides by N.
• Model-state memory with mixed precision: 2 (bf16 w) + 2 (bf16 g) + 4 (fp32 master) + 4 (m) + 4 (v) = 16 B/param (ZeRO paper accounting).
• 8-bit block-wise states (Dettmers et al.) cut m and v to about 2 B/param total with matching quality on the reported LLM and vision benchmarks.

Efficiency

• Cost: per parameter the fp32 step reads w, g, m, v and writes w, m, v: 28 B, so 196 GB for 7B parameters, about 60 ms at 3.35 TB/s (H100 SXM HBM3 bandwidth, NVIDIA spec). Model state is 16 B/param with mixed precision (ZeRO paper accounting). Bound: memory bandwidth, not compute.
• Levers:
  • fused multi-tensor kernel — one pass over all tensors instead of about 10 kernel launches per tensor (this page's count) — GPU training with many small tensors
  • shard states with ZeRO/FSDP — per-rank state divided by N ranks (ZeRO paper) — multi-GPU training
  • 8-bit block-wise m and v — states from 8 B to about 2 B/param (Dettmers et al.) — optimizer memory is the limit
  • factored second moment (Adafactor) — O(n+m) state for an n×m matrix (Shazeer and Stern) — very large matrices
  • global-norm clipping fused with the step — one fewer pass over gradients — every run with clipping
• Naive → standard → state of the art:
  • per-tensor Python loop · about 10 launches per tensor, launch-bound
  • fused or foreach multi-tensor kernel · one memory-bound pass, 28 B/param
  • fused kernel with sharded or 8-bit states · state divided by N ranks or reduced to about 2 B/param

Implementation notes

• Use the multi-tensor fused kernel (torch.optim.AdamW(fused=True), apex/DeepSpeed FusedAdam). A per-tensor Python loop launches about 10 kernels per parameter tensor and can cost more than the update itself.
• Keep m, v and the master weights in fp32. In bf16, v (which holds g² ≈ 1e-10) loses its meaning and ε dominates.
• Two parameter groups: decay for matmul weights, no decay for biases and norm γ/β. Embedding decay varies by recipe.
• Clip on the global norm across all shards and pipeline stages. Under FSDP use the module's clip_grad_norm_, which all-reduces the squared norms.
• Checkpoint the step count t along with m and v. Bias correction depends on it.

Pitfalls

• torch.optim.Adam(weight_decay=…) is L2 regularization, not AdamW.
• λ semantics differ. PyTorch and Optax shrink by lr·λ per step, so λ = 0.1 with lr = 3e-4 means a factor of 1 − 3e-5. The paper scales decay by the schedule multiplier only. Recipes do not port 1:1 between frameworks.
• β2 = 0.999 on large LLMs: rare large gradients after a long quiet stretch blow up m/√v and cause loss spikes. 0.95 is the common fix.
• ε too small with fp16 states underflows. ε too large (1e-4+) turns Adam into momentum SGD with lr/ε scaling.
• Changing batch size: the Adam learning rate scales roughly with √(batch ratio), not linearly as for SGD.

Variants and alternatives

In production

torch.optim.AdamW (foreach and fused CUDA kernels), DeepSpeed FusedAdam and CPUAdam (ZeRO-Offload), bitsandbytes AdamW8bit, Optax adamw. Trained GPT-3 (β2 0.95, λ 0.1), LLaMA 1-3, Mistral and most open LLMs.

Implementations

Sources

• Kingma and Ba, Adam: A Method for Stochastic Optimization, ICLR 2015. paper
• Loshchilov and Hutter, Decoupled Weight Decay Regularization, ICLR 2019. paper
• Touvron et al., LLaMA: Open and Efficient Foundation Language Models, 2023. paper
• Dettmers, Lewis, Shleifer, Zettlemoyer, 8-bit Optimizers via Block-wise Quantization, ICLR 2022. paper · code
• Shazeer and Stern, Adafactor: Adaptive Learning Rates with Sublinear Memory Cost, ICML 2018. paper
• You et al., Large Batch Optimization for Deep Learning: Training BERT in 76 minutes (LAMB), ICLR 2020. paper
• Liu et al., Swin Transformer, ICCV 2021 (AdamW training recipe). paper
• Radford et al., Robust Speech Recognition via Large-Scale Weak Supervision (Whisper), 2022 (AdamW with gradient-norm clipping). paper