Articles
Mixed-Precision Training Interview Questions: Overflow, Loss Scaling, and Stable Operations
Practice mixed-precision training questions with numeric overflow and underflow fixtures, GradScaler ordering, stable softmax, and bf16 trade-offs.
“The loss became inf, so increase the loss scale” is a revealing mixed-precision interview mistake. Loss scaling addresses gradients that become too small during a low-precision backward path; it cannot make an already-overflowed forward value finite. A strong answer first locates the first nonfinite tensor, then separates numeric range, rounding precision, stable formulas, and optimizer-step behavior.
This guide uses three original, small numeric fixtures rather than claiming a universal speedup from AMP. The fixtures illustrate explicit float16 arithmetic; autocast chooses operation-specific dtypes and may avoid some of these exact paths. The PyTorch AMP reference and AMP examples provide the behavior rules.
Interview prompt: three failures, three different repairs
An interviewer shows you a training loop that became faster after mixed precision but occasionally returns inf loss and skips optimizer updates. A teammate proposes turning up GradScaler. You are given three observations. First, an explicitly half-precision toy forward computes 10000 × 10000. Second, a tiny gradient from a backward operation is around 1e-8. Third, a custom softmax computes exp(20) and exp(19) before normalization. Which case is overflow, which is underflow, and which requires a stable formula? What evidence would distinguish them in a real model?
Treat each observation as a diagnostic probe, not a claim that PyTorch autocast would execute every listed operation in half precision. The default AMP approach keeps model and optimizer in their ordinary precision and lets autocast choose dtypes for eligible forward operations. PyTorch explicitly advises against manually calling .half() or .bfloat16() on the model or inputs when using autocast. Manually forcing all tensors low precision changes the experiment and can remove autocast's stability choices.
| Toy observation | Numeric failure | What scaling can do |
|---|---|---|
Explicit float16: 10000 × 10000 | Result 1e8 exceeds finite fp16 range | Cannot repair the forward inf |
A backward-path gradient near 1e-8 in fp16 | Can round or flush to zero | Scaling the loss before backward can keep some gradients representable |
Naive half-precision exp(20) then divide | Exponential overflows before normalization | Cannot repair; use a stable softmax/log-sum-exp path |
The standard finite maximum for float16 is 65,504, so 1e8 lies outside its range. bfloat16 uses a wider exponent range, similar in scale to float32, but has fewer significand bits than float16; it trades some precision for range. Choosing bf16 can address an fp16 range failure on supported hardware, but it does not guarantee the accuracy of every operation or solve a bad loss formula. PyTorch's AMP docs specifically warn that bf16-pretrained models can exceed fp16's numeric range.
Exercise 1: find the first nonfinite value
Start with a deliberately explicit half-precision test, outside autocast:
import torch
a = torch.tensor(10000.0, device="cuda", dtype=torch.float16)
product = a * a
print(product) # inf: the result is above fp16's finite range
The same mathematical product is 100,000,000, finite in float32. This is forward overflow, so multiplying the loss afterward cannot recover the lost number. To repair the actual model, trace the first nonfinite operation. Compare the corresponding input and output under fp32, fp16 autocast, and bf16 autocast if hardware supports it. Record per-op dtypes; “AMP is on” does not state the dtype of a particular op. Ask whether the input scale is unexpectedly large, whether a matrix reduction accumulates in the intended precision, and whether a custom op or explicit cast bypassed autocast.
Do not assume every fp16 GEMM simply accumulates in fp16. PyTorch's numerical accuracy notes describe fp32 intermediate accumulation and hardware-dependent reduced-precision reductions that can affect results. That is why a real answer names the kernel, dtype, hardware, and configuration before prescribing a flag. A targeted fp32 subregion can be appropriate when one operation is sensitive; widening the entire model is an experiment, not the only possible fix.
A frequent trap is to check only the final loss. If an early activation overflows, later normalization or clipping may yield NaN, but the first bad value is the cause. Add temporary finite-value checks at inputs, selected activations, logits, loss, and gradients; narrow by binary search through the graph. Capture the first failing batch and seed. If the same batch is finite in fp32 and bf16 but not fp16, that supports a range/precision hypothesis. If it is nonfinite in fp32 too, investigate model/data math or an excessive learning rate before blaming AMP.
Exercise 2: explain what GradScaler actually changes
Many low-precision backward operations inherit the dtype used for the corresponding forward operation. A gradient small enough to fall below the effective fp16 representable range can become zero; the parameter receives no useful update from that contribution. A loss scale S multiplies the loss before backward, so the chain-rule gradients are also multiplied by S. Before an optimizer step, the gradients are divided back by S. The optimizer should see the intended unscaled gradient if the scaled path stayed finite.
For an illustrative gradient of 1e-8, multiplying by 1024 yields 1.024e-5 during backward—a value much easier to represent in fp16—then unscaling restores the intended magnitude in the parameter gradient. This arithmetic is a conceptual example, not a promise that any particular kernel preserves the result exactly. If the forward value was already cast to zero or inf, scaling the later loss cannot recreate it. Likewise, a scale chosen too large can itself overflow gradients; dynamic scaling adjusts according to observed nonfinite gradients.
scaler = torch.amp.GradScaler("cuda")
optimizer.zero_grad(set_to_none=True)
with torch.autocast(device_type="cuda", dtype=torch.float16):
output = model(inputs)
loss = loss_fn(output, targets)
scaler.scale(loss).backward()
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
scaler.step(optimizer)
scaler.update()
This order matters. PyTorch's AMP examples specify unscaling before clipping so a threshold of 1.0 applies to the real gradients, not gradients enlarged by S. scaler.step(optimizer) checks for nonfinite gradients and skips optimizer.step() when it finds them; scaler.update() adjusts the scale afterward. A skipped optimizer step is a protective symptom, not evidence that the loss has been fixed. Log the scale trend, nonfinite count, and actual optimizer updates in the same time series as loss.
If you accumulate gradients over multiple microbatches, keep one scale across that effective batch, divide the loss by the accumulation count when appropriate, and unscale/clip/step/update only at the effective-batch boundary. Mixing scaled and unscaled gradients in one accumulation window loses the meaning of the sum. The official gradient-accumulation example makes this ordering explicit. A scheduler should usually advance in sync with successful optimizer updates for a step-based schedule; define that policy rather than allowing a skipped step to silently desynchronize the run.
Exercise 3: fix an unstable formula, not the loss scale
The naive softmax of logits [20, 19] computes exp(20) and exp(19). exp(20) is roughly 4.85 × 10^8, beyond fp16's finite range. Yet the desired probabilities are well behaved: subtract the maximum first, giving exponentials of [0, -1]. Then normalize 1 and e^-1, which yields approximately [0.731, 0.269]. A stable logsumexp or fused cross-entropy path applies this principle rather than materializing huge exponentials.
logits = torch.tensor([20.0, 19.0], device="cuda", dtype=torch.float16)
stable = torch.softmax(logits, dim=0)
# Do not implement softmax as logits.exp() / logits.exp().sum().
The snippet asks the library's softmax to handle the operation; it is not a universal guarantee about every backend, dtype, or custom implementation. If interviewing on a custom loss, derive the max-subtraction and test extreme logits, equal logits, very negative logits, and the backward gradient. For contrastive learning, small temperatures magnify logits and can expose the same instability. Normalize embeddings, check temperature, and use a stable cross-entropy implementation while preserving the intended objective; do not quietly clip scores until they look finite.
Other sensitive operations include reductions, variance-like calculations, norms, and probabilities near zero. Some should execute or accumulate in fp32 for numerical stability. With autocast, inspect the op-specific policy and output dtype rather than assuming every operation is half precision. If a custom function requires fp32, use a narrow autocast(enabled=False) subregion and explicitly cast its inputs to fp32, as shown in PyTorch's autocast documentation. Then benchmark the local widening and verify the gradients; an apparently finite loss can still have degraded accuracy.
Range, precision, and optimizer state are different questions
Range asks whether the magnitude can be represented at all. Precision asks whether nearby representable values are fine enough for your calculation. Optimizer state asks what dtype stores parameters, momentum, and other accumulated quantities. These are related but not interchangeable. Bfloat16's wider range can reduce overflow risk for large values, while its coarser mantissa can increase rounding error on small differences. Fp32 accumulation or master parameters may be important even when many forward matrix operations run lower precision.
An autocast training loop typically creates model and optimizer in default precision and casts eligible operations as needed; it does not mean every stored parameter or optimizer state has become fp16. A deliberately all-half model is a different setup with different risks. In the interview, inspect actual param.dtype, gradient dtype, optimizer state dtype, op outputs, and checkpoint contents. If an optimizer update is too small to change a low-precision parameter, loss scaling alone cannot solve the storage precision issue after gradients are unscaled.
Use a controlled comparison: fp32 baseline, fp16 autocast with scaler, and bf16 autocast where supported, all on the same fixed data subset and initial weights. Compare first nonfinite location, skipped-step frequency, gradient norms after unscale, validation quality, peak memory, and true examples per second. Hardware capability and kernel choice can change the performance outcome. A faster step that skips updates or diverges is not a training speedup.
A concise answer under interview pressure
For the 10000 × 10000 probe, say “forward fp16 range overflow; find or widen that operation—scaling the loss later cannot undo inf.” For the 1e-8 gradient probe, say “backward underflow is where loss scaling may help, provided the forward path remained valid and the scale itself does not overflow.” For the [20,19] logits, say “use stable softmax/log-sum-exp or fused loss; do not form huge exponentials.” Then explain the correct AMP step order: autocast forward/loss, scale backward, unscale before clipping, conditional optimizer step, update scale.
Finally, qualify the diagnosis. These are explicit half-precision toy calculations, not proof of the real model's autocast operator policy. Ask for the first bad tensor and a fp32/bf16 comparison on the same failing batch. That answer demonstrates numerical reasoning and a safe experimental plan, rather than reciting a list of AMP flags.