Mixed Precision Training

The previous lesson recorded an application fix as a controlled experiment. The same discipline applies when the changed artifact is a trained model: a faster run isn't an improvement if it corrupts updates, overflows, or lowers the declared held-out metric.

Suppose the order-support pipeline now needs a small carrier-event classifier that recognizes whether a delivery estimate is supported by a scan. You want to fine-tune it faster, so the proposed run changes its precision policy from FP32 to FP16 or BF16. This lesson builds the run evidence you would need before choosing that policy.

Make precision an explicit experiment parameter

Mixed precision training runs selected expensive operations in a compact floating-point format while retaining higher precision where training is fragile. In modern PyTorch, model parameters normally remain FP32 while autocast chooses lower-precision computation for eligible operations.^[1]

Start as you would in an experiment tracker: state the artifact under test and its guardrails before measuring candidates.

The job isn't to declare BF16 good or FP16 bad in the abstract. It's to understand what each format can lose, then measure the acceptable candidates under the same validation and hardware conditions.

A floating-point number has two limits

A floating-point value is similar to scientific notation: a sign, a scale, and significant digits. Its exponent controls range, meaning how tiny or large a magnitude it can represent. Its fraction (often called the mantissa in training discussions) controls resolution, meaning how close two neighboring values can be.

The important correction is easy to miss: BF16 improves range relative to FP16; it doesn't improve nearby resolution. BF16 has fewer fraction bits than FP16. That is why BF16 compute still normally updates FP32 parameters.

PyTorch exposes the exact format limits with torch.finfo. Run this on CPU; no accelerator is required to inspect the number system.

epsilon_at_1 is the spacing between 1.0 and the next representable value near it. smallest_normal and largest_finite describe range. FP16 gives finer resolution than BF16 near 1.0, but far less range.

Tiny updates need an FP32 home

Assume the classifier has a weight at 1.0. One optimizer step wants to subtract 0.0001. This is much larger than the smallest BF16 magnitude, but smaller than the spacing between BF16 values near 1.0.

Both 16-bit parameter values lose this update. The original mixed-precision recipe used an FP32 master copy of the weights so small updates accumulate instead of disappearing.^[2] Modern PyTorch AMP normally gets the same protection by leaving parameters in FP32 and autocasting eligible forward operations rather than converting the parameter storage itself.^[1]

Range decides whether gradients exist at all

Resolution is one issue. Range is another. Compare an extremely small gradient and a very large activation-like value when stored in FP16 and BF16.

FP16's smallest positive normal number is roughly $6.1 \times 10^{-5}$, and its subnormal floor is about $6.0 \times 10^{-8}$. A true gradient of $1.2 \times 10^{-8}$ becomes zero in FP16. At the other end, 100000 is beyond FP16's largest finite value of 65504, so it becomes Inf.

BF16 keeps the 8-bit exponent width of FP32, giving it a similar range and making these two magnitudes representable, although rounded. The BF16 training study documents that wider range as its main stability advantage over FP16.^[3]

FP16 uses loss scaling to rescue small gradients

For FP16, loss scaling moves gradient magnitudes into a representable interval during backpropagation. Multiply loss by a scale $S$; the chain rule multiplies each gradient by $S$ too. After backward, divide gradients by $S$ in FP32 before applying the optimizer step. The intended update has not changed.

For a true gradient of $1.2 \times 10^{-8}$:

Scaling too far causes overflow

A fixed scale that saves the smallest gradient may overflow a larger gradient in the same step. Dynamic scaling therefore has two outcomes: apply a finite, descaled update, or skip an overflowed step and reduce the scale.

PyTorch's torch.amp.GradScaler performs this scale, unscale, finite-check, skip, and update control flow for FP16 training. PyTorch also documents that if you inspect or clip gradients, you must call scaler.unscale_(optimizer) before clipping so thresholds apply to true gradient magnitudes.^[1]

Loss scaling isn't a general extension of FP16 range. It rescues small backward gradients that would underflow, but it can't make a forward activation above 65504 representable. PyTorch also warns that GradScaler may reduce its scale below 1 for overflow-prone models, so don't assume the scale always grows or stays above 1.^[1]

Compute low, update high

Loss scaling protects FP16 gradients from range failure. It doesn't make 16-bit parameter storage appropriate for tiny updates. Preserve the FP32 update path separately.

The original paper describes copying FP32 master weights into a low-precision compute copy.^[2] With ordinary AMP, PyTorch parameters remain FP32, autocast selects lower precision for eligible compute, and the optimizer updates the FP32 parameters directly.^[1]

This is the CUDA shape you would use for a real fine-tuning run. It isn't marked executable here because it needs an accelerator and a model workload:

For BF16, skipping GradScaler is a common policy because BF16's exponent range avoids the FP16 failure that loss scaling targets. It isn't a guarantee that every BF16 training job is stable. Bad data, unstable losses, overly large learning rates, or sensitive kernels can still produce non-finite values.

