Chapter 5 Stability Tricks in Large-Scale Training

5.1 Softmax Instabilities

5.1.1 Output-Side Z-Loss

When a training run at 50B tokens shows periodic loss spikes that slowly recover before eventually cas-cading into full divergence, the post-mortem almost always circles back to the output softmax normalizer. This topic surfaces at the senior and staff level whenever the interview pivots from architecture to training stability: “How do you stabilize a pre-training run at scale?” Getting the answer right requires understanding not just that the z-loss exists but precisely what it is regularizing and why the problem compounds as model scale grows. It is one of those interventions where the mechanism is simple, the cost is negligible, and not knowing about it is an immediate signal that you have not shipped a large training run.

The output distribution of a language model is a softmax over the full vocabulary $V$ :

where $u_{i} ( x )$ are the pre-softmax logits. The quantity $Z ( x )$ is the partition function. The cross-entropy loss is:

The numerical problem is $Z ( x )$. As the model trains and learns to make confident predictions, the logits $u_{j}$ grow in magnitude. If the mean logit magnitude reaches 6 with a vocabulary of 128k tokens, $Z \approx 128000 \cdot e^{6} \approx 5 \times 10^{7}$, giving $\log Z \approx 17 . 7$. In bfloat16 - the standard training dtype for frontier runs - the mantissa has only 7 bits, providing roughly 3 decimal digits of precision. When $Z$ is $5 \times 10^{7}$ and a token’s probability is $e^{u_{j}} / Z$ for a small $u_{j}$, the result underflows to zero. The optimizer receives a zero gradient for most of the vocabulary and loses the signal it needs to refine predictions across rare tokens.

The z-loss adds an auxiliary term to the objective that penalizes large values of $\log Z ( x )$:

where $\alpha$ is a small scalar. PaLM introduced this in production with $\alpha = 10^{- 4}$, explicitly as a training stability mechanism. Minimizing this term drives $\log Z \rightarrow 0$, which implies $Z \rightarrow 1$, which means the model is implicitly encouraged to keep its logits centered near zero.

The gradient of the z-loss with respect to logit $u_{j}$ is:

This correction is nonzero only when $\log Z \neq 0$ and is proportional to the token’s own probability - it applies the largest correction to the highest-probability token, which is exactly where the logit growth originates.

The z-loss does not change what the model learns to predict; it changes where the model places mass in logit space, keeping the partition function numerically tractable throughout training.

When $Z ( x ) \approx 1$, the softmax simplifies to $p_{i} \approx e^{u_{i}}$. There is no large-denominator cancellation, no log-sum-exp subtracting two enormous similar numbers. The Jacobian $\partial p_{i} / \partial u_{j} = p_{i} ( \delta_{i j} - p_{j} )$ remains wellconditioned because no probability mass has collapsed to zero. The model can still express confident predictions - it does so with moderate logits and a normalizer near one, rather than with huge logits and an offsetting huge denominator.

The alternative - relying on gradient clipping alone - fails for runs longer than a few hundred billion tokens. Gradient clipping controls the global gradient norm but does not prevent $Z$ from drifting. A model can accumulate a slowly increasing $\log Z$ that stays below the clipping threshold at each individual step but compounds over tens of thousands of steps until a single bad batch tips it into a spike. The z-loss prevents this drift by applying a persistent, differentiable pressure that grows proportionally with the drift itself.

The compute cost of the z-loss is essentially zero. $\log Z ( x )$ is already computed as part of the softmax during the forward pass. The auxiliary loss requires only a squaring, a scaling by $\alpha$, and the gradient correction during the backward pass. Across 500k training steps on 4096 H100s, this contributes nothing measurable to wall-clock time.

A Worked Example

Consider a 70B-parameter model: vocabulary 128k, trained in bfloat16, global batch 4M tokens per step, $\alpha = 10^{- 4}$ for the z-loss.

Without z-loss, at step 200k: Suppose the mean logit magnitude has drifted to $\bar{u} \approx 6 . 0$. Rough estimate: $Z \approx 128000 \cdot e^{6} \approx 5 \times 10^{7}$, so $\log Z \approx 17 . 7$. A token with logit $u_{j} = 0$ contributes probability $e^{0} / 5 \times 10^{7} = 2 \times 10^{- 8}$. In bfloat16, this underflows. The optimizer sees a zero gradient for all such tokens. Training continues, but the logit distribution is degenerating.

With z-loss, at step 200k: The auxiliary gradient $2 \alpha \log Z = 2 \times 10^{- 4} \times 17 . 7 \approx 3 . 5 \times 10^{- 3}$ is added per token at every step. This is small per step but has been applied continuously since step one. The cumulative effect has suppressed $\bar{u}$ to $\approx 2 . 0$, giving $Z \approx 128000 \cdot e^{2} \approx 9 . 4 \times 10^{5}$ and $\log Z \approx 13 . 8$. The probability of a token at logit zero is now $e^{0} / 9 . 4 \times 10^{5} \approx 10^{- 6}$, still representable in bfloat16. No loss spike occurred at step 200k where one would otherwise have been likely.

The $\log Z$ metric tells the story cleanly: a run without z-loss shows $\log Z$ trending from $\approx 5$ at step 1k to $\approx 18$ at step 200k. A run with z-loss shows $\log Z$ stabilizing below 10 after the first 10k steps and remaining bounded throughout.

