What is gradient loss scaling in mixed-precision training?

Gradient loss scaling rescues tiny FP16 backpropagation gradients that would otherwise round to zero. Multiply the loss by a factor S before backpropagation, which also multiplies derivatives by S. Then divide the parameter gradients by S before clipping and updating parameters. A dynamic gradient scaler additionally detects Inf/NaN gradients, skips unsafe optimizer updates and adjusts the scale.

This technique is based on Micikevicius et al., *Mixed Precision Training* (ICLR 2018). It is distinct from the general AMP workflow in Question 293 and the FP16-versus-BF16 format comparison in Question 294.

1. Why does FP16 loss scaling help?

• Small gradients disappear: FP16's limited exponent range can make a mathematically nonzero gradient or intermediate derivative become zero during low-precision backward calculations.
• Scaling loss early helps: The chain rule carries that scale throughout the backward computation, allowing some small derivatives to remain representable at the point where they are computed.
• Unscaling preserves the intended learning rate: Undo the multiplication in the resulting gradients before applying an optimizer step.
• Scaling an already-rounded zero cannot help: The loss must be scaled *before backward*, not after a tiny derivative was lost.
• The final `.grad` dtype is not the whole story: An ordinary PyTorch AMP model may keep FP32 parameters and FP32 parameter-gradient buffers, yet an intermediate backward operation may still run in FP16 and lose small derivatives there.

2. Basic equations

Let the original training loss be L(θ) and an ordinary gradient be g = ∂L/∂θ, for trainable parameters θ.

• Scaled loss: L_scaled = S × L(θ), with S > 0.
• Scaled gradient by chain rule: g_scaled = ∂(S×L)/∂θ = S × g.
• Restored gradient: g_restored = g_scaled / S.
• Normal SGD update: θ_new = θ_old − lr × g_restored.

In exact arithmetic, g_restored = g, so choosing S does not change the mathematical update. In real FP16 arithmetic, the scaled gradient is rounded but may avoid underflow.

What scaling cannot fix: Nonfinite values caused by a genuinely unstable model or operations overflowing in the forward pass may remain nonfinite regardless of loss scaling. The method does not cure every numerical problem.

3. Beginner-friendly numerical example: Saving a tiny FP16 gradient

Consider a toy FP32 parameter w_old = 0, a finite loss L = 0.125, and its true derivative with respect to this parameter:

g = 1 × 10⁻⁸.

Assume a relevant backward intermediate would be stored or computed as FP16 with round-to-nearest conversion. We isolate that step for teaching; not all PyTorch parameter gradients are FP16.

Step 1 — See the unscaled underflow.

• FP16's smallest theoretical positive subnormal is approximately 5.96046 × 10⁻⁸.
• Our 1 × 10⁻⁸ derivative is below half that value.
• Direct FP16 rounding gives `FP16(1 × 10⁻⁸) = 0`.
• If this were the only gradient contribution, the resulting parameter update would be zero.

Step 2 — Scale the loss before backward.

Choose S = 16,384 = 2¹⁴.

• L_scaled = 16,384 × 0.125 = 2,048.
• The backward derivative grows to g_scaled = 16,384 × 10⁻⁸ = 0.00016384.
• This scaled value exceeds FP16's minimum positive normal value, about 0.0000610352, so it is representable as a finite nonzero FP16 normal number.

Step 3 — Check the actual FP16 rounding of the scaled derivative.

In a conventional IEEE FP16 round-to-nearest conversion:

FP16(0.00016384) ≈ 0.00016379356384277344.

The value is not exact, but it did not become zero.

Step 4 — Unscale in a suitable FP32 update path.

g_restored = 0.00016379356384277344 / 16,384 ≈ 9.99716576 × 10⁻⁹.

• The ideal gradient is 1.0 × 10⁻⁸.
• Its recovered relative error is approximately 0.0283%.
• The meaningful small derivative has survived the FP16 intermediate calculation.

Step 5 — Verify the optimizer update.

Use illustrative SGD with lr = 0.1, FP32 w_old = 0:

• With scaling/unscaling: w_new = 0 − 0.1 × 9.99716576×10⁻⁹ ≈ −9.99716576×10⁻¹⁰.
• Ideal exact update: 0 − 0.1 × 10⁻⁸ = −1.0×10⁻⁹.
• If an unscaled FP16 intermediate had become zero: w_new = 0.

Interpretation: Scaling is not multiplying the effective learning rate by 16,384. The gradient is returned to its original scale before applying the update.

This toy calculation is a scalar illustration. Real backpropagation mixes data types, reduction order and kernel behavior; using FP32 accumulation or a scaler does not guarantee perfectly recovered gradients in all models.

4. Static loss scaling versus dynamic loss scaling

Static scaling chooses one fixed S for every training step.

