Why architecture matters here
Cross-entropy is not an arbitrary choice of loss; it is maximum-likelihood estimation. Minimizing −log p(target) over the corpus is exactly maximizing the probability the model assigns to the training text, and the loss value has a physical reading: it is the average number of nats the model needs to encode each token, with perplexity = e^loss as its intuitive twin. A loss of 2.0 means perplexity ~7.4 — the model is, on average, as uncertain as a fair choice among 7-8 tokens. Every scaling law, every training curve, every 'model X beats model Y' claim in pretraining is a statement about this one number, so understanding precisely what it computes is not optional.
The gradient structure is why training works at all. Because log-softmax's derivative telescopes, the gradient of the loss with respect to the logits is p − y: each wrong token's logit is pushed down in proportion to the probability the model gave it, and the correct token's logit is pushed up by 1 − p(correct). The signal is bounded in [−1, 1] per component, self-scaling (confident-and-right yields tiny gradients, confident-and-wrong yields large ones), and dense across the whole vocabulary — every logit receives signal at every position. No other common loss couples this stability with this density, and it is the reason logits + CE displaced every alternative output formulation.
The systems stakes are equally concrete. The logits tensor dwarfs every activation in the network — for many models the LM-head matmul plus CE is 20-30% of step memory if done naively — so the difference between a fused, chunked CE and a naive one is often the difference between fitting a batch and OOM. And because the reduction spans the vocabulary axis, tensor parallelism must handle it specially: get the vocab-parallel reduction wrong and the model trains on silently incorrect gradients.
The architecture: every piece explained
Top row: the forward path. The final hidden state h (shape [tokens, d_model]) meets the LM head W (d_model × vocab, often weight-tied to the embedding) to produce logits z. The log-sum-exp stage subtracts the row max m before exponentiating: log Σ e^z = m + log Σ e^(z−m). This is not a nicety — bf16 logits reach ±40 routinely, and e^40 overflows fp32; the max-subtraction makes the largest exponent exactly e^0 = 1. Log-softmax then is log p_i = z_i − m − log Σ e^(z−m), computed directly in log space rather than as log(softmax) — one fewer rounding boundary, no log-of-tiny-number underflow.
Middle row: what the loss actually averages. The NLL gather picks −log p[target] per position — one lookup, not a matmul against a one-hot. Masking zeroes positions labeled ignore_index (padding, prompt tokens in SFT, cross-document boundaries in packed batches) and — critically — divides by the count of unmasked tokens, not batch size, so examples with different mask ratios contribute comparable gradients. Label smoothing replaces the one-hot y with (1−ε) on the target and ε spread over the vocabulary; the gradient becomes p − y_smooth, which stops logits from growing unboundedly on easy tokens and calibrates confidence, at the cost of a constant floor added to the loss value. z-loss (the PaLM trick) adds λ·(log Z)², penalizing drift of the normalizer itself — cheap insurance against logit blow-up in bf16 at scale.
Bottom row: the memory architecture. The fused backward exploits ∂z = p − y: a fused kernel computes loss and this gradient in one pass over the logits, so the full log-probability tensor never needs to be stored for backward. Chunked CE goes further — it slices the token dimension, computes the LM-head matmul, loss, and gradient chunk by chunk, and accumulates, keeping peak memory at one chunk of logits instead of all of them. Vocab-parallel CE handles tensor parallelism, where each rank holds a vocabulary shard of W: each rank computes its local max and local Σe^(z−m); an all-reduce(max) then all-reduce(sum) reconstructs the exact global normalizer, and the rank owning the target token contributes the gather. Two small collectives replace ever materializing the full-vocab logits on any device.
End-to-end flow
Trace one training step of a 8B model, tensor-parallel over 4 GPUs, sequence length 8,192, vocab 128k. The final transformer layer emits h on every rank (replicated after the last all-reduce). Each rank multiplies h by its 32k-column shard of the LM head — four partial logits tensors, none of which is the full [tokens, 128k] monster.
The vocab-parallel CE kernel runs: each rank computes row-wise local maxima over its shard; one all-reduce(max) yields the true row max m. Each rank computes Σ e^(z−m) over its columns; one all-reduce(sum) yields the global normalizer Z. The rank whose shard contains each position's target token looks up z_target and contributes −(z_target − m − log Z); positions whose label is ignore_index (the packed batch has document boundaries and prompt tokens masked) contribute zero and are excluded from the denominator. The scalar loss is averaged over the ~52k unmasked tokens, not the 65k total slots.
Backward starts from that scalar. No stored softmax tensor exists, so the fused kernel recomputes p per chunk and emits ∂z = p − y_smooth (label smoothing ε = 0 here; pretraining usually skips it, SFT often uses 0.1) directly into the gradient buffer for the LM-head matmul's backward. Each rank's ∂z covers only its vocab shard — conveniently, exactly what its shard of W needs. The z-loss term adds 2λ·log Z · ∂log Z into the same buffer for pennies. Gradients flow on into the transformer stack.
On the logging side, the step's loss lands on the curve the team actually watches: smoothed token-averaged CE, reported alongside perplexity = e^loss and gradient norm. Three steps later the loss spikes 0.3 nats; because the pipeline logs the max-logit statistic too, the on-call sees max|z| jumped to 60 on a pathological document — the z-loss coefficient absorbs it, the curve recovers, and nobody restarts from checkpoint. That diagnosis is only possible because the loss was built as an observable system, not a black-box scalar.