What You Actually Decide in Practice

• Set $\alpha = 10^{- 4}$ as your starting point. This is PaLM’s value, independently validated by OLMo 2 and DCLM. Values above $10^{- 3}$ visibly hurt training loss by over-constraining the logit scale. Values below $10^{- 5}$ have negligible effect on $\log Z$ drift.

• Apply z-loss from step zero, not as a retrofit. Once $\log Z$ has drifted to 15-20, retroactively adding z-loss causes a transient spike as the logit distribution is pulled back. The intervention must be proactive.

• Log $\log Z ( x )$ per training step as a health signal. This scalar is free to compute - it is already computed during the forward pass. A $\log Z$ that is monotonically increasing with no plateau is a leading indicator of a future spike, visible days before the spike appears in the loss curve.

• Do not confuse z-loss with label smoothing. Label smoothing changes the target distribution, training the model to be less confident. Z-loss does not change what the model is trained to predict - it only constrains the logit scale. They solve different problems and can be used together.

• Treat z-loss as mandatory above 30B parameters. At 7B, $\log Z$ drift is slow enough that training often completes before a spike occurs. At 70B trained for 2T tokens, the cumulative drift is substantial. At 100B and beyond, in my experience, running without z-loss is gambling on the absence of a known failure mode.

How This Shows Up in Interviews

• “Your 70B pre-training run is showing periodic loss spikes at 50B and 150B tokens. What do you investigate?” Answer skeleton: first distinguish a data quality spike (a bad batch artifact, which recovers in 1-2k steps) from a structural drift issue (which recovers slowly and recurrs). Check gradient norms for correlated spikes. Then inspect $\log Z ( x )$ - if it has been trending upward for tens of thousands of steps before each spike, you have logit scale drift. The fix is to add z-loss from the next run’s start.

• “What is z-loss and when would you add it?” Answer skeleton: an auxiliary term $\alpha ( \log Z ( x ) )^{2}$ added to the cross-entropy loss; forces the output softmax normalizer toward one; prevents logit scale drift that causes numerical instability in bfloat16; use $\alpha = 10^{- 4}$; apply from step one; costs nothing in compute; does not change target predictions, only logit scale.

• Senior follow-up: “Why does $Z ( x )$ drifting cause specific problems in bfloat16 that would not appear in fp32?” A senior answer explains the mantissa bit difference directly: fp32 has 23 mantissa bits ($\approx$7 decimal digits of precision); bfloat16 has 7 mantissa bits ($\approx$3 decimal digits). When $Z ( x ) = 10^{7}$ and a token’s probability is $10^{- 8}$, fp32 represents this as a denormal but nonzero; bfloat16 rounds it to exact zero. The log-sum-exp trick prevents overflow in computing $\log Z$ itself, but it does not prevent underflow in the resulting probabilities $p_{j} = e^{u_{j}} / Z$ when $u_{j} \ll u_{\max}$. The model then receives zero gradients for those tokens and loses the signal needed to refine predictions over the low-probability vocabulary.

Key Takeaways

• The output softmax normalizer $Z ( x )$ grows naturally during training as the model becomes more confident; in bfloat16, large $Z$ causes token probabilities to underflow to zero and gradients to vanish for most of the vocabulary.

• Z-loss adds $\alpha ( \log Z )^{2}$ to the cross-entropy objective, forcing $\log Z \rightarrow 0$ without changing the model’s target distribution.

• Use $\alpha = 10^{- 4}$, apply from step one, and $\text{log} \log Z ( x )$ every step as an early-warning health metric.

• Pioneered by PaLM; subsequently adopted by Baichuan 2, DCLM, and OLMo 2; increasingly considered mandatory for runs above 30B parameters in bfloat16.

• Mechanistically distinct from label smoothing: z-loss controls logit scale, not target confidence.

5.1.2 Attention-Side Soft-Capping of Logits

When an attention head begins routing every query token to the same one or two key positions - what practitioners call attention collapse - the root cause is almost always unbounded logits feeding the attention softmax. This failure mode surfaces in senior and staff interviews whenever the conversation reaches training stability: “What happens to your attention mechanism late in training, and how do you prevent it from degrading?” Two distinct interventions exist: QK-norm and logit soft-capping. Knowing when to use each, and understanding the concrete trade-off between them, is what separates a strong answer from a rote recitation of technique names.

The attention computation for head $h$ at layer $\ell$ is:

The pre-softmax logits are $l_{i j} = q_{i} \cdot k_{j} / \sqrt{d_{k}}$. As training progresses, the norms of $q$ and $k$ grow - the model has no built-in mechanism to prevent this. With head dimension $d_{k} = 128$ and unit-RMS query and key vectors (achievable after warm-up), the maximum possible logit from a perfectly aligned pair is $\sqrt{d_{k}} \cdot \gamma_{q} \gamma_{k} \approx \sqrt{128} \approx 11 . 3$. But norms are not unit-RMS late in training; logits of 30-80 are observed in unstabilized runs. When one logit dominates a row by a large margin, the softmax concentrates: $\text{softmax} ( [ 60 , 0 , . . . , 0 ] ) \approx ( 1 , 0 , . . . )$. The head reduces to a positional lookup, its expressivity gone. Worse, the softmax Jacobian $\partial p_{i} / \partial l_{i j} = p_{i} ( \delta_{i j} - p_{j} )$ vanishes everywhere except the dominant position, and gradient flow through the attention mechanism effectively stops.

