Skip to content
ginsberg.si

Transformers from the Weights Up/Part 6 of 20

Transformers from the Weights Up, Part 6: Training: cross-entropy, backprop, AdamW

Training is forward, cross-entropy on the next token, one backward sweep, and an AdamW step. That is where the 6N FLOPs and the 16 bytes per parameter come from.

4 Nov 2026

  1. 1

    Forward

    Tokens [B, T] → logits [B, T, 128256]. Same computation as inference, except every position is scored, not just the last, and intermediate activations are kept for the backward pass.

  2. 2

    Loss

    Targets are the inputs shifted left by one: position t is graded on token t+1 (image 8). Average −log p(target) over all B·T positions.

  3. 3

    Backward

    Reverse-mode autodiff from the scalar loss: one sweep yields ∂L/∂θ for all 8,030,261,248 parameters, at about twice the forward cost.

  4. 4

    Update

    AdamW turns each gradient into a step using its running mean m and mean square v, then applies decoupled weight decay. Zero the gradients and repeat.

IMAGE 08thecatsatonthethe modelLOGITSOFTMAXmat3.174%floor1.922%idea−0.42%moon−1.21%highest logit becomes the tallest bar
Image 8. "the cat sat on the" goes to logits, then softmax: mat 74%, floor 22%, idea 2%, moon 1%.

Cross-entropy on next-token targets

For one position with logits z ∈ ℝ^V and target id y, the loss is the negative log of the probability the model gave the right token. Expanding the softmax gives a form that is also how it is computed stably (log-sum-exp with the max subtracted):

That gradient is the whole story at the top of the network: the predicted distribution minus the one-hot target. Every wrong token is pushed down in proportion to the probability it stole, the right token is pushed up by 1 − pᵧ, and the entries always sum to zero, because softmax can only move probability around. Over Llama-3-8B's vocabulary it is a [128256] vector per position, or [B, T, 128256] per batch.

(z) Cross-entropy and its gradient p − y, on a six-token vocabulary
tokenlogit zp = softmax(z) (bar) vs target y (outline)∂ℓ/∂z = p − ymat ← target0.0016.7%-0.833floor0.0016.7%+0.167rug0.0016.7%+0.167sofa0.0016.7%+0.167moon0.0016.7%+0.167idea0.0016.7%+0.167Σ(p − y) = -0.000 · ln 6 = 1.792 (uniform loss here)
Step
0
Loss ℓ = −ln p(mat)
1.7918 nats
Perplexity e^ℓ
6.000

Toy vocabulary and starting logits made up for illustration; every number is computed live. The target is 'mat'. Each Step is one gradient-descent update applied directly to the logits, z ← z − η(p − y), so you can watch the loss fall. In a real model the same p − y vector, 128,256 entries wide in Llama-3-8B, is what flows back into lm_head and from there into every layer.

From the logits the error enters the weights. With final hidden states H [T, 4096] and G = P − Y [T, 128256], lm_head gets dW_U = Gᵀ H ([128256, 4096], the same shape as the tensor) and the residual stream gets dH = G W_U ([T, 4096]), which is then carried down through all 32 blocks. The same mechanism trains Lesson 5's memories: for FFN slot i, the gradient on its value column is dW_down[:, i] = Σₜ aᵢ(xₜ) · δₜ, where δₜ is the error arriving at the FFN output. Only slots that fired are edited, in proportion to how hard they fired. Gradient descent writes key–value memories a tiny rank-T update at a time; ROME just does one on purpose.

Backprop is reverse-mode autodiff

The forward pass is a directed acyclic graph of primitive ops (matmul, add, SiLU, RMSNorm, softmax). The chain rule says the derivative of the scalar loss with respect to any node is the sum, over the node's consumers, of the consumer's gradient times the local derivative. Reverse mode evaluates that from the loss backwards, so one sweep gives the gradient for every input. Forward-mode differentiation would need one pass per parameter: 8 billion of them.

(aa) Reverse-mode autodiff on a five-op graph: forward, then backward
x1.5000∇ ?w₁0.8000∇ ?b-0.4000∇ ?a = w₁·x·∇ ?z = a + b·∇ ?h = SiLU(z)·∇ ?w₂1.5000∇ ?s = w₂·h·∇ ?L = −ln σ(s)·∇ ?

A one-neuron network (values made up; all arithmetic live): input x, weights w₁, w₂, bias b, SiLU, and a binary cross-entropy loss. Forward (blue edges) computes and caches values left to right. Backward (accent edges) starts from dL/dL = 1 and applies one local rule per node right to left: multiply swaps inputs, add copies, nonlinearities multiply by their slope. Gradients are shown in the lower half of each node.

Backprop through a matmul: where 6N comes from

