Prevent FP16 Gradient Underflow with Dynamic Loss Scaling

Mixed-precision training can reduce memory use and accelerate supported operations, but float16 introduces a numerical problem that is easy to miss: some gradients are too small to survive in FP16. They can round to zero before the optimizer gets a chance to use them.

Loss scaling addresses that problem by multiplying the loss before backpropagation, which multiplies the resulting gradients by the same factor. The gradients are divided by that factor before the optimizer update, so the intended update is unchanged when the arithmetic remains finite. Dynamic loss scaling adjusts the factor during training so you don’t have to guess one fixed value for the whole run.

This article builds the mechanism from a small example, then explains overflow handling, gradient clipping, accumulation, and the cases where loss scaling is unnecessary or cannot fix the underlying problem.

The problem is representability, not a weak learning signal

Floating-point formats can represent only a finite set of numbers. FP16 has a much narrower exponent range than FP32, so sufficiently small non-zero values cannot be represented as ordinary FP16 numbers. Values near the bottom of the range may become subnormal values with reduced precision, and still smaller values round to zero.

That matters during backpropagation because neural networks often produce gradients spanning many orders of magnitude. A gradient can be mathematically valid and useful while still being too small for the numeric format used to store or compute it.

Consider a simplified parameter whose true gradient is:

g = 0.00000003

Suppose that value becomes zero in the FP16 path used by a particular training computation:

true gradient:       0.00000003
stored FP16 gradient: 0

The optimizer cannot recover information that has already rounded to zero. Increasing the learning rate later does not solve this specific failure: zero multiplied by a larger learning rate is still zero.

This is different from a genuinely small gradient caused by the model or objective. Loss scaling does not decide that small gradients deserve larger parameter updates. It temporarily moves gradient values into a numerically safer range and then reverses that scaling before the optimizer step.

Scaling the loss scales every gradient

Let the original loss be L, the parameters be theta, and the loss scale be a positive constant S. Instead of differentiating L, compute:

L_scaled = S * L

By linearity of differentiation:

dL_scaled / dtheta = S * dL / dtheta

If the original gradient is g, backpropagation therefore produces S * g.

Take the earlier toy gradient and use a scale of 1024:

g          = 0.00000003
S          = 1024
S * g      = 0.00003072

The scaled value is much farther from zero. Before the optimizer uses it, divide by the same scale:

0.00003072 / 1024 = 0.00000003

Conceptually, the parameter update is the same as if the original gradient had been represented accurately from the start.

The sequence is:

ordinary loss
    -> multiply by S
    -> backward pass produces scaled gradients
    -> divide gradients by S
    -> optimizer uses unscaled gradients

The scale is a numerical device. It is not supposed to change the optimization objective.

Why one fixed scale is awkward

A large scale protects more small gradients from underflow, but it also pushes large gradients upward. If the scale is too large, an intermediate value or gradient can overflow to a non-finite value such as infinity.

A fixed scale therefore creates a tuning problem:

scale too small -> some useful small gradients may still underflow
scale too large -> scaled gradients may overflow

The useful range can also change during training. A scale that works near initialization may be unnecessarily conservative later, or a scale that works for many steps may overflow on an unusually large batch.

Dynamic loss scaling turns this into feedback rather than a one-time guess.

Dynamic loss scaling adapts to overflow

The exact policy is framework-specific, but the general mechanism is straightforward.

Start with a scale S. For each optimizer update:

1. multiply the loss by S;
2. run backpropagation;
3. unscale the gradients;
4. check whether the relevant gradients are finite;
5. perform the optimizer step only when they are finite;
6. update S according to the scaling policy.

A simplified policy might look like this:

if gradients contain inf or NaN:
    skip optimizer update
    reduce scale
else:
    apply optimizer update
    after enough successful updates, increase scale

Skipping an update after detected overflow matters. Once a scaled gradient has become infinity, dividing it by the scale does not reconstruct the finite gradient that should have existed. Updating parameters from corrupted gradients would defeat the purpose of the check.

Increasing the scale after a run of finite updates explores whether more of FP16’s available range can be used for small gradients. Reducing it after overflow moves the computation back toward a safer range.

This feedback loop is why the method is called dynamic loss scaling.

A small training-loop example

A typical mixed-precision training loop separates two concerns: autocasting chooses lower or higher precision for supported operations, while the gradient scaler manages loss scaling around backpropagation and the optimizer step.

In current PyTorch, a simplified CUDA example looks like this:

scaler = torch.amp.GradScaler("cuda")

for inputs, targets in loader:
    optimizer.zero_grad()

    with torch.autocast(device_type="cuda", dtype=torch.float16):
        predictions = model(inputs)
        loss = loss_fn(predictions, targets)

    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

The important part is the order, not the library syntax. scale(loss).backward() creates scaled gradients. scaler.step(optimizer) handles unscaling and checks for non-finite gradients before allowing the optimizer step. scaler.update() then adjusts the scale according to the scaler’s policy.

This example is intentionally minimal. Real training loops may also use gradient clipping, accumulation, multiple optimizers, distributed training, or custom gradient computations. Those features make the order of operations more important.

Unscale before inspecting or clipping gradients

Suppose the true gradient vector has norm 0.5, the loss scale is 1024, and you want to clip gradients to a maximum norm of 1.0.

Before unscaling, the stored gradient norm is conceptually:

0.5 * 1024 = 512

If you clip that scaled gradient against 1.0, the clipping code sees a norm of 512, not the true norm of 0.5. You would clip a gradient that should not have been clipped.