QK-norm applies a learned LayerNorm to the query and key projections before the dot product:

The LayerNorm fixes the per-component RMS of $q$ and $k$ to its learned scale parameter $\gamma$, regardless of how large the raw projections become. At initialization with $\gamma_{q} = \gamma_{k} = 1$, the maximum logit is bounded by $\sqrt{d_{k}}$. Training can increase $\gamma$, but the growth is controlled by gradient descent operating on the loss signal rather than by unchecked norm amplification.

Logit soft-capping takes the complementary approach: it bounds the dot products after they are computed. Given raw logits $l_{i j}$, the soft-capped logit is:

where $s$ is the soft-cap hyperparameter. For $| l_{i j} | \ll s$, $\tanh ( l / s ) \approx l / s$ and the logit passes through nearly unchanged. For $| l_{i j} | \gg s$, the tanh saturates at $\pm 1$ and the logit is clipped to $( - s , s )$.

QK-norm prevents logit growth at the input to the dot product by bounding vector norms; soft-capping prevents it at the output by clipping large values - the choice between them is not cosmetic, as they have measurably different effects on both stability and final quality.

QK-norm is the more principled intervention. By controlling the norms of $q$ and $k$ before the inner product, it addresses the root cause rather than the symptom. The model retains full expressivity over what it attends to, because the learned $\gamma$ parameters adjust the effective temperature of the attention distribution. Critically, QK-norm allows the use of more aggressive learning rates: because the feedback loop causing logit explosion is broken, the optimizer can push harder without triggering instability. NVIDIA’s ablation on training stability confirms this - adding QK-norm improves final perplexity relative to the baseline, specifically because it permits higher peak learning rates.

Soft-capping is simpler to implement and requires no additional learned parameters, but it introduces a nonlinearity with a gradient-killing property: the derivative of $s \cdot \tanh ( l / s )$ with respect to $l$ is $1 - \tanh^{2} ( l / s )$, which approaches zero as $| l | \rightarrow \infty$. This means that logits which have already grown past roughly $2 s$ receive vanishingly small gradients through the capping operation. Soft-capping prevents softmax collapse but does not force the model to represent attention distributions compactly - it merely hides the symptoms from the softmax while the underlying logit norms continue to drift. The same NVIDIA study finds that soft-capping at $s \in [ 30 , 50 ]$ either matches or slightly worsens baseline perplexity, because the tanh distortion changes the relative contrast between large logits, impeding the model’s ability to express sharply peaked attention distributions when they are genuinely needed.

Both QK-norm and soft-capping originated outside pure language model training. QK-norm was introduced for training ViT-22B (a 22-billion parameter vision transformer), where large-scale instability first became a pressing practical problem. It migrated into language through multimodal models - Chameleon, Idefics - before being adopted in pure-text production runs at DCLM and OLMo 2. Soft-capping was developed within the Gemma 2 training effort. This lineage matters: both techniques were motivated by empirical failure modes at scale, not by theoretical arguments.

A Worked Example

Consider a 13B-parameter model, 40 attention heads, $d_{k} = 128$, sequence length 4096, training in bfloat16.

Soft-capping behavior at different logit magnitudes (using Gemma 2’s local-attention cap $s = 50$):

• $l = 5$: $l^{\text{cap}} = 50 \cdot \tanh ( 0 . 10 ) \approx 4 . 99$ - unchanged, cap has no effect.

• $l = 30$: $l^{\text{cap}} = 50 \cdot \tanh ( 0 . 60 ) \approx 26 . 8$ - a 10.7% reduction; the relative ordering of large logits is slightly compressed.

• $l = 80$: $l^{\text{cap}} = 50 \cdot \tanh ( 1 . 60 ) \approx 46 . 1$ - a 42% reduction; a logit that would have completely dominated the softmax is now only moderately larger than the cap value.

• $l = 200$: $l^{\text{cap}} = 50 \cdot \tanh ( 4 . 00 ) \approx 49 . 9$ - effectively clipped to the cap; any logit beyond this is indistinguishable.

The implication: for a head where the maximum raw logit is 80 and the second-largest is 40, their softcapped values are 46.1 and 34.7, giving softmax probabilities roughly proportional to $e^{46 . 1}$ and $e^{34 . 7}$. The ratio $e^{46 . 1} / e^{34 . 7} = e^{11 . 4} \approx 9 \times 10^{4}$ - still extremely peaked, but no longer $e^{80} / e^{40} = e^{40} \approx 2 \times 10^{17}$. The softmax is no longer numerically degenerate.

QK-norm behavior at initialization: With $\gamma_{q} = \gamma_{k} = 1$, the RMS of each normalized query component is 1. The maximum logit is:

This bound holds throughout training, because LayerNorm always outputs unit per-component RMS scaled by $\gamma$. If the model needs sharper attention at a particular head, it will learn $\gamma > 1$ for that head - but the growth is slow, smooth, and controlled by the gradient, not by runaway norm amplification.

What You Actually Decide in Practice

• Choose QK-norm for new production runs. It addresses the root cause of logit explosion rather than the symptom. Implement it as a standard LayerNorm applied to the query and key projections before the dot product, with $\gamma$ initialized to 1.0 and no bias term. The additional parameter count is 2 scalars per head per layer - negligible.

