Skip to content
mlmentorship

Mixed precision: what's actually happening?

Beyond 'use BF16'. The senior answer explains what stays in FP32, why loss scaling exists for FP16, and the memory split.

Published · 8 min read ·Core ·Mixed

Systems / infrastructure · active recall

Practice before you read

35 minutes. Quantify the workload, identify the bottleneck, choose an architecture, and design failure recovery.

Back to my workbook How practice works

Systems / infrastructure · closed-book attempt

Mixed precision: what's actually happening?

Quantify the workload, identify the bottleneck, choose an architecture, and design failure recovery.

Problem

Context. A model-training workload seeks the speed or memory benefits of FP16 or BF16 without silently corrupting training. Activations, reductions, gradients, parameters, and optimizer state can have different numerical requirements. Task. Specify a precision policy, explain range and rounding failure modes, and describe how to validate it against a higher-precision baseline. Compare the role of loss scaling for FP16 and BF16, including the order of unscaling, clipping, and stepping. Constraints. Distinguish storage dtype, compute dtype, and accumulation dtype; automatic casting does not choose every policy for you. Account for nonfinite gradients and skipped steps, and do not assume every operation or parameter needs identical precision.

35:00recommended time

Closing or reloading clears the scratchpad. Only score, weak rubric dimensions, attempt count, and retry date can be stored locally.

30-second answer map

Trace one FP16 training step in the order that preserves small gradients without letting scaled or non-finite values corrupt the FP32 parameters.

One FP16 training step

Safe operation order for one FP16 mixed-precision training step

Safe operation order for one FP16 mixed-precision training step Persistent FP32 parameters feed an autocast forward pass in which eligible operations use FP16 and sensitive operations use FP32. The FP32 loss is multiplied by scale S before backward so small FP16 gradients survive. Gradients are converted to FP32 and divided by S before gradient clipping. A finite-value check then branches: non-finite gradients skip the optimizer update and reduce S, while finite unscaled gradients update the FP32 parameters and may eventually allow S to grow. The next step begins from the updated FP32 parameters. PERSISTENT STATE FP32 parameters + optimizer state the next valid update lands here AUTOCAST COMPUTE eligible ops: FP16 matmul · convolution sensitive ops: FP32 reductions · loss scale loss: L ← S · L then backward; small FP16 gradients shift up GUARD THE FP32 UPDATE cast to FP32 and unscale: g ← g / S clip only now, using the true gradient magnitude all gradients finite? no yes skip update reduce S; retry next step optimizer step update FP32 parameters
follow the center path in order. Scaling protects small FP16 gradients during backward, but it must be undone before clipping or interpreting their magnitude. Only finite, unscaled FP32 gradients may reach the optimizer; an overflow skips the update rather than poisoning the persistent parameters. BF16 normally removes the loss-scaling boxes, not the op-specific autocast or FP32 optimizer-state boundaries. Original schematic checked against Micikevicius et al. (2018) and the PyTorch AMP documentation.
Check: Why must clipping use unscaled gradients? Which state changes when a nonfinite update is skipped?
Full interview answer 8 min read

Asked in: ML breadth at LLM-team and infra interviews.

Explain how mixed-precision training assigns numerical formats while preserving useful gradients and stable optimization.

Problem

Context. A model-training workload seeks the speed or memory benefits of FP16 or BF16 without silently corrupting training. Activations, reductions, gradients, parameters, and optimizer state can have different numerical requirements.

Task. Specify a precision policy, explain range and rounding failure modes, and describe how to validate it against a higher-precision baseline. Compare the role of loss scaling for FP16 and BF16, including the order of unscaling, clipping, and stepping.

Constraints. Distinguish storage dtype, compute dtype, and accumulation dtype; automatic casting does not choose every policy for you. Account for nonfinite gradients and skipped steps, and do not assume every operation or parameter needs identical precision.

What an L4 answer sounds like

“Mixed precision uses FP16 or BF16 instead of FP32, which is faster and uses less memory. PyTorch has torch.cuda.amp for it.”

True, no mechanism. You’ve turned it on, never debugged it.

What an L5 answer sounds like

“One common mixed-precision recipe keeps FP32 parameters and optimizer state while autocasting eligible operations to BF16 or FP16. The exact policy depends on the device, kernel, framework, and operation:

  1. Persistent parameters. FP32 model parameters can already be the authoritative weights; a separate master copy is not always allocated.
  2. Operation-specific compute. Autocast chooses eligible lower-precision operations; it does not make every forward/backward value low-precision.
  3. Gradient accumulation. FP32 leaf parameters normally accumulate FP32 gradients, even though intermediate arithmetic may use lower precision. Casting after underflow cannot recover a lost gradient.
  4. Guarded update. Unscale before clipping, reject non-finite updates, and update the persistent parameters with the intended optimizer-state precision.

FP16 training commonly uses loss scaling before backward to preserve small intermediate gradients, followed by unscaling before the optimizer step. BF16’s wider exponent range often removes that need, but finite precision, overflow, and unstable operations still require checks. Preserving a gradient also does not guarantee its weight update is representable; the mixed-precision mechanism separates those cases.

Some operations stay in FP32 for stability:

  • LayerNorm/RMSNorm: variance computation needs precision.
  • Softmax: especially in attention; cast scores to FP32, softmax, cast back.
  • Loss function: usually FP32.

Frameworks (PyTorch autocast) handle most of the casting automatically.”

This is L5. You’ve described the full recipe with the stability ops named.