Memory savings need accounting, not slogans

When autocast runs an eligible operation in BF16 or FP16, its saved low-precision activation values use two bytes rather than four. Other operations stay in or return FP32 for numerical safety. Total training memory doesn't necessarily halve because FP32 parameters and optimizer moments may remain unchanged.

The next accounting exercise uses a deliberately small inventory: 100 million parameters, their gradients, two Adam moment buffers, and 800 million stored activation values. To make the arithmetic visible, assume those inventoried activations were saved in low precision for the AMP candidate. A real profile may retain FP32 values for some operations. This is a budget calculation, not a measured GPU profile.

For large models, sharding methods such as ZeRO and Fully Sharded Data Parallel (FSDP) address parameters, gradients, and optimizer-state memory that activation casting alone doesn't remove.^[4][5]

Distributed jobs add a communication dtype

When workers exchange gradients, the network payload has its own precision policy. FSDP's MixedPrecision configuration exposes param_dtype for computation and reduce_dtype for gradient reduction; they need not match.^[6]

A job can use BF16 for matrix computations and still communicate FP32 gradient payloads. Therefore a trustworthy run record separates compute_dtype, update_storage_dtype, and reduce_dtype rather than logging a single mixed_precision=true flag.

Decide from runs, not format preference

The final cell brings the lesson back to experiment tracking. The numbers below are illustrative recorded outcomes, not benchmark claims. They show the review rule you should apply after running the same classifier, data fingerprint, seed policy, held-out supported_evidence_f1 evaluation, and hardware profile for every precision configuration.

The right result isn't "BF16 wins because it's modern." Both scaled FP16 and BF16 pass this small fixture, and both require a real measurement on target hardware. BF16 is often simpler to operate because it commonly avoids loss scaling, but only the controlled run can justify promotion.

FP8 is a later optimization, not a default answer

FP8 reduces compute storage again, but its reduced range and resolution require managed scaling recipes. FP8 isn't one layout: the FP8 formats paper specifies complementary E4M3 and E5M2 encodings for deep-learning workloads. NVIDIA Transformer Engine documents a hybrid recipe that uses E4M3 during the forward pass and E5M2 during the backward pass, plus delayed, current, and block-scaling recipes for supported accelerators.^[7][8]

That is enough orientation here. Don't add FP8 to a training proposal until BF16 or scaled FP16 is measured, quality gates exist, and the team can operate the scaling policy. Precision work should reduce measured cost without creating unexplained convergence risk.

Mastery check

Key concepts

• Exponent range versus fraction resolution
• Why BF16 helps range but still needs FP32 update state
• FP16 loss scaling, finite checks, and overflow backoff
• AMP/autocast plus GradScaler control flow
• Activation memory versus parameter and optimizer-state memory
• Communication dtype as a separate distributed knob
• Precision selection as a logged experiment decision

Evaluation rubric

• Foundational: Reads torch.finfo results and explains why FP16 and BF16 fail in different ways.
• Foundational: Shows why a small update to an FP16 or BF16 stored parameter can disappear while FP32 preserves it.
• Intermediate: Explains loss scaling without claiming it changes the true optimization objective.
• Intermediate: Reads the AMP CUDA shape and places unscaling before gradient clipping.
• Intermediate: Computes a memory budget without promising that mixed precision halves total training memory.
• Advanced: Records compute, update-storage, and reduce dtypes separately in a distributed experiment.
• Advanced: Refuses to promote BF16 or scaled FP16 until comparable target-hardware measurements satisfy the declared held-out metric and numerical-stability gates.

Follow-up questions

Common pitfalls

BF16 is mistaken for an FP32 optimizer replacement

Symptom: A BF16-only parameter update stops improving loss even though gradients are finite. Cause: Wide range was confused with fine resolution near current weights. Fix: Keep FP32 update state under ordinary AMP and log the storage policy.

FP16 silently loses gradients

Symptom: Training appears stable but supported_evidence_f1 lags the FP32 baseline. Cause: Small unscaled FP16 gradients underflow to zero. Fix: Use GradScaler for FP16, track non-finite or skipped steps, and compare the declared held-out metric against the same baseline.

Gradient clipping sees scaled values

Symptom: Clipping behaves erratically or training diverges under FP16 AMP. Cause: The run clips gradients before scaler.unscale_(optimizer). Fix: Unscale first, then clip, then let the scaler perform or skip the optimizer step.

Memory claims omit optimizer state

Symptom: "Half-memory" planning fails when the job is scheduled. Cause: Only activation dtype changed while FP32 parameters and Adam moments remain large. Fix: Log a component-level memory profile or accounting budget, not a dtype slogan.

Distributed bandwidth remains high

Symptom: BF16 compute is enabled, but cross-worker traffic is still a bottleneck. Cause: Reduction payloads remain FP32. Fix: Inspect and record reduce_dtype separately, then measure held-out metric and communication changes before promotion.