• If using soft-capping, use $s = 30$-$50$. Gemma 2 uses $s = 50$ for local attention layers and $s = 30$ for global attention layers. Below $s = 10$, the tanh nonlinearity distorts attention patterns throughout training, not just for pathological logits, hurting quality from early on.

• Do not stack QK-norm and soft-capping without ablating. They are not additive - you are restricting logit magnitude through two mechanisms simultaneously, and the interaction can over-constrain attention heads during early training when the model needs to form its initial routing patterns. Pick one.

• Monitor per-head 99th-percentile logit magnitude as a training health metric. A logit distribution that is monotonically increasing with no plateau signals incipient collapse, often 5k-20k steps before it manifests as a loss spike. Log this scalar per layer per checkpoint.

• Apply whichever technique you choose from step one. Retrofitting QK-norm mid-run causes a distribution shift in the attention patterns that spikes the loss; the model must re-learn its routing structure under the new normalization regime. Retrofitting soft-capping is less disruptive but still causes a transient perturbation. Neither is a viable mid-run rescue operation.

How This Shows Up in Interviews

• “You’re seeing attention heads collapse to attending a single position after 20B tokens. What’s happening and what do you do?” Answer skeleton: query and key norms have grown unchecked, producing logits large enough for the softmax to concentrate near one-hot; QK-norm prevents this by normalizing $Q$ and $K$ before the dot product, bounding the logit magnitude from step one; alternatively, soft-capping clips post-dot-product logits via a tanh; add QK-norm from the next run’s start; in the current run, monitor per-head logit statistics and consider reducing the learning rate to slow the norm growth.

• “Compare QK-norm and soft-capping. When would you choose each?” Answer skeleton: QK-norm is preferred - it addresses root cause, adds minimal parameters, improves perplexity, and enables higher learning rates; soft-capping is simpler (no learned parameters) and can be retrofitted into existing attention implementations without changing the projection layers, but it can worsen perplexity due to tanh distortion of logit ordering; choose QK-norm for greenfield runs; soft-capping if you need to stabilize an existing architecture with minimal code changes.

• Senior follow-up: “Logit soft-capping’s gradient through the tanh vanishes for large logits. Doesn’t that prevent the model from learning to reduce a logit that’s already past the cap?” A senior answer identifies this precisely: once a logit exceeds $\approx 2 s$, the gradient of the loss with respect to that logit is attenuated by $( 1 - \tanh^{2} ( l / s ) ) \approx 0$, meaning the optimizer receives almost no signal to reduce it. The soft-cap prevents the softmax from seeing the extreme value, but the underlying logit norm can continue to grow unchecked. This is why soft-capping addresses the symptom but not the cause: in a run without any QK-norm, raw logit norms may still drift upward over the training horizon even with soft-capping in place. QK-norm avoids this by preventing norm growth at the source, maintaining well-conditioned gradients throughout.

Key Takeaways

• Unchecked growth of query and key norms produces large attention logits; when a logit dominates a row, the softmax collapses to near-one-hot and gradient flow through the attention mechanism stops.

• QK-norm normalizes $Q$ and $K$ via LayerNorm before the dot product, bounding logit magnitude at $\gamma \sqrt{d_{k}}$; it improves perplexity and allows higher learning rates.

• Soft-capping clips post-dot-product logits to $( - s , s )$ via $s \cdot \tanh ( l / s )$; it is simpler to implement but can worsen quality because tanh distorts the relative contrast between large logits and kills gradients for already-large values.

• Both techniques originate from large-scale vision and multimodal training (QK-norm from ViT-22B, softcapping from Gemma 2) and migrated into pure-text LLM training as model scale demanded greater stability.

• Apply from step one; do not stack both without ablating; monitor per-head 99th-percentile logit magnitude as an early-warning health metric.

5.2 Numerical Precision and Training Health

5.2.1 BF16, FP16, and Mixed-Precision Pitfalls

Every pre-training run on modern hardware mixes at least two floating-point formats, and the assignment of tensors to formats is not arbitrary. Interviewers at infrastructure-focused roles ask about mixed precision because getting it wrong produces either silent numerical failures that look like architecture bugs or unnecessary memory bloat that limits achievable model scale. The question surfaces at the mid-level as “walk me through your mixedprecision training setup,” and at the senior level as “why do you sometimes see NaN values propagate even in a bfloat16 run that has gradient clipping enabled.”

Three floating-point formats dominate LLM training today. Each trades dynamic range against mantissa precision:

• FP32 (float32): 1 sign + 8 exponent + 23 mantissa = 4 bytes. Dynamic range $\approx 1 . 2 \times 10^{- 38}$ to $3 . 4 \times 10^{38}$. Roughly 7 decimal digits of precision. The traditional default; now reserved for optimizer state and master weights.

• FP16 (float16): 1 sign + 5 exponent + 10 mantissa = 2 bytes. Maximum representable value: $65504$. Minimum positive normal: $\approx 6 \times 10^{- 5}$. Roughly 3 decimal digits of precision. The narrow 5-bit exponent is the critical limitation for large-model training.