Apply the multiply rule to a whole matrix. For Y = XW with X [T, K], W [K, N] and upstream gradient dY [T, N], matching shapes forces the answer:

Each is a matmul with the same three dimensions as the forward product, so each costs the same 2TKN FLOPs. Forward 2, backward 4: Lesson 4's 6N FLOPs per token, now derived rather than asserted. dW sums over all T tokens in the batch, which is why bigger batches give a less noisy gradient for the same weight read. For Llama-3-8B, 6 × 8.03B ≈ 4.8 × 10¹⁰ FLOPs per training token, ignoring attention's context-length term.

(ab) One linear layer, three matmuls: where 6N comes from
Forward: Y = X WX[4,096, 4,096]×W[4,096, 14,336]=Y[4,096, 14,336]Y: 2TKNdX: 2TKNdW: 2TKNFLOPs so far: 2 · T·K·N
This product
2 × 4,096 × 4,096 × 14,336 = 0.481 TFLOP
Forward + backward for this layer
1.443 TFLOP = 6 × 58.72M weights × 4,096 tokens

Llama-3-8B's up_proj, written as Y = XW with W stored as [K, N] = [4096, 14336] (PyTorch stores the transpose, [out, in]; the algebra is identical). T is tokens in the micro-batch. Each of the three products is 2·T·K·N FLOPs; the shaded sweep is the output being filled. Forward 2, backward 2 + 2: six FLOPs per weight per token.

A useful sanity check on the rule: the Llama 3 paper reports its 405B flagship used 3.8 × 10²⁵ training FLOPs on 15.6T tokens. 6 × 405 × 10⁹ × 15.6 × 10¹² = 3.79 × 10²⁵. The same rule gives the 8B, on its "about 15T" tokens, ≈ 7.2 × 10²³ (Lesson 4). The paper also reports 38–43% BF16 model-FLOPs utilization for the 405B on H100s; at 40% of the H100 SXM's 989 TFLOP/s, 7.2 × 10²³ FLOPs is roughly 5 × 10⁵ GPU-hours. That last figure is our estimate, not a reported number.

AdamW: the update rule and its state

Plain SGD steps θ ← θ − ηg. Adam (Kingma and Ba, 2014) keeps two exponential moving averages per parameter, a mean m and an uncentered variance v, and divides one by the square root of the other. AdamW (Loshchilov and Hutter, 2017) changes one line: weight decay is applied directly to θ instead of being added to g.

PieceWhat it doesTypical value
m (first moment)Momentum: averages gradient direction over ~1/(1 − β₁) steps, cancelling noise and oscillationβ₁ = 0.9 (≈ 10 steps)
v (second moment)Per-parameter RMS gradient. Dividing by √v makes the step size roughly η for every weight, whatever its gradient scaleβ₂ = 0.95 is common for LLMs (GPT-3 used it); the Adam paper's default is 0.999
Bias correctionm and v start at 0 and would be too small early; dividing by 1 − βᵗ fixes it. At t = 1, m̂ = g and v̂ = g², so the first step is η·sign(g)automatic
εGuards the division when v ≈ 01e-8 (Adam paper default)
Decoupled decay λShrinks every weight by ηλθ per step, uniformly. Added to g instead, it would be divided by √v̂ and weights with large gradients would barely be regularizedLlama 3's scaling-law runs: decay = 0.1 × the learning rate at each step
(ac) SGD vs momentum vs Adam(W) on an ill-conditioned bowl
x = 0minimum (4, 1)loss vs step (log scale)
SGD (η = 0.45)
step 0: (-7.00, -2.20), loss 2.17e+1
SGD + momentum 0.9
step 0: (-7.00, -2.20), loss 2.17e+1
Adam (η = 0.25)
step 0: (-7.00, -2.20), loss 2.17e+1

Toy loss f = ½(0.02(x − 4)² + 4(y − 1)²), curvature 200× larger in y than in x, like the very different gradient scales across a transformer's tensors. Same start, 160 steps, all trajectories computed live. SGD's step is capped by the steep direction (η < 2/4) so it crawls along the shallow one. Adam divides each coordinate by its own RMS gradient, so both coordinates move about η per step. AdamW's decoupled decay (slider) also pulls every weight toward 0, here visibly shifting where x settles.

Training memory: 16 bytes per parameter

Serving Llama-3-8B in bf16 needs 2 bytes per parameter: 16.06 GB. Training it with mixed-precision AdamW keeps five copies' worth of state for every weight, and none of it can be streamed away because every step touches all of it.

TensorPrecisionBytes/paramLlama-3-8B
Weights used in forward/backwardbf16216.06 GB
Gradientsbf16216.06 GB
Master copy of weights (what AdamW updates)fp32432.12 GB
Adam mfp32432.12 GB
Adam vfp32432.12 GB
Total model states16128.48 GB

