Mixed Precision Training

The last chapter turned access-policy supervised fine-tuning (SFT) into a real training system: a response-token loss, a batch budget, resumable checkpoints, and a held-out export rule on policy_pass_rate. Now the same run needs an arithmetic decision.

Suppose a naive FP16 candidate keeps update state in FP16. A weight near 1.0 can miss its 1e-4 update.

A different batch can fail in forward: an activation near 1e5 crosses FP16's 65504 ceiling. Nothing about the model's task changed; the number system did.

Before changing dtype, trace which values disappear, which values overflow, and which quality gates move.

Mixed precision training runs selected expensive operations in a compact floating-point format while keeping higher precision where training is fragile.

Current PyTorch AMP (Automatic Mixed Precision) 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}

Don't declare BF16 good or FP16 bad in the abstract. See what each format can lose, then measure the survivors under the same validation and hardware conditions.

A floating-point number has two limits

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

Before reading the table, predict the failure for each value in the opening trace. Which format has enough range for 1e5? Which format has enough spacing near 1.0 for a 1e-4 step? Those are separate questions.

The easy mistake is treating "16-bit" as one thing. 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 limits with torch.finfo. Run this on CPU; you don't need an accelerator to inspect the number system.

epsilon_at_1 is the gap between 1.0 and the next representable value near it. smallest_normal and largest_finite describe range. FP16 is finer than BF16 near 1.0, and far narrower in 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 standalone magnitudes, but its fraction has only 7 bits, so a tiny change to a weight near 1.0 can still round away. FP32 update state preserves those accumulated changes.

Tiny updates need an FP32 home

Take one access-policy classifier weight sitting at 1.0. One optimizer step wants to subtract 0.0001. That update is much larger than the smallest BF16 magnitude, and still smaller than the spacing between BF16 values near 1.0.

Predict the stored value before running the cell. A format can represent 1e-4 as a standalone number and still lose it when subtracting from a weight near 1.0.

Both 16-bit parameter values lose this update.

The original mixed-precision recipe kept 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} Ordinary PyTorch AMP gets the same protection by leaving 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 failure. Range is another. Compare an extremely small gradient and a large activation-like value when stored in FP16 and BF16.