• BF16 (bfloat16): 1 sign + 8 exponent + 7 mantissa = 2 bytes. Dynamic range identical to FP32 ($\approx 3 . 4 \times 10^{38}$). Roughly 3 decimal digits of precision. Developed by Google Brain in 2018 specifically because deep learning training needs dynamic range more than it needs mantissa resolution.

An 8-bit format, FP8, exists in two variants on H100 hardware: E4M3 (4 exponent bits, 3 mantissa bits, used for weights and forward-pass activations, maximum $\approx 448$) and E5M2 (5 exponent bits, 2 mantissa bits, used for gradients, maximum $\approx 57344$). FP8 requires per-tensor dynamic scale factors and is production-viable on H100s via NVIDIA Transformer Engine, but it is not yet the default at most organizations.

The standard production recipe is BF16 mixed precision:

• Forward and backward compute (all matrix multiplications): BF16. Tensor cores accumulate partial sums internally in FP32 and return a BF16 result.

• Master parameters: FP32.

• Optimizer states (Adam’s first moment $m_{t}$ and second moment $v_{t}$): FP32.

• Gradient accumulation across micro-batches: FP32; per-step gradients passed to the optimizer are FP32.

This asymmetry follows from a key observation: per-layer matmuls are robust to the rounding noise introduced by 7 mantissa bits, but optimizer momentum terms that accumulate small corrections over hundreds of thousands of steps require 23 mantissa bits to avoid signal loss. Under AdamW, the per-parameter update is proportional to $m_{t} / \sqrt{v_{t} + \epsilon}$. When $m_{t} \approx 10^{- 5}$ and $v_{t} \approx 10^{- 4}$, the ratio is approximately $0 . 03$. In BF16, a value of $10^{- 5}$ rounds to zero. In FP32, it is preserved.

Use BF16 where you need throughput; keep FP32 wherever you accumulate over time.

The core argument against FP16 for large-model training is dynamic range. FP16’s maximum representable value is $65504$. During training, gradient magnitudes, activation outputs, and pre-softmax logits can easily exceed this threshold. In naive attention (without FlashAttention’s safe-softmax), a pre-softmax score of $12$ contributes $e^{12} \approx 162000$ to the partition function before the log-sum-exp normalization step. If that intermediate exponentiation happens in FP16, it overflows to inf. Once inf enters the compute graph, any subsequent multiplication produces NaN, which propagates forward and backward through every layer in that batch. The run is dead.

BF16 avoids this entirely. Its 8 exponent bits match FP32’s ceiling. The precision sacrifice - 7 mantissa bits versus 23 - means values are rounded to approximately 0.7% of their true magnitude, which has no measurable effect on the final quality of a 10B-parameter model trained on a trillion tokens. The same 0.7% rounding applied to an optimizer state accumulated over 500k steps would progressively corrupt the effective learning signal, which is why optimizer states stay in FP32.

Loss scaling is the standard workaround for training in FP16: scale the loss by a constant ($\sim 2^{15}$ to $2^{20}$) before the backward pass to lift gradient magnitudes into FP16’s representable range, then unscale before the optimizer step. Dynamic loss scaling adjusts this factor step by step to prevent both overflow and underflow. This works, but it adds a second failure mode: if the scale is too large, gradients still overflow; if too small, they underflow to zero silently. BF16 eliminates the need for loss scaling entirely, removing an entire class of training instability.

A Worked Example

Consider a 13B-parameter model: $d = 5120$, 40 layers, 40 attention heads, $d_{\text{head}} = 128$, trained with AdamW in BF16 mixed precision on 8 H100 80 GB GPUs.

Static memory footprint (per the 16-bytes-per-parameter rule):

• FP32 master parameters: $13 \times 10^{9} \times 4 = 52$ GB

• FP32 gradients: $13 \times 10^{9} \times 4 = 52$ GB

• FP32 Adam $m_{t}$: $13 \times 10^{9} \times 4 = 52$ GB

• FP32 Adam $v_{t}$: $13 \times 10^{9} \times 4 = 52$ GB

• BF16 compute copy of parameters: $13 \times 10^{9} \times 2 = 26$ GB Total static tensors: $234$ GB. This far exceeds a single H100’s 80 GB. With ZeRO stage-3 across 8 GPUs, the static footprint per GPU drops to $\approx 29$ GB, leaving $\sim 50$ GB per GPU for activations. The often-cited “16 bytes per parameter” accounts for the four FP32 copies (params, gradients, $m_{t}$, $v_{t}$) but not the BF16 compute copy; the true figure under BF16 mixed precision is $\approx 18$ bytes per parameter.

FP16 overflow scenario: Suppose we use FP16 for compute without loss scaling. At step 1,500, a long-document batch pushes attention logits to $s_{i j} = 11 . 9$ before the $1 / \sqrt{d_{\text{head}}}$ scaling at one head. With grad_clip set to 1.0 but the overflow occurring in the forward pass, no gradient norm check saves us. The raw exponent $e^{11 . 9} \approx 147000$ exceeds FP16’s ceiling of $65504$; the result is inf. The attention softmax for that row returns NaN. The NaN propagates through all 40 layers in the backward pass. After the optimizer step, the affected parameters contain NaN. At step 1,501, the loss reads NaN. The run is unrecoverable without a rollback.

BF16 with the same batch: BF16’s exponent range extends to $3 . 4 \times 10^{38}$. The value $e^{11 . 9} \approx 147000$ is represented without overflow. FlashAttention’s tiled online-softmax further avoids materializing the full exponentiated row in any single buffer. Training continues uninterrupted.