Why an fp32 master copy? Adam moves each weight by roughly η per step. Late in a cosine schedule η is 10⁻⁵ or smaller, while a bf16 weight near 0.01 has only 8 significant bits (Lesson 1's bf16 panel), so neighbouring representable values are about 6 × 10⁻⁵ apart: the update would round to zero. The update is applied in fp32 and the bf16 copy is re-cast from it. Llama 3 goes further and accumulates gradients across micro-batches and reduces them across data-parallel workers in FP32, for numerical stability.

(ad) Bytes per parameter: serving vs training Llama-3-8B
per GPU2 B2 B4 B4 B4 B80 GB (one H100)bf16 weights: 2 B → 16.06 GBbf16 grads: 2 B → 16.06 GBfp32 master weights: 4 B → 32.12 GBfp32 Adam m: 4 B → 32.12 GBfp32 Adam v: 4 B → 32.12 GB
Training, mixed-precision AdamW
16 bytes/param × 8.03B = 128.5 GB
Per GPU
128.5 GB: does not fit
vs serving in bf16
8× the 16.06 GB inference footprint

Model-state memory only, for 8,030,261,248 parameters (GB = 10⁹ bytes). Mixed-precision Adam is 2 + 2 + 12 = 16 bytes/param, the accounting in the ZeRO paper (Rajbhandari et al., 2019); pure fp32 is also 16, so mixed precision saves compute and activation memory, not optimizer state. Full sharding (ZeRO-3 / FSDP) divides model states by the number of data-parallel GPUs. Activations come on top and scale with batch × sequence × layers.

Lesson 6: key numbers

11.76
nats of loss at uniform initln 128,256; perplexity 128,256
p − y
∂ℓ/∂logits[128256] per position, sums to 0
6N
training FLOPs per token2N forward + 4N backward; ≈ 4.8e10 for 8B
≈ 8N
with activation checkpointingone recomputed forward
16 B
per param, mixed-precision AdamW2 + 2 + 4 + 4 + 4 (ZeRO's 2Ψ + 2Ψ + 12Ψ)
128.5 GB
Llama-3-8B model statesvs 16.06 GB to serve in bf16
3 × 10⁻⁴
Llama 3 8B peak learning rateLlama 3 paper, Table 3
0.9 / 0.95
common β₁ / β₂ for LLMsAdam paper default β₂ = 0.999

Lesson 6: self-check

1. A freshly initialized Llama-3-8B gives near-uniform next-token probabilities. What training loss should the first steps report?

Answer. About ln 128,256 ≈ 11.76 nats per token (perplexity ≈ 128,256). If the first logged loss is far from that, suspect initialization, label shifting, or masking.

2. At one position the target token has p = 0.7 and another token has p = 0.2. What are the gradients on those two logits?

Answer. p − y: the target gets 0.7 − 1 = −0.3 (descent raises its logit), the other gets +0.2 (descent lowers it). All V entries of p − y sum to 0.

3. Why is the backward pass of a linear layer about twice the forward cost?

Answer. It needs two matmuls of the same size as the forward one: dX = dY Wᵀ (to keep backpropagating) and dW = Xᵀ dY (to update the weight). Each is 2TKN FLOPs, so forward 2 + backward 4 = 6 FLOPs per weight per token.

4. Two weights have gradients of 1e-6 and 10 on the very first step. How far does Adam move each?

Answer. About η for both. With bias correction at t = 1, m̂ = g and v̂ = g², so the step is η·g/|g| = η·sign(g) (ε aside). Adam's step size is set by η, not by the gradient's scale; that is the point of dividing by √v̂.

5. Why does AdamW apply weight decay to θ directly instead of adding λθ to the gradient?

Answer. Added to the gradient, the decay term would be divided by √v̂ like everything else, so weights with historically large gradients would be regularized much less than others. Decoupled decay shrinks every weight at the same rate ηλ, which is what weight decay is supposed to mean (Loshchilov and Hutter).

6. You want to fully fine-tune Llama-3-8B with mixed-precision AdamW on 80 GB GPUs. What is the floor on GPU count from model states alone?

Answer. 16 B × 8.03B ≈ 128.5 GB, so at least 2 GPUs with full sharding (≈ 64.2 GB each), leaving under 16 GB each for activations, buffers and fragmentation. In practice you would use more GPUs, activation checkpointing, or a method that trains far fewer parameters.

7. Gradients for an FFN value column vᵢ (Lesson 5) are Σₜ aᵢ(xₜ)·δₜ. What does that say about which memories a batch edits?

Answer. Only slots whose coefficient aᵢ was non-negligible on some token in the batch get a meaningful update, scaled by how strongly they fired and pointed along the error δ arriving at the FFN output. Rarely-firing slots are rarely edited, which is one reason rare facts need many exposures to stick.

Lesson 6: resources