Training Backprop

The batch dimension (32) does not appear in the gradient shape — it gets summed out during the X^T · dY multiply. Every sample in the batch contributes to the same gradient tensor.

The gradient of L with respect to W₁ always has the same shape as W₁: [768, 3072]. It doesn't matter how many layers come after — the chain rule multiplies everything together and the batch dimension gets summed out.

Every gradient element is the same size as the parameter — so the backward pass computes exactly as many gradient values as there are parameters. For a 7B model, that's 7 billion gradient elements per training step.

Since b is broadcast to every row of Y, the gradient flows back from every row. The chain rule requires summing over the batch: ∂L/∂b = ∑_i ∂L/∂Y_i. This is a sum, not a mean — the loss is already averaged over the batch if needed.

This is why activation memory dominates training — every linear layer must save its input for the backward pass. A transformer with 32 layers saves these activations 32 times. Activation checkpointing trades compute for memory by recomputing these instead of storing them.

Cross-entropy is logarithmic — being very wrong is dramatically more expensive than being slightly wrong. This is by design: the −log function has an asymptote at 0, so assigning near-zero probability to the correct class is catastrophically penalized. This is also why label smoothing helps — it prevents the model from trying to push p_c all the way to 1.0.

javascript
function crossEntropy(logits, target) {
  const maxZ = Math.max(...logits);
  const exps = logits.map(z => Math.exp(z - maxZ));
  const sumExp = exps.reduce((a, b) => a + b);
  const logProb = (logits[target] - maxZ) - Math.log(sumExp);
  return -logProb;
}

Line 2 is the bug. Computing Math.exp(z) directly on raw logits overflows to Infinity when z > ~709 (FP64 limit). The fix is to subtract the max logit first: logits.map(z => Math.exp(z - maxZ)). This is the log-sum-exp trick — it produces the same softmax probabilities but avoids overflow. Every production softmax implementation uses this trick.

Perplexity has an intuitive interpretation: the model is "as confused as if it were choosing uniformly between ~12 options at each token." A perfect model has perplexity 1 (loss = 0). GPT-3 achieves ~20 perplexity on common benchmarks. A model that always guesses uniformly over V=50k tokens has perplexity 50,000.

The gradient pushes the correct class logit up (negative gradient = increase) and all other logits down (positive gradient = decrease). The magnitudes sum to zero: −0.341 + 0.242 + 0.099 = 0. This is a consequence of probabilities summing to 1.

This is O(seq²) — the quadratic bottleneck of attention. With 32 heads, the total is 32 × 4.19M = 134M elements. At FP32 that's 536 MB just for one layer's attention scores.

The forward pass has only 2 matmuls (QK^T and PV), so the backward pass does roughly 2× the FLOPs of the forward. This "backward is ~2× forward" rule holds approximately across the entire transformer.

The Jacobian of softmax is ∂p_i/∂z_j = p_i(δ_ij − p_j), which can be written as diag(p) − pp^T. This is a dense matrix — changing any one logit z_j affects all probabilities. This is fundamentally different from ReLU (diagonal Jacobian) and is one reason softmax is relatively expensive in the backward pass.

The order is: dV (only needs P and dO, both available), dP (needs dO and V), dS (needs dP and P from softmax backward), dQ (needs dS and K), dK (needs dS and Q). Note that dV and dP can be computed in parallel since they have no dependency on each other, but dS must come before dQ and dK.

At the midpoint of the cosine decay, the LR is exactly halfway between lr_max and lr_min. The cosine shape decays slowly at first, then fast in the middle, then slowly again at the end — matching the intuition that late training needs fine-grained updates.

javascript
function cosineSchedule(step, warmup, total, lrMax, lrMin) {
  if (step < warmup) {
    return lrMax * step / warmup;
  }
  const progress = (step - warmup) / (total - warmup);
  return lrMin + 0.5 * (lrMax - lrMin) * (1 + Math.cos(Math.PI * progress));
}

When the learning rate is too high, each gradient step overshoots the minimum. The loss increases, which produces larger gradients, which produce even larger updates — a positive feedback loop that quickly reaches Infinity, then NaN. This is training divergence. The fix is always to reduce the LR or add warmup. The critical LR depends on the model architecture, batch size, and loss surface curvature.

That's ~0.056% of the total 15T tokens spent just ramping up the learning rate. It seems like a waste, but without warmup the first few hundred steps would produce garbage updates that can permanently damage the model. The investment pays for itself.

Note: The uncorrected m₃ = 0.744 would severely underestimate the true running mean. The correction factor 1/(1 − 0.729) ≈ 3.69 compensates for the zero initialization. After ~30 steps, β₁^t becomes negligible and the correction is effectively 1.

The bias correction for v is even more dramatic: the correction factor is ~333× at step 3. This is because β₂ = 0.999 means v accumulates very slowly, so it takes many more steps to reach its steady state. This large correction is why the first few Adam steps behave differently — and why warmup is important.

This is why training is so memory-hungry. The model weights in FP16 are only 14 GB, but the optimizer states add 56 GB — 4× more than the model itself! This is the main reason ZeRO-style optimizer sharding exists: split these 56 GB across GPUs instead of duplicating on each one.

Line 4 (and line 5) is the bug. The bias correction is missing: mHat = m should be mHat = m / (1 - b1**t), and similarly for vHat. Without bias correction, the first ~100 steps have severely underscaled updates because m and v are initialized to 0 and slowly warm up. After many steps, β^t → 0 and the correction becomes negligible, which is why it "converges to correct behavior later."