What You Actually Decide in Practice

• Default to BF16 for all compute; never use FP16 for LLM training without dynamic loss scaling. On H100s, BF16 tensor core throughput is roughly $10 \times$ faster than FP32 and comparable to FP16, with none of FP16’s overflow risk. The choice is effectively free.

• Optimizer states and master parameters must stay in FP32. Switching Adam’s $m_{t}$ and $v_{t}$ to BF16 causes visible training quality degradation above 7B parameters, even in short 100B-token runs. The 3- digit mantissa precision rounds away the small per-step corrections that build up meaningful parameter movement over 500k steps.

• Use FlashAttention for numerically stable attention in BF16. Naive attention that materializes the full $N \times N$ attention matrix in BF16 can produce catastrophic cancellation at sequence lengths above 4k. FlashAttention’s tiled algorithm maintains the running softmax max and sum in FP32 within each CUDA kernel tile, avoiding this without a format switch.

• Consider FP8 for inference-side quantization before adopting it for training. FP8 inference (quantizing a trained BF16 model) is low-risk and well-supported. FP8 training requires Transformer Engine, per-tensor scale factor calibration, and empirical validation against a BF16 baseline for 10k steps before committing. The throughput gain ($\sim 1 . 5$-$2 \times$ over BF16) is real but the engineering overhead is substantial.

• Monitor activation_max and the pre-clip gradient norm from step zero. These two scalars are your early-warning system for format overflow. A activation_max that is trending toward $10^{4}$ in BF16 (still far from overflow, but growing) combined with an elevated gradient norm signals that the model is entering a regime where even BF16’s range may be stressed; in FP8 mode, the same trend would indicate imminent scale-factor miscalibration.

How This Shows Up in Interviews

• “Walk me through your mixed-precision training setup for a 70B model run.” Answer skeleton: BF16 for all forward and backward matmuls; FP32 master weights, gradients, and Adam optimizer states; FlashAttention for numerically stable attention; gradient clipping at 1.0. Explain that FP16 requires dynamic loss scaling and is avoided for this exact reason. Mention the 16-18 bytes per parameter static footprint and why ZeRO stage-2 or -3 is required at 70B.

• “Why does training with FP16 sometimes produce NaN values that BF16 does not?” Answer skeleton: FP16 has a 5-bit exponent capping representable values at $65504$. Attention logits, pre-softmax activation values, or backward-pass gradient products can exceed this threshold, overflowing to inf, which any subsequent multiplication converts to NaN. BF16’s 8-bit exponent matches FP32’s dynamic range and eliminates this overflow path; the loss in mantissa precision (7 bits vs. 23) does not matter for per-layer computation.

• Senior follow-up: “Your team wants to enable FP8 training end-to-end on H100s to improve throughput. What is your rollout plan and what failure modes do you specifically instrument?” A senior answer covers: (1) start with FP8 only in the feed-forward matmuls, leaving attention, layer norm, and softmax in BF16/FP32; (2) instrument per-tensor amax_history to detect scale-factor miscalibration - if the rolling maximum underestimates by $2 \times$, overflow occurs silently; (3) run a parallel BF16 baseline for 10k steps at small scale to confirm loss curves match within 0.1 nats; (4) expand to full FP8 only after the subset is validated; (5) flag any step where nan_frac $> 0$ for immediate investigation - in FP8 mode a single miscalibrated tensor contaminates the entire forward pass.

Key Takeaways

• FP16 is dangerous for LLM training above 7B parameters without dynamic loss scaling; its 5-bit exponent overflows to inf at values above $65504$, which multiplies to NaN and propagates through every downstream tensor in the batch.

• BF16 matches FP32’s dynamic range with half the memory, making it the standard compute dtype for frontier training runs; the mantissa precision loss is inconsequential for per-layer computation.

• Optimizer states and master weights must stay in FP32; the 23-bit mantissa precision is required to accumulate small per-step corrections faithfully over hundreds of thousands of optimizer steps.

• The standard static memory budget under AdamW BF16 mixed precision is approximately 18 bytes per parameter: 4 (FP32 params) + 4 (FP32 gradients) + 4 (Adam $m_{t}$) + 4 (Adam $v_{t}$) + 2 (BF16 compute copy); activations scale separately with batch size and sequence length.

• FP8 training on H100s with Transformer Engine yields roughly $1 . 5$-$2 \times$ throughput gains over BF16 but requires per-tensor scale factor calibration and careful monitoring of amax_history; most teams adopt it only after the BF16 pipeline is stable and throughput-limited.

5.2.2 Loss Spikes: Detection and Mitigation

A loss spike that goes undetected for 3,000 steps forces a checkpoint rollback that costs several GPU-hours and invalidates the data-mix accounting for that training window. Staff-level interviewers expect a concrete playbook: not “watch the loss curve” but a named set of metrics, specific alert thresholds, and tiered responses that distinguish recoverable anomalies from structural damage. The topic surfaces as “walk me through your training health monitoring setup” or directly as “your 70B run spiked at 150B tokens and isn’t recovering - what do you do?”

A loss spike is a sudden increase of 0.1 nats or more in training loss that persists for more than 50 consecutive steps. Three root causes account for nearly all production instances, and they are distinguishable by which diagnostic metric moves first.