The useful prediction is asymmetric: FP16 should lose the small gradient and overflow the large value, while BF16 should keep both magnitudes with some rounding.

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, so it has a similar range and these two magnitudes stay representable, although rounded. The BF16 training study treats that wider range as the main stability advantage over FP16.^{[3]Reference 3A Study of BFLOAT16 for Deep Learning Training.https://arxiv.org/abs/1905.12322}

Loss scaling can rescue the tiny backward gradient. It can't repair a forward activation that already overflowed.

FP16 uses loss scaling to rescue small gradients

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

Gradient clipping belongs after that unscale step. Clipping a gradient of 0.02 while it's still multiplied by 1024 treats it as 20.48, so a threshold of 1.0 would shrink a gradient that was already safely below the intended limit.

For a true gradient of $1.2 \times 10^{-8}$ and a teaching scale $S = 1024$, predict the order before reading the table. The scale must enter before the low-precision backward pass. Multiplying a gradient after FP16 already rounded it to zero can't recover the lost value.

That $1.23 \times 10^{-5}$ value sits below FP16's normal floor of $6.1 \times 10^{-5}$, so the rescue uses subnormals. A larger scale, including PyTorch's default GradScaler start of $2^{16} = 65536$, would push the same gradient into the normal range. The mechanism is the same either way.

Scaling too far causes overflow

A 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.

Before running the next cell, choose the outcome for the quiet pair and the spiky pair at S = 1024. The tiny gradient needs the scale, but the 1e2 spike becomes 1.024e5, past FP16's maximum.

The quiet step keeps scale 1024. The spiky step skips the update and halves the scale. That control flow is what GradScaler automates:

In current PyTorch, torch.amp.GradScaler performs this scale, unscale, finite-check, skip, and update control flow for FP16 training.^{[1]Reference 1Automatic Mixed Precision package - torch.amphttps://docs.pytorch.org/docs/stable/amp.html}

Its default starting scale is 65536.0, it doubles after 2000 finite steps, and it multiplies by 0.5 after overflow. If you inspect or clip gradients, call scaler.unscale_(optimizer) first 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.^{[1]Reference 1Automatic Mixed Precision package - torch.amphttps://docs.pytorch.org/docs/stable/amp.html}

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} A BF16-pretrained model can hit that activation wall when its activations don't fit FP16's 65504 cap.

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. Keep the FP32 update path.

Ordinary AMP does that without a separate master-weight tensor: parameters stay 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}

The 2018 recipe instead copied FP32 master weights into a low-precision compute copy each step.^{[2]Reference 2Mixed Precision Training.https://arxiv.org/abs/1710.03740} Both approaches keep high-precision update state.

They don't keep the same stored parameter dtype, which matters as soon as you start counting bytes.

The boundary to watch is the optimizer update. Compute can use fewer bits, but the value that accumulates tiny changes needs the FP32 path.

This is the CUDA shape you'd use for a real access-policy fine-tune. It isn't marked executable here because it needs an accelerator and a model workload. GradScaler(..., enabled=False) turns scale, unscale, and step into the ordinary FP32 control flow, so BF16 and FP16 can share one loop:

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 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 jobs often use gradient accumulation: several micro-batches contribute to gradients before a single optimizer.step. Before looking at the table, identify that one logical boundary. The scaler contract follows the 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 also divide 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.

Predict the total before looking at the inventory. If only the 800 million saved activations switch from four bytes to two, the activation line halves, but the parameter, gradient, and Adam lines stay put.

The next inventory is deliberately small: 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.

The 2018 paper could say mixed precision "nearly halves" training memory because activations dominated and weights, activations, and gradients were stored in FP16 beside an FP32 master copy.^{[2]Reference 2Mixed Precision Training.https://arxiv.org/abs/1710.03740}

Ordinary AMP, as in the budget above, often leaves parameters and Adam moments in FP32. The next chapter's 16-byte-per-parameter count is yet another recipe: low-precision parameter and gradient copies plus an FP32 master and FP32 Adam moments.

Write down which copies exist before using a bytes-per-parameter slogan.

For large models, sharding methods such as ZeRO and Fully Sharded Data Parallel (FSDP) address the parameter, gradient, 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 mixed-precision 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. A trustworthy run record therefore separates compute_dtype, update_storage_dtype, and reduce_dtype rather than logging a single mixed_precision=true flag.

Decide from runs, not format preference

Once the arithmetic is clear, the question becomes attribution: did the dtype change the result, or did the run change around it? Freeze the access-policy artifact, data, seeds, eval suite, and hardware, then record the rest of the comparison contract before timing anything.

That record turns a speed or memory number into evidence someone else can reproduce. The small contract below carries the artifact, baseline, data, seed, evaluation, hardware, and release gates through the example.

The 0.93 floor matches the SFT chapter's step-200 policy_pass_rate. Keep the previous export rule too: format_pass_rate still has to clear 0.98 before policy rate is allowed to decide. Precision work doesn't get to invent a friendlier metric.

RL can expose a BF16 train-inference mismatch

BF16's wide exponent range makes it a strong default candidate for many training jobs, including this SFT run. 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 next fixture brings the SFT lesson back to experiment tracking. The numbers 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 policy_pass_rate, and target-GPU profile for every precision configuration. One extra BF16 run changes hardware on purpose so the comparison filter has something to reject.

Before reading the output, predict which candidate is ineligible even though its examples-per-second value is highest. A changed accelerator breaks the comparison before speed enters the decision.

Don't promote BF16 because it sounds newer. 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

The fixture leaves two candidates. The next temptation is an 8-bit path. FP8 can cut operand storage again, but it isn't a smaller FP16 switch: its encoding, scale granularity, conversion rules, and kernel path all become part of the training contract.

FP8 has two complementary encodings. E4M3 uses 4 exponent bits and 3 fraction bits, reaches a maximum finite magnitude of 448, and doesn't encode infinity. E5M2 uses 5 exponent bits and 2 fraction bits, reaches 57,344, and keeps infinity encodings.

E4M3 buys nearby resolution; E5M2 buys range. Those limits come from the OCP FP8 specification and the original FP8 formats paper.^{[8]Reference 8FP8 Formats for Deep Learning.https://arxiv.org/abs/2209.05433}

The encoding still doesn't tell you how a training stack scales a tensor. Transformer Engine's tensor-scaled hybrid recipe uses E4M3 for forward activations and weights, then E5M2 for backward gradients.^{[9]Reference 9Using FP8 and FP4 with Transformer Enginehttps://docs.nvidia.com/deeplearning/transformer-engine/user-guide/examples/fp8_primer.html}

Its current-scaling recipe computes an absolute maximum (amax) from the current tensor; delayed scaling estimates that scale from amax history. Both use one FP32 scale per tensor, but current scaling reads the tensor once to find amax and again to scale and cast it, while delayed scaling trades that read for history that can lag a new spike.^{[9]Reference 9Using FP8 and FP4 with Transformer Enginehttps://docs.nvidia.com/deeplearning/transformer-engine/user-guide/examples/fp8_primer.html}

MXFP8 changes the scale granularity rather than pretending one global scale fits every value. The OCP MX specification defines MXFP8 as FP8 elements with one E8M0 power-of-two scale shared by each block of 32 values.

Transformer Engine's MXFP8 recipe uses E4M3 by default, supports a hybrid E4M3-forward/E5M2-backward option, and requires Blackwell (SM100) or later for its native path.^{[9]Reference 9Using FP8 and FP4 with Transformer Enginehttps://docs.nvidia.com/deeplearning/transformer-engine/user-guide/examples/fp8_primer.html}

Storage format and accumulation are separate questions. The FP8 paper expects math to produce wider outputs, while OCP MX says a general dot product should return Float32.^{[8]Reference 8FP8 Formats for Deep Learning.https://arxiv.org/abs/2209.05433}

OCP leaves internal precision and operation order implementation-defined. Transformer Engine documents round-to-nearest-even for current-scaling conversion.^{[9]Reference 9Using FP8 and FP4 with Transformer Enginehttps://docs.nvidia.com/deeplearning/transformer-engine/user-guide/examples/fp8_primer.html}

OCP also distinguishes saturation from overflow: an out-of-range E4M3 value can clamp to 448 under saturation or become NaN under overflow, while E5M2 can clamp to 57,344 or become infinity. Count saturation, non-finite values, scale history, and held-out errors; a finite saturated value can still be wrong.^{[9]Reference 9Using FP8 and FP4 with Transformer Enginehttps://docs.nvidia.com/deeplearning/transformer-engine/user-guide/examples/fp8_primer.html}

Hardware and software decide whether those rules are cheap or even available. Verify accelerator model and count, topology, CUDA and PyTorch versions, Transformer Engine version, tensor dimensions, kernel layout, and any fallback path before calling an FP8 run comparable. A lower-bit label doesn't prove a native Tensor Core path or a faster step.

Don't add FP8 to a training proposal until BF16 or scaled FP16 is measured under the evidence contract above.

Keep model, optimizer, data, seeds, batch shapes, and evaluation fixed. Compare step time, examples per second, peak memory, scaling overhead, saturation and skipped-step counts, loss curves, and the held-out format_pass_rate and policy_pass_rate gates. Precision work should reduce measured cost without creating unexplained convergence risk.

Common pitfalls

A precision failure usually points to a boundary that moved. Start with the symptom, identify which value or state crossed it, then check the smallest change that restores the intended contract.

BF16 is mistaken for an FP32 optimizer replacement

A BF16-only parameter update can stop improving loss even while every gradient is finite. Wide exponent range doesn't provide fine spacing near the current weight, so keep FP32 update state under ordinary AMP and log the storage policy.

FP16 silently loses gradients

Training can look stable while policy_pass_rate lags the FP32 baseline when small unscaled FP16 gradients underflow to zero. Use GradScaler, track non-finite or skipped steps, and compare the declared held-out metric against the same baseline.

Gradient clipping sees scaled values

If clipping behaves erratically or training diverges under FP16 AMP, inspect whether the run clips before scaler.unscale_(optimizer). Unscale first, then apply gradient clipping, and let the scaler perform or skip the optimizer step.

GradScaler steps every micro-batch

If effective batch size collapses or later micro-batches write into already-unscaled gradients, unscale_ / step / update ran inside the accumulation window. Scale and backward every micro-batch, then unscale, clip, step, and update once per optimizer step.

Memory claims omit optimizer state

"Half-memory" planning fails when only activation dtype changed while FP32 parameters and Adam moments remain large. Log a component-level memory profile or accounting budget instead of a dtype slogan.

Distributed bandwidth remains high

BF16 compute can be enabled while cross-worker traffic remains a bottleneck because reduction payloads stayed FP32. Inspect and record reduce_dtype separately, then measure held-out quality and communication changes before promotion.