Mixed Precision Training

The previous chapter built a supervised fine-tuning (SFT) loop with a response-token loss and held-out evaluation. Now we want that training loop to run faster and fit larger batches into VRAM without corrupting the updates that make the network learn.

Supported GPU Tensor Core paths can accelerate 16-bit matrix multiplications (GEMMs), and a 16-bit tensor uses half the bytes of an equally shaped FP32 tensor. Neither fact predicts the whole job's speed or traffic: hardware, kernels, shapes, casts, and the FP32 baseline matter. Calling .half() changes parameter storage, not just compute. That can lose both small gradients and small updates to existing weights.

Suppose a naive FP16 candidate keeps update state in FP16. A weight near 1.0 misses its 1e-4 update because the format can't represent a change that small near one.

A different batch fails in forward execution: an activation or attention logit near 1e5 crosses FP16's ceiling of 65504 and turns into Inf. Nothing about the model's task changed; the number system did.

These are distinct numerical failures: one loses a subtle update to an existing weight; the other can't represent the magnitude of the signal at all.

Mixed precision training assigns different dtypes to different operations and state. A conventional recipe uses compact GEMM operands while retaining FP32 update state and sensitive accumulations.^{[1]Reference 1Mixed Precision Training.https://arxiv.org/abs/1710.03740} It still needs range checks; FP32 updates can't repair an already-overflowed low-precision forward operation.

In a conventional PyTorch Automatic Mixed Precision (AMP) loop, create FP32 model parameters and their optimizer yourself, then let autocast choose an operation-specific dtype inside the forward region.^{[2]Reference 2Automatic 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, understand how bits are allocated, and measure the survivors under identical validation and hardware conditions.

A floating-point number has two limits

A finite, normal binary floating-point value has a sign bit $s$, an exponent field $e$, and a fraction field $m$:

$x = (-1)^s \times 2^{e - \text{bias}} \times \left(1 + \frac{m}{2^{\text{fraction\_bits}}}\right)$

This equation uses the implicit leading 1 of a normal number. Subnormals use a leading 0 and the minimum normal exponent; zero, infinity, and NaN have separate encodings.

Its exponent controls dynamic range: how tiny or enormous a magnitude it can represent before hitting underflow or overflow. Its fraction controls numerical resolution: how close two neighboring numbers can be, which dictates machine epsilon ($\epsilon$), the gap from 1.0 to the next larger representable value.

In a 16-bit budget, you spend 1 bit on the sign, leaving exactly 15 bits to split between the exponent and the fraction. FP16 and BF16 make opposite design bets with those 15 bits:

The easy mistake is treating "16-bit" as one uniform format:

1. FP16 (IEEE 754 half-precision) dedicates 10 bits to the fraction and 5 to the exponent. Its largest finite value is $65504$, smallest positive normal is $2^{-14} \approx 6.10 \times 10^{-5}$, and smallest positive subnormal is $2^{-24} \approx 5.96 \times 10^{-8}$. Round-to-nearest can round a value below that subnormal minimum up to it. A value below half the minimum rounds to zero; the exact midpoint also rounds to zero under ties-to-even.
2. BF16 (Brain floating point) keeps FP32's 8 exponent bits and bias, with only 7 fraction bits. The minimum normal matches FP32, but its maximum finite value is slightly smaller and its encoded subnormal minimum is $2^{-133}$, versus FP32's $2^{-149}$. Conversion needn't truncate: Cloud TPU uses nearest-even rounding and flushes BF16 subnormals to zero.^{[3]Reference 3Improve Your Model's Performance with bfloat16https://docs.cloud.google.com/tpu/docs/bfloat16} Wide range usually makes FP16-style loss scaling unnecessary, not underflow impossible.
3. BF16 trades away resolution for range. With only 7 fraction bits, BF16 has about 2.1 decimal digits of precision. Its machine epsilon near 1.0 is $2^{-7} \approx 7.81 \times 10^{-3}$, which is eight times coarser than FP16 ($2^{-10} \approx 9.77 \times 10^{-4}$) and over 65,000 times coarser than FP32.

PyTorch exposes these limits directly with torch.finfo. Run this on CPU; you don't need an accelerator to inspect the arithmetic boundaries:

import torch

formats = (
    ("FP32", torch.float32),
    ("FP16", torch.float16),
    ("BF16", torch.bfloat16),
)

print("format  epsilon_at_1  smallest_normal  largest_finite")
for label, dtype in formats:
    info = torch.finfo(dtype)
    print(f"{label:<6}  {info.eps:>12.1e}  {info.tiny:>15.1e}  {info.max:>14.1e}")

print(f"bf16_min_normal_matches_fp32={torch.finfo(torch.bfloat16).tiny == torch.finfo(torch.float32).tiny}")
print(f"bf16_resolution_coarser_than_fp16={torch.finfo(torch.bfloat16).eps > torch.finfo(torch.float16).eps}")

format  epsilon_at_1  smallest_normal  largest_finite
FP32         1.2e-07          1.2e-38         3.4e+38
FP16         9.8e-04          6.1e-05         6.6e+04
BF16         7.8e-03          1.2e-38         3.4e+38
bf16_min_normal_matches_fp32=True
bf16_resolution_coarser_than_fp16=True

epsilon_at_1 is the gap from 1.0 to the next larger representable value. Spacing isn't constant across all numbers: it scales with the current power of two. smallest_normal marks where numbers lose their implicit leading 1 and enter the subnormal range. FP16 is finer than BF16 near 1.0, but drastically narrower in range.^{[4]Reference 4Type Infohttps://docs.pytorch.org/docs/stable/type_info.html}

Inspect both neighbors before judging an update

Take one weight sitting at 1.0. One optimizer step wants to subtract 0.0001 ($10^{-4}$). BF16 easily represents that magnitude on its own, yet applying it to a much larger stored weight can erase the change.

This is swamping: rounding loses a small addition to a larger value. For an exactly representable weight, compare the proposed update with half the gap to its neighbor in the update's direction. Equality needs the rounding tie rule. There isn't one exact threshold equal to $\epsilon |w| / 2$ for every weight and both directions.

At the power-of-two boundary 1.0, the gap below is half the gap above:

Ask a sharper question than "is the update below epsilon?": which neighbor could it reach? nextafter lets us inspect both sides.^{[5]Reference 5torch.nextafterhttps://docs.pytorch.org/docs/2.14/generated/torch.nextafter.html}

import torch

for label, dtype, decrement in (("FP16", torch.float16, 0.0003),
                                ("BF16", torch.bfloat16, 0.0025)):
    one = torch.tensor(1.0, dtype=dtype)
    below = torch.nextafter(one, torch.tensor(0.0, dtype=dtype))
    above = torch.nextafter(one, torch.tensor(float("inf"), dtype=dtype))
    after = one - torch.tensor(decrement, dtype=dtype)
    assert decrement < torch.finfo(dtype).eps / 2 and after != one
    print(f"{label}: gap_below={1.0 - below.item():.8f}, "
          f"gap_above={above.item() - 1.0:.8f}, decremented={after.item():.8f}")

subnormal = torch.tensor(4e-8).to(torch.float16)
assert subnormal > 0
print(f"below_min_subnormal_can_round_up={subnormal.item():.2e}")

FP16: gap_below=0.00048828, gap_above=0.00097656, decremented=0.99951172
BF16: gap_below=0.00390625, gap_above=0.00781250, decremented=0.99609375
below_min_subnormal_can_round_up=5.96e-08

Both decrements survive even though each is below half the upward gap. Our smaller 1e-4 decrement below still rounds away in either 16-bit format.

Predict the stored value before running the cell below. A 16-bit format can represent 1e-4 as a standalone number and still discard it completely when updating a weight near 1.0:

weight = torch.tensor([1.0], dtype=torch.float32)
update = torch.tensor([1.0e-4], dtype=torch.float32)

for label, dtype in formats:
    before = weight.to(dtype)
    after = before - update.to(dtype)
    changed = bool(after.item() != before.item())
    print(f"{label}: stored_after_step={after.item():.8f}, update_survived={changed}")

FP32: stored_after_step=0.99989998, update_survived=True
FP16: stored_after_step=1.00000000, update_survived=False
BF16: stored_after_step=1.00000000, update_survived=False

Both 16-bit parameter values lose this particular update. Repeated deterministic updates below the directional rounding threshold can leave a coordinate unchanged despite nonzero gradients; not every 16-bit run freezes at initialization.

For our ordinary AMP baseline, keep a high-precision update path:

Baseline storage policy: Keep parameters and optimizer moments in FP32 while eligible operations use reduced precision. This protects the demonstrated update; it doesn't make arithmetic exact or every operation low precision.

The original 2018 recipe maintained an FP32 master copy and an FP16 compute copy, while explicitly noting that master weights weren't universally necessary.^{[1]Reference 1Mixed Precision Training.https://arxiv.org/abs/1710.03740} Ordinary PyTorch AMP lets FP32 parameters remain the optimizer's update destination. Autocast selects eligible operation dtypes and can cache lower-precision weight conversions; it doesn't automatically convert existing low-precision parameters back to FP32.^{[2]Reference 2Automatic Mixed Precision package - torch.amphttps://docs.pytorch.org/docs/stable/amp.html}

Other storage policies need different safeguards. Ozkara et al.'s 2025 study trained LLMs with BF16 state and stochastic rounding, using temporary FP32 update arithmetic.^{[6]Reference 6Stochastic Rounding for LLM Training: Theory and Practicehttps://arxiv.org/abs/2502.20566} A small change can survive probabilistically rather than disappearing at every nearest-rounding step. That is a separate recipe with its own measured evidence, not something ordinary autocast enables. FP32 state is our baseline choice, not a mathematical requirement for all training systems.

Range decides whether gradients exist at all

Resolution causes failure when an update is too small relative to an existing weight. Range causes failure when a value is too small or too large to exist on its own.

Backpropagation can produce small magnitudes, but their distribution depends on the model, normalization, residual paths, and loss. Inspect it rather than assuming every deep model's gradients lie in a prescribed range. The following two values are arithmetic fixtures.

Compare an extremely small gradient ($1.2 \times 10^{-8}$) and a large activation-like magnitude ($1.0 \times 10^5$) across FP16 and BF16:

values = torch.tensor([1.2e-8, 1.0e5], dtype=torch.float32)

for label, dtype in (("FP16", torch.float16), ("BF16", torch.bfloat16)):
    cast = values.to(dtype)
    print(
        f"{label}: small={cast[0].item():.2e}, "
        f"large={cast[1].item():.2e}, "
        f"all_finite={bool(torch.isfinite(cast).all())}"
    )

fp16 = values.to(torch.float16)
bf16 = values.to(torch.bfloat16)
assert fp16[0] == 0 and torch.isinf(fp16[1])
assert bf16[0] > 0 and torch.isfinite(bf16).all()

FP16: small=0.00e+00, large=inf, all_finite=False
BF16: small=1.20e-08, large=9.98e+04, all_finite=True

FP16 fails at both ends:

1. Underflow: $1.2 \times 10^{-8}$ is below half FP16's minimum positive subnormal, so nearest rounding produces zero. Subnormal handling also depends on the operation and hardware; this CPU cast isn't a GPU-kernel test. A later multiplication can't reconstruct a value already lost.
2. Overflow: Casting $100000$ to FP16 produces infinity. Non-finite activations can contaminate later losses or gradients. That doesn't inevitably destroy every parameter: GradScaler can skip an affected optimizer's update. Skipping prevents that update, not the original forward overflow.

BF16 keeps these two cast results finite and the small value nonzero, with rounding error. This fixture doesn't prove every BF16 activation, gradient, or accumulated update is safe.^{[7]Reference 7A Study of BFLOAT16 for Deep Learning Training.https://arxiv.org/abs/1905.12322}

Loss scaling can rescue an underflowing backward gradient in FP16. But loss scaling can't repair a forward activation that already overflowed past $65504$.

FP16 uses loss scaling to rescue small gradients

For FP16, loss scaling shifts gradient magnitudes out of the underflow zone and into the representable dynamic range during backpropagation.^{[1]Reference 1Mixed Precision Training.https://arxiv.org/abs/1710.03740}

The mathematical principle relies directly on the linearity of the derivative. If you multiply the scalar loss $\mathcal{L}$ by a positive scale factor $S$ before backpropagation:

$\mathcal{L}_{\text{scaled}} = S \cdot \mathcal{L}$

By the chain rule, every gradient computed with respect to any parameter $w$ is scaled by that exact same factor:

$\nabla_w \mathcal{L}_{\text{scaled}} = \nabla_w (S \cdot \mathcal{L}) = S \cdot \nabla_w \mathcal{L}$

Multiplying by $S = 1024$ shifts magnitudes by $\log_2(1024) = 10$ binary exponent steps, or roughly $3.01$ decades. Our tiny $1.2 \times 10^{-8}$ gradient becomes $1.23 \times 10^{-5}$: representable as an FP16 subnormal, still below its minimum normal. This CPU fixture preserves it; a path that flushes subnormals would need a larger scale or another dtype.

Once the backward pass finishes, unscale the gradients by dividing by $S$ in FP32 before the optimizer updates this baseline's FP32 parameters:

$g_{\text{true}} = \frac{1}{S} \cdot g_{\text{scaled}}$

In exact arithmetic, scaling and unscaling cancel. Finite arithmetic introduces rounding error; the fixture below measures it rather than promising bit-identical updates.

Gradient clipping must always occur after unscaling. Suppose a parameter's true gradient norm is $0.02$, well within a healthy max_norm = 1.0 threshold. If you attempt to clip before unscaling, the clipping utility inspects the scaled gradient ($0.02 \times 1024 = 20.48$). Seeing $20.48 \gt 1.0$, it divides all gradients by $20.48$. After subsequent unscaling, your true gradient becomes $0.02 / 20.48 \approx 0.000977$, slashed by a factor of twenty due to a premature clipping check!

For a true gradient of $1.2 \times 10^{-8}$ and a teaching scale $S = 1024$, trace the lifecycle:

The scale must enter before the backward pass. Multiplying a gradient by $S$ after FP16 has already rounded it to zero can't recover the lost value; multiplying zero by 1024 is still zero.

The next test isolates the scale, quantize, and unscale cycle before we exercise dynamic backpropagation:

true_grad = torch.tensor([1.2e-8], dtype=torch.float32)
scale = 1024.0

plain_fp16 = true_grad.to(torch.float16)
scaled_fp16 = (true_grad * scale).to(torch.float16)
recovered_fp32 = scaled_fp16.to(torch.float32) / scale

print(f"plain_underflowed={plain_fp16.item() == 0.0}")
print(f"scaled_visible={scaled_fp16.item() > 0.0}")
print(f"scaled_is_normal={scaled_fp16.item() >= torch.finfo(torch.float16).tiny}")
print(f"recovered_grad={recovered_fp32.item():.2e}")
print(f"recovery_relative_error={abs(recovered_fp32.item() - true_grad.item()) / true_grad.item():.3%}")

plain_underflowed=True
scaled_visible=True
scaled_is_normal=False
recovered_grad=1.20e-08
recovery_relative_error=0.077%

Dynamic loss scaling as a feedback controller

A fixed scale can trade underflow for overflow. With $S = 65536$, a gradient of $1.0$ scales beyond FP16's maximum and rounds to infinity. A smaller scale reduces that risk but may lose smaller gradients. Which magnitudes survive is a property of the actual backward path.

The 2018 paper discussed dynamic scale selection as future work.^{[1]Reference 1Mixed Precision Training.https://arxiv.org/abs/1710.03740} Current PyTorch implements an adaptive controller. For one ordinary optimizer, its default policy works as follows:^{[8]Reference 8GradScaler implementation (PyTorch 2.14.0)https://github.com/pytorch/pytorch/blob/v2.14.0/torch/amp/grad_scaler.py}

1. Start with an aggressive scale factor $S = S_0$ (PyTorch defaults to $2^{16} = 65536.0$) and initialize a consecutive successful steps counter $c = 0$.
2. Compute forward pass under autocast, generating scalar loss $\mathcal{L}$.
3. Backpropagate on scaled loss $\mathcal{L}_{\text{scaled}} = S \cdot \mathcal{L}$. Intermediate gradients follow operation dtypes; FP32 parameters accumulate FP32 .grad buffers.
4. Unscale gradients in FP32: $g \leftarrow \frac{1}{S} g_{\text{scaled}}$.
5. Check the optimizer's gradients for non-finite values. PyTorch combines checking and unscaling by device and dtype. Its ordinary optimizer-step path reads flags with .item(), which can synchronize; AMP-aware optimizer paths can consume device flags directly.
6. Branch on finiteness:
   • If ANY gradient contains an Inf or NaN:
     • Skip the optimizer step completely. Don't update weights, AdamW first moments ($m$), or second moments ($v$).
     • Back off the scale: $S \leftarrow S \times \text{backoff\_factor}$ (default $0.5$, halving the scale).
     • Reset the success counter: $c \leftarrow 0$.
   • If ALL gradients are finite:
     • Clip gradient norms if configured (now safely operating on true unscaled magnitudes).
     • Execute optimizer.step(), updating the FP32 parameters in our AMP baseline.
     • Increment the success counter: $c \leftarrow c + 1$.
     • Check growth window: If $c == \text{growth\_interval}$ (default 2,000 clean steps), increase the scale: $S \leftarrow S \times \text{growth\_factor}$ (default $2.0$), and reset $c \leftarrow 0$.

For this one-optimizer contract, don't apply only its finite tensors. GradScaler skips that optimizer's step as a unit. Multiple optimizers are different: each makes an independent skip decision, followed by one shared scaler update.^{[9]Reference 9Automatic Mixed Precision Exampleshttps://docs.pytorch.org/docs/stable/notes/amp_examples.html} If your algorithm needs joint atomicity, implement that policy explicitly.

Try that distinction directly. We inject infinity into one optimizer's gradient to test skip behavior; this doesn't simulate a GPU overflow kernel.

import torch

good = torch.nn.Parameter(torch.tensor(1.0))
bad = torch.nn.Parameter(torch.tensor(1.0))
optimizers = [torch.optim.SGD([p], lr=0.1) for p in (good, bad)]
scaler = torch.amp.GradScaler("cpu", init_scale=128)
scaler.scale(good).backward()
scaler.scale(bad).backward()
bad.grad.fill_(float("inf"))  # authored fault injection
for optimizer in optimizers:
    scaler.step(optimizer)
scaler.update()  # once after both step decisions
assert good.item() < 1.0 and bad.item() == 1.0
print(f"good_stepped={good.item() < 1.0}; bad_stepped={bad.item() != 1.0}; "
      f"next_scale={scaler.get_scale():.0f}")

good_stepped=True; bad_stepped=False; next_scale=64

Here is a live CPU demonstration of this feedback loop in action:

model = torch.nn.Linear(1, 2, bias=False)  # parameters stay FP32
with torch.no_grad():
    model.weight.fill_(1.0)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
scaler = torch.amp.GradScaler("cpu", init_scale=1024, growth_interval=2)

for name, large_grad in (("quiet", 0.02), ("spike", 100.0),
                         ("quiet_again", 0.02), ("quiet_twice", 0.02)):
    optimizer.zero_grad(set_to_none=True)
    before = model.weight.detach().clone()
    state_before = {k: v.clone() for k, v in optimizer.state[model.weight].items()}
    with torch.autocast("cpu", dtype=torch.float16):
        prediction = model(torch.ones(1, 1))
        loss = (prediction.float() * torch.tensor([1.2e-8, large_grad])).sum()

    scaler.scale(loss).backward()
    scaler.unscale_(optimizer)
    gradient = model.weight.grad.clone()
    finite = bool(torch.isfinite(gradient).all())
    if finite:
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    scaler.step(optimizer)  # skips the entire Adam update if non-finite
    scaler.update()

    changed = not torch.equal(model.weight, before)
    state = optimizer.state[model.weight]
    if not finite:
        assert not changed
        assert all(torch.equal(state[k], value) for k, value in state_before.items())
    else:
        assert changed and gradient[0].item() > 0
    print(f"{name}: grad={gradient.flatten().tolist()}, changed={changed}, "
          f"adam_step={int(state['step'])}, next_scale={scaler.get_scale():.0f}")

quiet: grad=[1.1990778148174286e-08, 0.0200042724609375], changed=True, adam_step=1, next_scale=1024
spike: grad=[1.1990778148174286e-08, inf], changed=False, adam_step=1, next_scale=512
quiet_again: grad=[1.1990778148174286e-08, 0.0200042724609375], changed=True, adam_step=2, next_scale=512
quiet_twice: grad=[1.1990778148174286e-08, 0.0200042724609375], changed=True, adam_step=3, next_scale=1024

The first quiet step succeeds and updates weights. The spiky step detects inf after unscaling, completely skips the optimizer update, leaves Adam moments untouched, and halves the scale from 1024 to 512. Two subsequent quiet steps succeed, triggering the growth threshold (growth_interval=2) and restoring the scale to 1024.

In current PyTorch, torch.amp.GradScaler encapsulates this entire state machine.^{[2]Reference 2Automatic Mixed Precision package - torch.amphttps://docs.pytorch.org/docs/stable/amp.html} Its default starting scale is 65536.0, it doubles after 2000 consecutive successful steps, and it multiplies by 0.5 on any non-finite step. If you inspect or clip gradients, call scaler.unscale_(optimizer) first so your thresholds evaluate true gradient magnitudes.^{[2]Reference 2Automatic Mixed Precision package - torch.amphttps://docs.pytorch.org/docs/stable/amp.html}

Loss scaling can't repair a forward overflow. Casting a forward value beyond FP16's rounding range can produce infinity; scaling the later loss doesn't change that cast. Backoff below 1 can help a scaled backward overflow, but won't fix this forward cause. BF16 pretraining suggests checking range compatibility, not an inevitable fine-tuning collapse.^{[2]Reference 2Automatic Mixed Precision package - torch.amphttps://docs.pytorch.org/docs/stable/amp.html}

Check the loss separately from its gradients. GradScaler's ordinary step decision examines optimizer gradients, not whether the unscaled loss is finite.^{[8]Reference 8GradScaler implementation (PyTorch 2.14.0)https://github.com/pytorch/pytorch/blob/v2.14.0/torch/amp/grad_scaler.py} An infinite additive term can make the loss invalid without changing its parameter derivative. Predict whether the scaler skips this authored fault:

import torch

weight = torch.nn.Parameter(torch.tensor(1.0))
optimizer = torch.optim.SGD([weight], lr=0.1)
scaler = torch.amp.GradScaler("cpu", init_scale=128)
loss = weight.square() + torch.tensor(float("inf"))  # authored forward fault
scaler.scale(loss).backward()
scaler.unscale_(optimizer)
finite_loss = bool(torch.isfinite(loss))
finite_gradient = bool(torch.isfinite(weight.grad).all())
scaler.step(optimizer)
scaler.update()
assert not finite_loss and finite_gradient and weight.item() < 1.0
print(f"finite_loss={finite_loss}; finite_gradient={finite_gradient}")
print(f"scaler_stepped={weight.item() < 1.0}; next_scale={scaler.get_scale():.0f}")

finite_loss=False; finite_gradient=True
scaler_stepped=True; next_scale=128

The scaler steps because the unscaled gradient is 2.0. Its overflow controller has no evidence to reject that gradient. This fixture deliberately proceeds to expose the boundary; the training sketch below rejects a non-finite loss before backward. A finite loss and finite gradients still don't certify correct data or successful learning.

Compute low, update high

Loss scaling can protect small backward gradients. It doesn't resolve deterministic rounding of tiny parameter updates. Keep the FP32 storage policy for this baseline.

With an FP32 model, ordinary AMP does that without a separate master-weight tensor: autocast selects eligible compute dtypes, and the optimizer updates the original parameters.^{[2]Reference 2Automatic 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.^{[1]Reference 1Mixed 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. In this baseline, FP32 storage lets the demonstrated tiny changes accumulate. Other update policies need their own accuracy checks.

This is the CUDA shape for a real fine-tune. It isn't marked executable here because it needs an accelerator, model, criterion, and dataloader. Create the optimizer from FP32 model parameters already on the GPU. GradScaler(..., enabled=False) bypasses scaling and skip detection; it doesn't turn BF16 compute into FP32 compute. The disabled branch therefore needs its own non-finite-gradient policy:

dtype = torch.bfloat16  # compare against torch.float16 in a controlled run
if dtype == torch.bfloat16 and not torch.cuda.is_bf16_supported():
    raise RuntimeError("This CUDA device doesn't support the chosen BF16 path")
scaler = torch.amp.GradScaler("cuda", enabled=(dtype == torch.float16))

for batch, target in dataloader:
    optimizer.zero_grad(set_to_none=True)
    with torch.autocast(device_type="cuda", dtype=dtype):
        logits = model(batch.cuda())
        loss = criterion(logits, target.cuda())

    if not torch.isfinite(loss).all().item():
        raise FloatingPointError("Non-finite forward loss; optimizer not stepped")
    scaler.scale(loss).backward()
    scaler.unscale_(optimizer)
    finite = all(torch.isfinite(p.grad).all().item()
                 for p in model.parameters() if p.grad is not None)
    if not finite and not scaler.is_enabled():
        raise FloatingPointError("Non-finite BF16 gradients; optimizer not stepped")
    if finite:
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0, error_if_nonfinite=True)
    scaler.step(optimizer)
    scaler.update()

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. The per-parameter .item() checks above synchronize with the GPU; production implementations can aggregate checks, but shouldn't silently drop the safety decision.

Autocast is operation-specific, not a blanket cast. On CUDA, eligible matrix multiplications can use the selected low precision while operations such as cross_entropy autocast to FP32. Backward follows the dtypes used by the corresponding forward operations, even though an FP32 parameter's accumulated .grad buffer is FP32. In-place operations and operations with an explicit dtype don't follow the same autocast eligibility rules. Keep backward outside the autocast region and don't call model.half() to enable AMP.^{[2]Reference 2Automatic Mixed Precision package - torch.amphttps://docs.pytorch.org/docs/stable/amp.html}

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

Keep the same loss scale throughout the whole accumulation window. Unscale only once after every contribution has arrived, then step and update once. A non-finite contribution invalidates the whole optimizer step. Save and restore scaler.state_dict() along with the model and optimizer when resuming.^{[9]Reference 9Automatic Mixed Precision Exampleshttps://docs.pytorch.org/docs/stable/notes/amp_examples.html}

Dividing each mean loss by accum_steps is correct only when those micro-batches deserve equal weight. SFT often averages over valid response tokens, and their counts can differ. For one worker, sum each micro-batch's valid-token losses and divide by the whole window's valid-token count, then let GradScaler multiply that normalized loss. Don't divide the GradScaler scale by hand. For distributed training, also account for whether the collective sums or averages gradients.

This CPU autograd counterexample uses one token in the first micro-batch and three in the second. Predict why averaging two means overweights the first token. Squared error keeps the arithmetic visible; token cross-entropy has the same normalization issue.

import torch

targets = (torch.tensor([0.0]), torch.tensor([2.0, 2.0, 2.0]))
total_tokens = sum(y.numel() for y in targets)

def accumulated_gradient(mode):
    weight = torch.tensor(1.0, requires_grad=True)
    for y in targets:
        losses = (weight - y).square()
        loss = losses.mean() / len(targets) if mode == "mean_of_means" else losses.sum() / total_tokens
        loss.backward()
    return weight.grad.item()

weight = torch.tensor(1.0, requires_grad=True)
full_loss = (weight - torch.cat(targets)).square().mean()
full_loss.backward()
wrong = accumulated_gradient("mean_of_means")
correct = accumulated_gradient("token_weighted")
print(f"full_batch_grad={weight.grad.item():.1f}")
print(f"mean_of_means_grad={wrong:.1f}, token_weighted_grad={correct:.1f}")
assert correct == weight.grad.item() and wrong != correct

full_batch_grad=-1.0
mean_of_means_grad=0.0, token_weighted_grad=-1.0

For a short final window, use its actual contribution count rather than the planned window size. An all-masked window has zero valid tokens: reject or skip it before dividing. Advance a per-update learning-rate schedule according to your successful-update policy, not blindly for every skipped step.

Memory savings need accounting, not slogans

One of the most persistent myths in deep learning is the slogan: "Mixed precision cuts model memory in half."

End-to-end savings depend on which storage actually changes. Ordinary AMP with an FP32 model can leave persistent parameter and optimizer bytes unchanged while shrinking some saved tensors. Other recipes also change weights, gradients, optimizer state, or sharding. Doubling the feasible batch isn't guaranteed.

Start with three explicit inventories: unsharded AdamW, FP32 moment storage, one full-size gradient buffer, and no AMSGrad buffer. Scalar counters, workspaces, conversion caches, and allocator overhead are excluded. The activation row assumes the same set of values is stored at the chosen dtype; it isn't a description of every tensor autocast saves.

Notice the static tally:

1. Pure FP32: 4B weight + 4B gradient + two 4B moment buffers = 16 bytes per parameter. AdamW's second raw moment isn't the gradient variance.
2. FP32 model with AMP: The same parameter, gradient, and moment storage totals 16 bytes per parameter. Converted weight tensors may be cached for an autocast region or retained for backward, adding storage beyond this tally.
3. Master-copy inventory: Low-precision weights (2B) and gradients (2B), an FP32 master (4B), and FP32 moments (8B) total 16 bytes per parameter. An additional full-size FP32 gradient staging buffer makes this particular inventory 20 bytes. Actual framework copies and sharding determine the per-rank cost.

Autograd saves tensors needed for backward, not every intermediate. Some hidden-state storage grows like $bshl$ (batch, sequence length, hidden width, layers); materialized attention scores also have a quadratic sequence-length term. Fused attention and recomputation change that inventory.^{[10]Reference 10Reducing Activation Recomputation in Large Transformer Modelshttps://arxiv.org/abs/2205.05198} No fixed activation percentage follows from the model family.

Inspect a small case before assuming all saved values are 16-bit. This CPU hook records dtypes, not peak GPU memory.^{[11]Reference 11Autograd: saved_tensors_hookshttps://docs.pytorch.org/docs/2.14/autograd.html#torch.autograd.graph.saved_tensors_hooks} The explicit .float() before squaring creates an FP32 saved tensor alongside the BF16 linear-path input. The pack hook returns a detached tensor to avoid retaining a reference cycle with its input.

import torch

layer = torch.nn.Linear(4, 4)
saved_dtypes = set()
def record(tensor):
    saved_dtypes.add(str(tensor.dtype))
    return tensor.detach()

with torch.autograd.graph.saved_tensors_hooks(record, lambda tensor: tensor):
    with torch.autocast("cpu", dtype=torch.bfloat16):
        output = layer(torch.ones(2, 4))
        loss = output.float().square().mean()
    loss.backward()
assert saved_dtypes == {"torch.bfloat16", "torch.float32"}
print(f"saved_dtypes={sorted(saved_dtypes)}")

saved_dtypes=['torch.bfloat16', 'torch.float32']

Halving the bytes of an unchanged set of activation values halves that component. The next authored inventory assumes 100 million parameters and 800 million saved activation values, all of which we choose to cast. It isn't a profiled model:

def gib(values: int, bytes_per_value: int) -> float:
    return values * bytes_per_value / (1024 ** 3)

parameter_values = 100_000_000
activation_values = 800_000_000

fp32_budget = {
    "parameters": gib(parameter_values, 4),
    "gradients": gib(parameter_values, 4),
    "adam_moments": gib(parameter_values * 2, 4),
    "activations": gib(activation_values, 4),
}
amp_budget = {
    **{name: value for name, value in fp32_budget.items() if name != "activations"},
    "activations": gib(activation_values, 2),
}

print(f"fp32_total_gib={sum(fp32_budget.values()):.2f}")
print(f"amp_total_gib={sum(amp_budget.values()):.2f}")
print(f"activation_saving_gib={fp32_budget['activations'] - amp_budget['activations']:.2f}")
print(f"total_reduction={(1 - sum(amp_budget.values()) / sum(fp32_budget.values())):.1%}")

fp32_total_gib=4.47
amp_total_gib=2.98
activation_saving_gib=1.49
total_reduction=33.3%

The 2018 paper's roughly halved-memory discussion used FP16 weights, activations, and gradients plus an FP32 master, with activation-dominated memory.^{[1]Reference 1Mixed Precision Training.https://arxiv.org/abs/1710.03740} It isn't a promise for this AMP configuration or every Transformer workload.

Ordinary AMP, as in the budget above, leaves parameters and Adam moments in FP32. Those FP32 parameters, gradients, and two moments already total 16 bytes per parameter. A different 16-byte recipe stores low-precision parameter and gradient copies plus an FP32 master and FP32 Adam moments. Equal totals don't imply equal copies.

The inventory excludes temporary workspaces, cached autocast weight conversions, allocator overhead, and optimizer-specific buffers. Measure peak allocated and reserved GPU memory during representative steady-state steps before treating the budget as capacity planning.

For large models where model state exceeds single-GPU capacity, sharding frameworks like ZeRO and Fully Sharded Data Parallel (FSDP) can partition parameters, gradients, and optimizer moments across GPUs. Which states are partitioned depends on the chosen strategy or stage.^{[12]Reference 12ZeRO: Memory Optimizations Toward Training Trillion Parameter Models.https://arxiv.org/abs/1910.02054[13]Reference 13PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel.https://arxiv.org/abs/2304.11277} That addresses storage mixed-precision compute alone may leave unchanged.

Distributed jobs add a communication dtype

When workers exchange gradients, communication has its own precision policy. PyTorch's FullyShardedDataParallel (FSDP1) MixedPrecision configuration exposes param_dtype for forward and backward computation and reduce_dtype for gradient reduction. The two fields may differ; other distributed APIs have their own configuration types.^{[14]Reference 14FullyShardedDataParallelhttps://docs.pytorch.org/docs/stable/fsdp.html}

The next calculation is the size of one gradient tensor to reduce, not bytes sent across the network. Actual wire traffic also depends on the collective algorithm, worker count, sharding, and topology.

gradient_values = 100_000_000
fp32_reduce_gib = gib(gradient_values, 4)
bf16_reduce_gib = gib(gradient_values, 2)

print(f"gradient_payload_fp32_gib={fp32_reduce_gib:.2f}")
print(f"gradient_payload_bf16_gib={bf16_reduce_gib:.2f}")
print(f"payload_reduction={(1 - bf16_reduce_gib / fp32_reduce_gib):.0%}")

gradient_payload_fp32_gib=0.37
gradient_payload_bf16_gib=0.19
payload_reduction=50%

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.

FP32 tensor storage alone doesn't define the baseline's arithmetic. PyTorch can use TF32 for CUDA GEMMs or convolutions, and some FP16/BF16 GEMMs allow intermediate reductions at reduced precision. Record those controls too: reduced-precision partial sums can overflow even when the final mathematical result fits.^{[15]Reference 15Numerical Accuracyhttps://docs.pytorch.org/docs/stable/notes/numerical_accuracy.html}

That record supports reproducing a speed or memory result. For the authored precision comparison below, choose format_pass_rate >= 0.98 and policy_pass_rate >= 0.93 as validation gates. These lab thresholds aren't universal release standards. Measure time over warmed-up, synchronized GPU steps; asynchronous launch time alone isn't step latency.^{[16]Reference 16CUDA semanticshttps://docs.pytorch.org/docs/2.14/notes/cuda.html}

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.^{[17]Reference 17Defeating 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 two rounding boundaries. At the first, FP16 merges nearby logits while BF16 separates them. At the second, the reverse happens. A favorable example for one dtype isn't evidence that it always reduces mismatch.

import torch

def policy_gap(dtype: torch.dtype, first: float, second: float) -> float:
    training_logits = torch.tensor([first, 1.0], dtype=torch.float32)
    rollout_logits = torch.tensor([second, 1.0], dtype=torch.float32)
    training_cast = training_logits.to(dtype).to(torch.float32)
    rollout_cast = rollout_logits.to(dtype).to(torch.float32)
    training_probability = torch.softmax(training_cast, dim=0)[0]
    rollout_probability = torch.softmax(rollout_cast, dim=0)[0]
    gap = abs(training_probability.item() - rollout_probability.item())
    return gap

for first, second in ((1.0038, 1.0040), (1.0004, 1.0006)):
    gaps = [policy_gap(dtype, first, second) for dtype in (torch.float16, torch.bfloat16)]
    print(f"logits={first},{second}: FP16_gap={gaps[0]:.7f}, BF16_gap={gaps[1]:.7f}")

assert policy_gap(torch.float16, 1.0038, 1.0040) == 0
assert policy_gap(torch.bfloat16, 1.0038, 1.0040) > 0
assert policy_gap(torch.float16, 1.0004, 1.0006) > 0
assert policy_gap(torch.bfloat16, 1.0004, 1.0006) == 0

logits=1.0038,1.004: FP16_gap=0.0000000, BF16_gap=0.0019531
logits=1.0004,1.0006: FP16_gap=0.0002441, BF16_gap=0.0000000

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.

Apply both quality gates before comparing speed

The next fixture contains invented outcomes to test the comparison rule, not benchmark results. Assume model, data, seeds, software, batch shape, and evaluation are held fixed. One run changes hardware, another fails format validity, and another contains an invalid metric. Predict which rows must be excluded even if their throughput looks attractive.

from dataclasses import dataclass
from math import isfinite

@dataclass(frozen=True)
class PrecisionRun:
    run_id: str
    policy_pass_rate: float
    format_pass_rate: float
    nonfinite_steps: int
    examples_per_second: float
    hardware: str = "same-gpu"

runs = (
    PrecisionRun("fp32", 0.93, 0.99, 0, 800),
    PrecisionRun("fp16_unscaled", 0.88, 0.99, 3, 1240),
    PrecisionRun("fp16_scaled", 0.93, 0.99, 0, 1190),
    PrecisionRun("bf16", 0.93, 0.99, 0, 1310),
    PrecisionRun("other_hardware", 0.93, 0.99, 0, 1770, "different-gpu"),
    PrecisionRun("bad_format", 0.96, 0.97, 0, 1400),
    PrecisionRun("invalid_metric", float("nan"), 0.99, 0, 1500),
    PrecisionRun("invalid_throughput", 0.93, 0.99, 0, float("inf")),
)

def eligible(run):
    rates = (run.policy_pass_rate, run.format_pass_rate)
    return (run.hardware == "same-gpu"
            and all(isfinite(rate) and 0 <= rate <= 1 for rate in rates)
            and run.policy_pass_rate >= 0.93 and run.format_pass_rate >= 0.98
            and run.nonfinite_steps == 0
            and isfinite(run.examples_per_second) and run.examples_per_second > 0)

accepted = [run.run_id for run in runs if eligible(run)]
rejected = [run.run_id for run in runs if not eligible(run)]
assert accepted == ["fp32", "fp16_scaled", "bf16"]
print(f"quality_and_comparability_pass={accepted}")
print(f"rejected={rejected}")

quality_and_comparability_pass=['fp32', 'fp16_scaled', 'bf16']
rejected=['fp16_unscaled', 'other_hardware', 'bad_format', 'invalid_metric', 'invalid_throughput']

Both scaled FP16 and BF16 pass this toy filter. Neither result establishes a performance win: the numbers were invented. On real runs, compare repeated throughput and peak-memory measurements only among eligible candidates, then check loss curves and held-out slices across seeds. A few early FP16 scaler backoffs can be normal; the zero-nonfinite rule here is a deliberately strict fixture gate, not a universal convergence criterion.

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.

The OCP FP8 specification defines 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. Other formats with similar names can use different special-value rules.^{[18]Reference 18OCP 8-bit Floating Point Specification (OFP8), Revision 1.0https://www.opencompute.org/documents/ocp-8-bit-floating-point-specification-ofp8-revision-1-0-2023-12-01-pdf-1}

E4M3 buys nearby resolution; E5M2 buys range. The original FP8 formats paper motivates using them for different tensors rather than treating FP8 as a single format.^{[19]Reference 19FP8 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.^{[20]Reference 20Using FP8 and FP4 with Transformer Enginehttps://docs.nvidia.com/deeplearning/transformer-engine/examples/fp8_primer.html} The API and recipe documentation here were checked against version 2.20.2 on October 10, 2026; pin the implementation used by an actual experiment.

Its current-scaling recipe computes an absolute maximum (amax) from the current tensor; delayed scaling estimates that scale from amax history. Both use tensor-wide FP32 scaling factors. Current scaling needs the current amax before quantization, which can require extra passes unless fused with other operations; delayed scaling instead risks history that lags a new spike.^{[20]Reference 20Using FP8 and FP4 with Transformer Enginehttps://docs.nvidia.com/deeplearning/transformer-engine/examples/fp8_primer.html}

MXFP8 changes the scale granularity. OCP MX defines MXFP8 as FP8 elements with one E8M0 power-of-two scale shared by each block of 32 values.^{[21]Reference 21OCP Microscaling Formats (MX) Specification, Version 1.0https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf}

Transformer Engine's MXFP8 recipe uses E4M3 by default and also accepts a hybrid E4M3-forward/E5M2-backward configuration.^{[22]Reference 22Transformer Engine recipe APIhttps://docs.nvidia.com/deeplearning/transformer-engine/api/common.html} Its native hardware path was introduced with Blackwell.^{[20]Reference 20Using FP8 and FP4 with Transformer Enginehttps://docs.nvidia.com/deeplearning/transformer-engine/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.^{[19]Reference 19FP8 Formats for Deep Learning.https://arxiv.org/abs/2209.05433[21]Reference 21OCP Microscaling Formats (MX) Specification, Version 1.0https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf}

OCP MX leaves internal precision and operation order implementation-defined, so the stored dtype alone doesn't specify the numerical result of a kernel.^{[21]Reference 21OCP Microscaling Formats (MX) Specification, Version 1.0https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf}

OCP FP8 also distinguishes saturating from non-saturating conversion: a positive value beyond the format's rounding range can clamp to 448 or become NaN in E4M3, 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.^{[18]Reference 18OCP 8-bit Floating Point Specification (OFP8), Revision 1.0https://www.opencompute.org/documents/ocp-8-bit-floating-point-specification-ofp8-revision-1-0-2023-12-01-pdf-1}

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.

Check your understanding

Evaluation rubric

• Foundational: Distinguish a tiny standalone value from a tiny update to a weight near 1.0, using FP16 and BF16 range and spacing.
• Intermediate: Trace scale, backward, unscale, finite check, clipping, step or skip, and scale update. Explain why an overflowed Adam step must leave weights and moments unchanged.
• Intermediate: Normalize unequal-token micro-batches by the full effective batch and keep one loss scale through the accumulation window.
• Advanced: Separate activation, parameter, optimizer, and communication bytes, then compare measured precision candidates without changing hardware or quality criteria.

Follow-up questions