1. Data batch anomaly. A batch containing corrupted tokens, binary-encoded text, heavily duplicated content, or extreme-length documents with out-of-distribution statistics produces a locally very large loss gradient. The gradient is large but non-destructive: the model absorbs the update and the loss recovers within 100-500 steps as subsequent normal batches dilute the optimizer state.

2. Gradient explosion or logit drift. The global gradient norm persistently exceeds the clip threshold (or the clip threshold is set too loosely), causing large parameter updates that overshoot the loss basin over multiple steps. Alternatively, the softmax logit scale drifts upward as described in the z-loss subsection, gradually destabilizing the gradient signal until a threshold is crossed. Both sub-causes share the signature of a non-recovering spike with no single obvious culprit batch.

3. Numerical precision failure. FP16 or FP8 overflow introduces inf or NaN values into the compute graph. The loss may jump to extreme values (tens of nats or literal NaN) and does not recover without a rollback and configuration change.

The practical monitoring set is five scalars logged every training step at negligible overhead:

• loss: the primary signal, evaluated on the current batch.

• grad_norm: the global $\ell_{2}$ norm of all gradients before clipping. A clipped norm tells you nothing about whether an explosion occurred; the pre-clip norm does.

• log_Z: the mean log-partition-function across the batch, already computed during the forward pass as part of the cross-entropy loss. The leading indicator for logit drift.

• loss_ema_ratio: current loss divided by an exponential moving average of the previous 1,000 steps. A ratio above $1 . 5$ triggers a Tier 1 alert; above $2 . 5$ triggers Tier 2.

• nan_frac: fraction of gradient elements that are NaN or Inf. Any nonzero value triggers an immediate halt.

Root cause determines the correct response; treating all spikes as identical wastes GPU hours on rollbacks that are unnecessary and misses interventions that are urgent.

The value of the triage protocol is specificity. Three response tiers map onto the three root causes:

Tier 1 - skip and continue. Applied when grad_norm spikes to more than $3 \times$ its running median in a single step, loss_ema_ratio is above $1 . 5$, and the loss begins returning toward the baseline within 200 steps. This pattern is consistent with a data anomaly. The correct action is to log the shard index of the offending batch, optionally skip the optimizer update for that step (retain the gradient for accounting purposes), and continue without rollback. Introducing a full rollback for a self-correcting anomaly destroys training throughput with no benefit.

Tier 2 - rollback to checkpoint. Applied when the loss rises for more than 500 steps without returning below the pre-spike level, or when log_Z has increased by more than 3 units in a 300-step window. This pattern indicates structural damage to the model or optimizer state. Roll back to the checkpoint immediately preceding the onset of the gradient norm elevation, remove the offending data range from the training sequence, and resume. If log_Z was the leading signal, add or increase the z-loss coefficient before resuming.

Tier 3 - rollback and reconfigure. Applied when nan_frac exceeds zero, or the loss jumps to values above 10 nats without gradient norm explanation. This indicates a numerical precision failure. Roll back and identify the source tensor (activation monitoring or a per-layer nan_frac breakdown will localize it), then adjust the precision configuration - switching a problematic layer back from FP8 to BF16, or from BF16 to FP32 - before resuming.

The prior approach of relying on gradient clipping alone addresses only Tier 1 partially. Gradient clipping caps the norm of each update but does not prevent logit scale drift, does not flag data quality issues upstream of the gradient computation, and does not detect NaN propagation that originates in the forward pass.

A Worked Example

A 13B-parameter model: $d = 5120$, 40 layers, BF16 compute, AdamW, global batch 2M tokens per step, base learning rate $3 \times 10^{- 4}$, cosine decay, grad_$\mathtt{clip} = 1 . 0$. The run is at step 60,000 when two anomalies occur.

Scenario 1 - data anomaly at step 42,100:

• loss: $2 . 41 \rightarrow 2 . 89$ (single step, $+ 0 . 48$ nats)

• Pre-clip grad_norm at step 42,100: $7 . 3$ (running median over prior 5,000 steps: $1 . 1$; ratio $= 6 . 6 \times$)

• loss over next 400 steps: $2 . 89 \rightarrow 2 . 72 \rightarrow 2 . 58 \rightarrow 2 . 44$ (recovering)

• log_Z: stable at $8 . 3 \pm 0 . 2$ before and after the spike

• nan_frac: 0 throughout

Diagnosis: data anomaly (Tier 1). The gradient norm is $6 . 6 \times$ the median at a single step; the loss is recovering. Action: log the batch shard index for data quality inspection, skip the optimizer update at step 42,100, continue training. No rollback. Estimated cost: zero additional GPU time; one flagged shard for offline investigation.

Scenario 2 - logit drift at step 58,500:

• log_Z trend: $8 . 1$ at step 57,000, $9 . 6$ at step 57,500, $11 . 9$ at step 58,000, $15 . 2$ at step 58,500 (rising for 1,500 steps)

• loss at steps 58,200 onward: $2 . 38 \rightarrow 2 . 65 \rightarrow 2 . 99 \rightarrow 3 . 41 \rightarrow 3 . 88$ (diverging)

• grad_norm: elevated but not extreme, $1 . 8$-$2 . 6$ (only $1 . 7 \times$ median, not a single-step spike)

• nan_frac: 0