javascript
function adamStep(theta, grad, m, v, t, lr, b1, b2, eps) {
  m = b1 * m + (1 - b1) * grad;
  v = b2 * v + (1 - b2) * grad * grad;
  const mHat = m / (1 - Math.pow(b1, t));
  const vHat = v / (1 - Math.pow(b2, t));
  theta = theta - lr * mHat / (Math.sqrt(vHat) + eps);
  return { theta, m, v };
}

With batch=1, μ_B = x, so (x − μ_B) = 0 for every element. The variance is also 0. After normalization, x̂ = 0/√ε ≈ 0 everywhere — the layer outputs γ·0 + β = β regardless of input. This is why LayerNorm was invented for sequence models: it normalizes across features (the d dimension) instead of across the batch, so it works with any batch size, even 1.

Note: PyTorch uses the convention where momentum=0.1 means "keep 10% of old, take 90% of new" — the opposite of many textbooks. This is why the running mean converges quickly to recent batch means. At eval time, this accumulated μ_run replaces the batch mean to ensure deterministic outputs.

BatchNorm computes statistics across the batch for each feature. This creates three problems for transformers: (1) different sequence lengths in a batch have different valid positions, (2) autoregressive generation uses batch=1, and (3) batch statistics couple samples in unwanted ways. LayerNorm sidesteps all of this by normalizing across the d-dimensional feature vector for each token independently.

RMSNorm saves the mean computation (one fewer reduction kernel) and has fewer parameters (no β — only γ). The 2023 LLaMA paper showed it trains just as well as LayerNorm while being ~8% faster in wall-clock time.

Each GPU processes a micro-batch of 4 samples, accumulates gradients over 8 steps (32 effective per GPU), and 4 GPUs contribute in parallel. The optimizer step happens once per 128 samples — identical to training with batch size 128 on a single GPU (if it fit in memory).

For a 32-layer model, that's 64 GB of activation memory saved. This is the main practical benefit of mixed precision — not the 2× faster matmuls on tensor cores (though that helps too), but the halved activation memory that lets you train larger models or use larger batch sizes.

Multiplying the loss by S scales all gradients by S (by the chain rule). The original 1×10⁻⁹ becomes 6.55×10⁻⁵, which FP16 represents with no problem. After the backward pass, we divide all gradients by S before the optimizer step, recovering the true gradient value in FP32.

That's 14 bytes per parameter, or 2× + 4× + 4× + 4× = 14× the FP16 model size. This is the famous "14 bytes per parameter" rule for mixed-precision Adam training. A single 80GB A100 can't even hold the optimizer states for a 7B model — you need at least 2 GPUs with ZeRO Stage 2.

~1M tokens per step is a common target for LLM training. LLaMA 3 used ~4M tokens per step. The tokens-per-step determines how many optimizer steps you need: to train on 1T tokens at 1M tokens/step, you need 1,000,000 steps.

Gradient clipping preserves the direction of the gradient vector but caps its magnitude. Every parameter's gradient is scaled by the same factor, so the relative magnitudes are preserved. This is much better than per-parameter clipping, which distorts the direction.

javascript
function clipGradNorm(grads, maxNorm) {
  const norm = Math.sqrt(grads.reduce((s, g) => s + g * g, 0));
  if (norm <= maxNorm) return grads.slice();
  const scale = maxNorm / norm;
  return grads.map(g => g * scale);
}

A flat loss with tiny gradient norm means the optimizer has no useful signal to follow. This is different from convergence (where loss would be near the theoretical minimum). Common causes: (1) the LR schedule decayed too fast and hit ~0 too early, (2) the LR is so small that updates are negligible, (3) the model landed in a saddle point. The fix is usually to restart with a higher minimum LR or longer warmup.

Clipping activates whenever the norm exceeds max_norm = 1.0. The very first logged norm is 1.2 > 1.0, so clipping was already active at step 100 (and likely from the start). The escalating norms (1.2 → 89.7 → NaN) suggest the model is diverging despite clipping — the gradients are growing faster than clipping can contain. This needs a lower learning rate, not just more aggressive clipping.

A 125M model on 10B tokens is quite small — under 2 hours on 8 A100s. For comparison, LLaMA 3 70B on 15T tokens takes ~30 million GPU-hours.

Easily fits on a single GPU! With 78 GB remaining for activations, you can use very large batch sizes or long sequences. In fact, the bottleneck for a 125M model is usually compute (GPU utilization), not memory. You'd use 8 GPUs for speed, not memory.

Smaller models can tolerate higher learning rates because their loss landscape is smoother with fewer parameters. GPT-2 Small (124M) was originally trained with lr = 2.5×10⁻⁴, but modern practice with cosine schedules suggests ~6-10×10⁻⁴ works well for this scale.

~19K steps is a short training run. With 2000 warmup steps, that means warmup is ~10% of total training — typical for this scale. The cosine schedule decays over the remaining 17K steps.

The correct order: FP16 forward → Scale loss → FP16 backward → Unscale grads → Clip gradients → FP32 Adam step. Scaling must happen before backward (so gradients stay in FP16 range). Unscaling must happen before clipping (clip on true gradient magnitude). The Adam step uses FP32 master weights to maintain precision.

A 125M model on 10B tokens costs ~$35 on spot instances. Scaling this up: a 7B model on 1T tokens costs ~$200K, and LLaMA 3 405B on 15T tokens cost an estimated $100M+. Compute cost scales as O(N × D), so both model size and dataset size are cost multipliers.