Learning objective: trace one FP16 training step in the order that preserves small gradients without letting scaled or non-finite values corrupt the FP32 parameters.

One FP16 training step

Unscale before clipping; skip nonfinite parameter updates

Unscale before clipping; reject nonfinite updates before steppingIllustrative gradient ten is scaled by one hundred. Unscale before clipping gives one; clipping the scaled gradient first gives 0.01 after unscaling. A nonfinite result must skip the parameter step.Unscale before clipping to preservethe intended thresholdCorrect: unscale, then clip1000101Wrong: clip, then unscale100010.01Nonfinite values stop before theparameter step.
Read it this way: Illustrative gradient ten is scaled by one hundred. Unscale before clipping gives one; clipping the scaled gradient first gives 0.01 after unscaling. A nonfinite result must skip the parameter step. Which order leaves the intended clipped gradient one instead of 0.01?
Reference drawing

One FP16 training step

Unscale before clipping; skip nonfinite parameter updates

Unscale before clipping; reject nonfinite updates before steppingIllustrative gradient ten is scaled by one hundred. Unscale before clipping gives one; clipping the scaled gradient first gives 0.01 after unscaling. A nonfinite result must skip the parameter step.Unscale before clipping to preservethe intended thresholdCorrect: unscale, then clip1000101Wrong: clip, then unscale100010.01Nonfinite values stop before theparameter step.
Derivation
Safe operation order for one FP16 mixed-precision training step Persistent FP32 parameters feed an autocast forward pass in which eligible operations use FP16 and sensitive operations use FP32. The FP32 loss is multiplied by scale S before backward so small FP16 gradients survive. Gradients are converted to FP32 and divided by S before gradient clipping. A finite-value check then branches: non-finite gradients skip the optimizer update and reduce S, while finite unscaled gradients update the FP32 parameters and may eventually allow S to grow. The next step begins from the updated FP32 parameters. PERSISTENT STATE FP32 parameters + optimizer state the next valid update lands here AUTOCAST COMPUTE eligible ops: FP16 matmul · convolution sensitive ops: FP32 reductions · loss scale loss: L ← S · L then backward; small FP16 gradients shift up GUARD THE FP32 UPDATE cast to FP32 and unscale: g ← g / S clip only now, using the true gradient magnitude all gradients finite? no yes skip update reduce S; retry next step optimizer step update FP32 parameters
follow the center path in order. Scaling protects small FP16 gradients during backward, but it must be undone before clipping or interpreting their magnitude. Only finite, unscaled FP32 gradients may reach the optimizer; an overflow skips the update rather than poisoning the persistent parameters. BF16 normally removes the loss-scaling boxes, not the op-specific autocast or FP32 optimizer-state boundaries. Original schematic checked against Micikevicius et al. (2018) and the PyTorch AMP documentation.
Check: Why must clipping use unscaled gradients? Which state changes when a nonfinite update is skipped?
Read it this way: Illustrative gradient ten is scaled by one hundred. Unscale before clipping gives one; clipping the scaled gradient first gives 0.01 after unscaling. A nonfinite result must skip the parameter step. Which order leaves the intended clipped gradient one instead of 0.01?

What an L6 answer sounds like

“…practical things that bite people:

Use a byte ledger. Two FP32 Adam moment arrays cost eight bytes per parameter; FP32 parameters and their FP32 gradient buffers add four bytes each. Activation savings and temporary casts depend on the implementation. Mixed precision and state sharding address different memory terms.

BF16 changes range, not mathematical stability. It often avoids FP16 range failures, but it does not prevent every NaN or guarantee a throughput ordering. Check device support, kernel behavior, reductions, and task quality.

FP8 (H100+) is the next step. Two variants: E4M3 (more mantissa, less range, for forward) and E5M2 (more range, for backward gradients). Per-tensor scaling factors that update during training. Frameworks like Transformer Engine handle the bookkeeping.

Mixed precision can introduce silent quality regressions. Differences appear in long-context attention (numerical stability of softmax over many positions), in very deep networks, and in RL where reward signals are subtle. Always evaluate on the production task, not just on training loss.

Gradient scaler + grad clipping interact. If you clip gradients, do it on the unscaled gradients (after scaler.unscale_() in PyTorch). Clipping the scaled gradient gives wrong results.”

Tells that get you a strong-hire vote

  • You name the FP32 master weights + low-precision compute pattern explicitly.
  • You explain why FP16 needs loss scaling but BF16 doesn’t.
  • You list ops that stay in FP32 for stability.
  • You acknowledge optimizer state as the memory dominator, not weights.
  • You bring up FP8 for H100-era training.

Tells that get you down-leveled

  • “Just use AMP” with no further detail.
  • Suggesting FP16 over BF16 on modern hardware.
  • No mention of loss scaling for FP16.
  • Confusion between mixed precision (training) and quantization (inference).

Common follow-up

“You said optimizer state dominates memory. Can you mixed-precision the optimizer too?”

The L6 answer:

“Yes, but carefully. Adam’s second moment v is the largest concern; storing it in BF16 loses precision and causes training instability. Approaches: 8-bit Adam (Dettmers et al.) quantizes optimizer state to INT8 with per-block scales, recovering most of the memory with minor quality loss. AdaFactor uses a factored approximation of v with much less memory. For pretraining at scale, ZeRO sharding of FP32 optimizer state across data-parallel ranks is more common than quantizing it.”

Related: Mixed precision training, Walk me through how you’d train a 100B parameter model, Quantization.

Found an error or stale claim? Report a correction or send a private note.