Diagnosis: logit scale drift (Tier 2). The run was not configured with z-loss and the softmax normalizer has grown to $e^{15 . 2} \approx 4 \times 10^{6}$. At this scale, BF16’s 7-bit mantissa rounds most token probabilities to zero, collapsing the gradient signal for the majority of the vocabulary. Action: roll back to the checkpoint at step 57,000 (the last checkpoint before the log_Z acceleration began), add z-loss with $\alpha = 10^{- 4}$, resume. The log_Z stabilizes below 10 within 3,000 steps; the loss returns to the pre-spike trajectory by step 5,000 after resumption.

What You Actually Decide in Practice

• Save checkpoints every 1,000 steps; retain the last 10. On a 512-GPU cluster training a 70B model with a 4M-token global batch, 1,000 steps represents roughly 20-30 minutes of wall time. Retaining 10 checkpoints provides a three-to-five-hour rollback window for diagnosing slowly developing issues at the cost of approximately 400 GB of checkpoint storage per saved snapshot.

• Set grad_$\mathtt{clip} = 1 . 0$ for 7B-70B models and log the pre-clip norm every step. The clipped norm is uninformative. The pre-clip norm distribution tells you when something is structurally wrong: a median around $0 . 8$-$1 . 2$ with an occasional $5 \times$ outlier is normal; a median that steadily rises from $1 . 0$ to $3 . 0$ over 10,000 steps is an early signal of logit drift or an overheated learning rate.

• Log log_Z from step zero as a first-class health metric. As detailed in the z-loss subsection, a monotonically increasing log_Z is the clearest multi-step leading indicator of a future divergent spike. If log_Z has exceeded 12 and has been rising for 5,000 steps, add or increase the z-loss coefficient immediately; do not wait for the spike.

• Implement skip-batch logic for Tier 1 events. If loss_ema_ratio exceeds 2.0 at any step, skip the optimizer update for that step. Still run the forward and backward passes to compute the gradient norm for diagnostic purposes, then discard the gradient before optimizer.step(). This prevents a single catastrophic batch from damaging the Adam momentum estimates while the data shard investigation proceeds.

• Stop the run immediately on nan_frac $>$ 0 and do not attempt to continue. A single NaN in the gradient contaminates every parameter touched by that gradient through subsequent momentum accumulation. The affected optimizer state cannot be repaired; only a rollback to a clean checkpoint resolves it. Continuing through NaN values corrupts the master weight copy in a way that may not be visible in the loss for hundreds of steps.

How This Shows Up in Interviews

• “Your 70B pre-training run develops a loss spike at 150B tokens that has not recovered after 3,000 steps. How do you diagnose and fix it?” Answer skeleton: first check whether grad_norm showed a single-step extreme outlier before the spike (data anomaly) or a gradual elevation (logit drift or LR issue). Check whether log_Z has been rising for thousands of steps before the visible spike - if so, the root cause is logit scale instability. Roll back to the last checkpoint before the log_Z acceleration, add z-loss with $\alpha = 10^{- 4}$ if not already present, and resume. If neither the gradient norm nor log_Z explains the spike, inspect the data shard window for corrupted documents or a pipeline error that changed the data distribution.

• “What training health metrics do you log beyond the loss curve?” Answer skeleton: the five-metric set - loss, pre-clip grad_norm, log_Z, loss_ema_ratio, and nan_frac. Explain what each indicates and at what threshold it triggers an action. Mention that per-layer gradient norm histograms (sampled every 100 steps rather than every step to limit overhead) are added when a run is exhibiting unexplained instability and you need to localize which layer is the source.

• Senior follow-up: “Design an automated spike detection and rollback system for a production pretraining run. What would you automate versus leave to human judgment?” A senior answer covers: automate Tier 1 skip-batch on loss_ema_ratio $> 2 . 0$; automate immediate halt on nan_frac $> 0$; automate paging the on-call engineer for Tier 2 events with a pre-populated diagnostic summary (log_Z trend over the prior 5,000 steps, gradient norm percentile history, data shard identifier for the steps around the spike onset). Explain why fully automated Tier 2 rollback is risky: the rollback destination requires understanding whether the spike is a data anomaly that already passed (in which case you roll back past it) or a configuration deficiency that persists (in which case rolling back without configuration change is pointless). The human judgment step is cheap; the cost of a wrong automated rollback is not.

Key Takeaways

• Loss spikes fall into three tiers by root cause - data anomaly (self-recovers within 500 steps, skip-batch response), logit drift or gradient explosion (requires rollback to last stable checkpoint), and numerical precision failure (requires rollback and configuration change) - and each demands a different response.

• The pre-clip gradient norm, log_Z, and nan_frac are the three non-loss metrics that distinguish these tiers; log all five health metrics at every training step.

• Save checkpoints every 1,000 steps and retain at least 10; this bounds the rollback cost at one to three hours of compute while providing a multi-hour diagnostic window.

• A log_Z that is monotonically rising and has exceeded 12 is a reliable multi-thousand-step leading indicator of a divergent spike; the correct response is to add z-loss before the spike occurs, not to roll back after it does.

• Fully automated skip-batch (Tier 1) is safe and worth implementing; automated Tier 2 rollbacks require human sign-off because the correct rollback destination depends on diagnosing the root cause, not just detecting that a spike occurred.