Mixed Precision Training

A training-speed change is an experiment, not an automatic improvement. A faster run isn't better if it corrupts updates, overflows, or lowers the declared held-out metric.

Suppose the deploy assistant now needs a small trace-evidence classifier that recognizes whether an incident summary is supported by logs. You want to fine-tune it faster, so the proposed run changes its precision policy from FP32 to FP16 or BF16. Build 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. Current PyTorch AMP examples create the model and optimizer in default precision, then let autocast choose an operation-specific dtype inside the forward region.^{[1]Reference 1Automatic Mixed Precision package - torch.amphttps://docs.pytorch.org/docs/stable/amp.html}

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's 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.

If BF16 has FP32-like range, why don't we store model updates only in BF16?

Answer

Range and resolution are different. BF16 can represent very small magnitudes, but its fraction has only 7 bits, so tiny changes to a weight near 1.0 can still round away. FP32 update state preserves those accumulated changes.

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]Reference 2Mixed Precision Training.https://arxiv.org/abs/1710.03740} Current PyTorch AMP gets the same protection in its ordinary pattern by creating parameters in default precision and autocasting eligible forward operations rather than converting parameter storage before the optimizer step.^{[1]Reference 1Automatic Mixed Precision package - torch.amphttps://docs.pytorch.org/docs/stable/amp.html}

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]Reference 3A Study of BFLOAT16 for Deep Learning Training.https://arxiv.org/abs/1905.12322}

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.

In current PyTorch, 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]Reference 1Automatic Mixed Precision package - torch.amphttps://docs.pytorch.org/docs/stable/amp.html}

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]Reference 1Automatic Mixed Precision package - torch.amphttps://docs.pytorch.org/docs/stable/amp.html}

An FP16 step overflows after backward. Should the optimizer apply the finite gradients and drop only the non-finite tensors?

Answer

No. Skip the whole optimizer step, reduce the loss scale, and retry on a later batch. Applying a partial update changes the optimizer contract; clipping must also happen only after gradients are unscaled to their true magnitudes.

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]Reference 2Mixed Precision Training.https://arxiv.org/abs/1710.03740} With ordinary AMP, PyTorch parameters remain FP32, autocast selects lower precision for eligible compute, and the optimizer updates the FP32 parameters directly.^{[1]Reference 1Automatic Mixed Precision package - torch.amphttps://docs.pytorch.org/docs/stable/amp.html}

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.

GradScaler and gradient accumulation share one optimizer-step boundary

The sketch above assumes one micro-batch per optimizer step. Real SFT and distributed jobs often use gradient accumulation: several micro-batches contribute to gradients before a single optimizer.step. The scaler contract still keys off that optimizer step, not each micro-batch.

If you step/update every micro-batch, the effective batch size collapses and the scale adapts on incomplete gradients. If you unscale before the accumulation window finishes, later micro-batches write into already-unscaled grad buffers. Count skipped non-finite optimizer steps in telemetry; a silent skip storm is a failed precision experiment even when the run "completes."

When loss is averaged over the accumulation window, keep the usual loss / accum_steps scaling and let GradScaler multiply that micro-loss. Don't double-count accumulation by also dividing the GradScaler scale by hand.

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.

Saved activations use two bytes instead of four, yet total measured memory falls by only one third. Is mixed precision broken?

Answer

No. Parameters, gradients, and Adam moments may still use FP32 storage, so only one part of the inventory halved. Profile each component before deciding whether activation casting, sharding, checkpointing, or a smaller batch targets the real bottleneck.

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]Reference 4ZeRO: Memory Optimizations Toward Training Trillion Parameter Models.https://arxiv.org/abs/1910.02054[5]Reference 5PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel.https://arxiv.org/abs/2304.11277}

Distributed jobs add a communication dtype

When workers exchange gradients, the network payload has its own precision policy. Current PyTorch FSDP MixedPrecision configuration exposes param_dtype for forward and backward computation and reduce_dtype for gradient reduction; the two fields may differ.^{[6]Reference 6FullyShardedDataParallelhttps://docs.pytorch.org/docs/stable/fsdp.html}

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

RL can expose a BF16 train-inference mismatch

BF16's wide exponent range makes it a strong default candidate for many training jobs, but range isn't the only concern. Reinforcement-learning fine-tuning compares or updates a training policy using trajectories generated by an inference policy. Different kernels, batching shapes, or execution paths can produce nearby FP32 logits before casting. BF16's coarse spacing may round those nearby values to different sides of a representable boundary, so rollout probabilities no longer match the probabilities used during training.

Qi et al. report this failure mode in LLM reinforcement-learning fine-tuning and find that using FP16 consistently reduced the train-inference mismatch and improved optimization in their tested tasks, algorithms, and frameworks.^{[7]Reference 7Defeating the Training-Inference Mismatch via FP16https://arxiv.org/abs/2510.26788} That's evidence for an RL-specific candidate, not a universal reversal of the BF16 guidance above. FP16 still needs loss scaling for small gradients, finite forward values, and the same controlled quality checks as every precision policy.

The toy below doesn't reproduce a full RL run. It isolates the rounding boundary. Two execution paths produce first logits 1.0038 and 1.0040. FP16 rounds both to the same value, while BF16 rounds one down and the other up, creating a probability gap.

An RL precision contract should therefore add a direct train-versus-rollout comparison: record logits or log probabilities for matched prompts, measure divergence under the actual training and generation stacks, and check reward and held-out quality across seeds. Promote scaled FP16 over BF16 only when it closes that measured mismatch without introducing overflow or quality regressions.

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 target-GPU profile for every precision configuration. One extra BF16 run changes hardware on purpose so the comparison filter has something to reject.

The right result isn't "BF16 wins because it's modern." Both scaled FP16 and BF16 pass this small fixture, and both require real measurements under the declared contract. The faster run_bf16_other_hardware value can't rank against them because its accelerator changed. BF16 is often simpler to operate because it commonly avoids loss scaling, but only a comparable 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 2.16.0 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.^{[8]Reference 8FP8 Formats for Deep Learning.https://arxiv.org/abs/2209.05433[9]Reference 9Using FP8 and FP4 with Transformer Enginehttps://docs.nvidia.com/deeplearning/transformer-engine/user-guide/examples/fp8_primer.html}

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

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.