• Benefit: Simple, reproducible and sometimes adequate when gradient magnitudes are predictable.
• Limitation: A value large enough to preserve tiny gradients may overflow elsewhere; the right scale can change throughout training.

Dynamic scaling changes S based on whether relevant gradients are finite:

1. Choose an initial scale S.
2. Scale the loss and run backward.
3. Check parameter gradients for Inf or NaN during unscale/step processing.
4. If all are finite: apply the optimizer update; count the successful step.
5. If any are nonfinite: skip that optimizer's update so invalid gradients do not corrupt parameters.
6. After an overflow: multiply S by a backoff_factor, commonly 0.5, and reset the success counter.
7. After enough consecutive successful steps: multiply S by a growth_factor, commonly 2, once the growth_interval threshold is reached.

Important: Dynamic scaling is a feedback mechanism, not an oracle that guarantees no overflow in intermediate computations. PyTorch also does not guarantee the scale stays above 1.

5. Numerical example: Overflow, skipped update and backoff

This is a separate toy iteration from the tiny-gradient example.

• Suppose a gradient g_big = 10 occurs in an FP16 backward path.
• The current scale is S = 8,192.
• Scaled value: 10 × 8,192 = 81,920.
• FP16's maximum finite value is 65,504, so this conversion can overflow to +Inf.

Correct dynamic behavior:

1. The gradient finiteness check detects the nonfinite result.
2. Skip the optimizer step. If w_old=1.0, it remains 1.0 on this failed update.
3. Apply backoff_factor=0.5: S_new = 8,192 × 0.5 = 4,096.
4. On a later attempt, if a comparable gradient value occurs, 10 × 4,096 = 40,960, which fits FP16.
5. If all checked gradients are finite, the optimizer may step using the unscaled gradient. Under illustrative SGD with lr=0.01, w_new = 1−0.01×10=0.9.

The last update is hypothetical: a new minibatch or changed parameters could produce a different true gradient. Skipping an unsafe update is preferable to writing Inf/NaN into model weights.

Scale growth illustration: Starting at S=4,096, with growth_factor=2 and an illustrative growth_interval=2, two consecutive successful optimizer steps can raise the scale to 8,192. If that causes another overflow, the scaler backs off again. Real configurations often use longer growth intervals.

6. The correct PyTorch GradScaler order

This representative CUDA FP16 code uses the documented modern PyTorch API; model, criterion and train_loader must be defined.

PYTHON
import torch

model = model.cuda()  # Keep default FP32 parameters
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
scaler = torch.amp.GradScaler(
    "cuda",
    init_scale=2.0**14,
    growth_factor=2.0,
    backoff_factor=0.5,
    growth_interval=2000,
)

for inputs, targets in train_loader:
    inputs, targets = inputs.cuda(), targets.cuda()
    optimizer.zero_grad(set_to_none=True)

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

    # Scale before backward (backward outside autocast).
    scaler.scale(loss).backward()

    # Explicitly unscale before clipping or inspecting gradients.
    scaler.unscale_(optimizer)
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

    # Only calls optimizer.step() when gradients are finite.
    scaler.step(optimizer)
    scaler.update()

What the key functions do:

• scaler.scale(loss) multiplies the loss by the current scale.
• .backward() propagates scaled derivatives through the graph.
• scaler.unscale_(optimizer) divides this optimizer's gradients by S and checks them for nonfinite values.
• clip_grad_norm_ then sees original-scale gradients.
• scaler.step(optimizer) applies the update only if gradients are finite. Since explicit unscale_ was already called, it does not unscale this optimizer a second time.
• scaler.update() adjusts the scale based on the results.

Do not call `unscale_` twice for the same optimizer in one step. If no clipping or gradient inspection is needed, omit explicit unscale_; scaler.step(optimizer) performs that work automatically. Calling bare optimizer.step() with still-scaled gradients bypasses safety checks and gives incorrect updates.

The older torch.cuda.amp.GradScaler API is deprecated in favor of torch.amp.GradScaler("cuda").

7. Why unscale before gradient clipping?

Clipping should enforce a limit on the true gradient norm, not on its artificially amplified value.

Suppose g=[3,4] with L2 norm 5 and a desired maximum norm of 1.

Correct:

• Unscaled norm sqrt(3²+4²)=5.
• Clip factor 1/5=0.2.
• Final gradient [3,4]×0.2=[0.6,0.8], with norm 1.

Incorrect: Clip before undoing the loss scale.

• For S=16, gradients are [48,64], with norm 80.
• Clipping these directly to norm 1 yields [0.6,0.8].
• Subsequent unscaling divides by 16 again: `[0.0375,0.05]`.
• Its resulting norm is only 0.0625, much less than the intended maximum when clipping the original gradient.