The correct conceptual order is:

scaled backward
    -> unscale gradients
    -> clip using the intended threshold
    -> optimizer step

With PyTorch’s scaler, that can be written as:

scaler.scale(loss).backward()
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
scaler.step(optimizer)
scaler.update()

The same rule applies to most logic that interprets gradient magnitudes. If you log gradient norms, apply custom clipping, or inspect .grad values for debugging, know whether those values are currently scaled.

Keep one scale across a gradient-accumulation window

Gradient accumulation builds one effective update from several microbatches. Scaling introduces an extra invariant: gradients that are being added together must use a compatible scale.

Suppose two microbatches produce gradients g1 and g2. If both use the same scale S, the accumulated scaled gradient is:

S * g1 + S * g2 = S * (g1 + g2)

One division by S recovers the intended accumulated gradient.

If the scale changes between microbatches, you instead get:

S1 * g1 + S2 * g2

There is no single scale you can divide by to recover g1 + g2 in general.

For that reason, keep gradients scaled with the same factor while building one effective batch. Unscale, step, and update the scale only when the accumulation window is complete. If your training framework manages this automatically, its documented accumulation pattern should determine where scale updates occur.

Loss scaling does not make every low-precision problem disappear

Loss scaling solves a specific problem: gradients becoming too small for an FP16 computation path. Several related failures need different fixes.

Forward-pass overflow is a different problem

Scaling happens around the loss and backward pass. It does not repair activations that already overflowed during the forward pass. If a model produces non-finite activations before the loss is scaled, investigate the operations, input range, initialization, normalization, precision choices, and other sources of instability.

A bad optimization setup remains bad

A learning rate that is too aggressive can still destabilize training after gradients are correctly unscaled. So can a broken loss function, malformed data, or an invalid custom operation. A scaler may detect non-finite gradients and skip steps, but repeated scale reductions are a symptom worth investigating, not proof that the scaler is fixing the root cause.

Scaling cannot recover values lost before scaling takes effect

The useful multiplication must affect gradient generation before the problematic low-precision values are rounded away. Multiplying already-underflowed gradients after backpropagation would only multiply zeros.

This is why loss scaling is applied to the loss before backward(), rather than as a post-processing step on finished gradients.

FP16 and bfloat16 do not have the same need for scaling

float16 and bfloat16 are both 16-bit formats, but they allocate their bits differently. FP16 provides more fraction precision, while bfloat16 uses an exponent width comparable to FP32 and therefore has a much wider dynamic range than FP16.

That wider exponent range makes gradient underflow from limited range much less of a concern for bfloat16 training. As a result, loss scaling is commonly associated with FP16 mixed-precision training rather than being a universal requirement for every 16-bit format.

This does not mean bfloat16 is numerically identical to FP32 or automatically appropriate for every operation. It means the specific FP16 range problem that motivates loss scaling is substantially reduced.

Likewise, using mixed precision does not imply that every tensor in a training system is stored in the same low-precision format. Frameworks and optimizers may keep selected operations, accumulations, parameters, or optimizer state in higher precision. Treat the documented behavior of your stack as the source of truth rather than assuming one dtype applies everywhere.

Diagnose scaling problems from behavior, not one scale value

The absolute scale value is not a model-quality metric. A larger scale does not mean the model is learning better, and a smaller scale does not by itself mean training is broken.

More useful signals are the behavior around the scale:

• Does the training loss follow a plausible trajectory?
• Are optimizer steps being skipped frequently because of non-finite gradients?
• Do non-finite values already appear in forward activations or the unscaled loss?
• Did instability begin after a change to the learning rate, data, model, optimizer, or custom kernels?
• Are gradient operations such as clipping happening after unscaling?
• With accumulation, is the scale kept constant until the effective update is ready?

An occasional skipped update can be part of a dynamic scaler finding a workable range. Persistent overflow is different. If the scale repeatedly falls while training remains non-finite, look beyond the scaler.

When dynamic loss scaling is worth using

Dynamic loss scaling is a good fit when training uses FP16 for backward computations and the framework supports a well-tested scaler. It removes most of the manual work of selecting a fixed scale and adapts when gradient magnitudes change over the run.

A fixed scale can still make sense in a tightly controlled setup where its safe range has been validated, but it shifts responsibility to you: the scale must be large enough to protect useful small gradients without causing damaging overflow across representative training conditions.

If training uses FP32 throughout, loss scaling usually adds no benefit for FP16 underflow because FP16 is not involved. If training uses bfloat16, first check whether your framework and hardware recommend a scaler for that path; don’t carry over an FP16 recipe automatically.

There is also a simpler engineering rule: if mixed precision does not provide a meaningful memory or throughput benefit for your workload, full-precision training may be preferable to adding numerical machinery you don’t need.

Use the scaler as a numerical adapter

The cleanest mental model is to treat dynamic loss scaling as an adapter between the gradient magnitudes produced by the training objective and the numeric range available to FP16 backpropagation.

It temporarily magnifies the loss so small gradients have a better chance of surviving, reverses that magnification before the optimizer update, and backs off when the scaled computation overflows. Done correctly, it protects numerical information without intentionally changing the optimization objective.

When adding it to a training pipeline, verify the operation order first: scale before backward, keep a consistent scale across one accumulation window, unscale before magnitude-dependent gradient operations, skip corrupted updates, and change the scale only at the appropriate update boundary. That sequence is more important than any particular initial scale value.