Rule: Finish backward, call scaler.unscale_(optimizer) once, then clip, then call scaler.step(optimizer) and scaler.update().

8. How does loss scaling interact with gradient accumulation?

When accumulating gradients across multiple microbatches, one effective batch must use a consistent loss scale from beginning to end.

Example: Two equal-sized microbatches with original per-microbatch mean gradients g₁=1×10⁻⁸ and g₂=2×10⁻⁸.

• For an equally weighted aggregate mean, divide each microbatch's loss by 2.
• Use one fixed scale S for both backwards.
• The accumulated scaled gradient is S×(g₁+g₂)/2.
• After both backward calls, unscale once to recover the intended aggregate (1×10⁻⁸+2×10⁻⁸)/2 = 1.5×10⁻⁸ (up to finite-precision rounding).
• Clip if desired; make one optimizer step, call scaler.update() once, then zero gradients before the next effective batch.

Why this matters: Unscaling or changing S after only the first microbatch mixes differently scaled gradients in the same buffers; they can no longer be correctly recovered by one division.

• Do not call unscale_ until all microbatches belonging to an optimizer update have been backpropagated.
• Do not call update() inside the accumulation window.
• Apply scale growth/backoff and Inf/NaN skip decisions at effective-batch granularity.
• If microbatches differ in size and the objective is a sample-wise mean, weight their losses by the actual counts rather than blindly averaging two unequal means.

9. BF16: Does it need loss scaling too?

• BF16's much wider exponent range makes FP16-style tiny-gradient underflow less common for ordinary neural-network values.
• Consequently, BF16 AMP training often does not require a gradient scaler, unlike some FP16 workloads.
• A conditional setup can use torch.amp.GradScaler("cuda", enabled=(amp_dtype == torch.float16)) when the rest of the precision policy supports this decision.
• BF16 is not immune to instability: Its coarser rounding, numerically fragile operations, bad data and excessively large updates can still harm convergence.
• Specialized or distributed optimizers may apply separate scaling policies; follow their documented conventions.

The FP16 and BF16 bit layouts and representable numbers were the subject of Question 294, so this section focuses only on scaling necessity.

10. Debugging and operational considerations

• Check skipped updates: Some skipped steps during initial dynamic calibration can be expected, but frequent skipping can slow or stall progress.
• Monitor the scale: Use scaler.get_scale() and instrument actual optimizer-step counts. Do not treat an optimizer return value of None as proof of a skipped step: many optimizers return None even when they step normally.
• Look for nonfinite gradients: If instability persists at very low scales, examine the model, data normalization, loss function and problematic operations.
• No after-the-fact repair: Once an FP16 intermediate was rounded to zero, increasing a downstream parameter gradient cannot reconstruct it.
• Multiple optimizers: A scaler can decide whether to skip each optimizer's step based on its own gradients; call scaler.step for each and scaler.update after all have been handled.
• Save state when resuming: Checkpoint the scaler state together with the model and optimizer to reproduce the dynamic scaling trajectory.
• Avoid double unscaling: Explicit unscale_ is once per optimizer per optimizer step.
• Unscale before inspecting gradients: Clipping, manual gradient modifications and regular diagnostic norms should use original-scale gradients.
• Do not treat loss scaling as a speed boost: It preserves numerical information while allowing an FP16 training path; kernel throughput and memory benefits come from lower-precision operations themselves.

11. Applications and limitations

• Object detection and vision: Loss scaling historically helped some FP16 object detectors whose small backward values otherwise became zero.
• Large-scale sequence models: FP16 transformers and recurrent networks can require scaling under particular activation and gradient distributions.
• Generative systems: FP16 training of GANs and diffusion networks may benefit if useful derivatives are lost to underflow.
• Hardware sensitivity: Rounding, subnormal support and intermediate dtypes depend on the accelerator and kernels; verify behavior instead of assuming every gradient is actually stored in FP16.
• Not a universal fix: A single global scale may preserve some derivatives but overflow others; it cannot repair unrelated optimizer or model-design failures.

Curriculum boundary: Question 293 introduced AMP and general training memory. Question 294 compared FP16 and BF16 numeric formats. Question 295 is specifically about loss-scaling, detecting unsafe gradients and preserving optimizer semantics. Question 296 will cover activation checkpointing to save GPU memory, a separate efficiency technique.

Key takeaway: Scale the loss before backpropagation, unscale gradients before clipping and weight updates, skip the step if scaled gradients become nonfinite, and use one stable scale per effective batch. This preserves useful tiny FP16 gradients without changing the intended optimization step.

Research and documentation: Micikevicius et al., *Mixed Precision Training* (ICLR 2018); PyTorch, *Automatic Mixed Precision Examples*; PyTorch, *AMP/GradScaler API*; NVIDIA, *Training with Mixed Precision*.