21. Mixed-Precision-Safe Manual Loss Scaling

Problem

In fp16 mixed-precision training, gradients that are small in magnitude (and fp16 has far less dynamic range near zero than fp32) can silently underflow to zero during backward, stalling learning for whatever those gradients were supposed to update. The standard fix, loss scaling, is what torch.cuda.amp.GradScaler automates: multiply the loss by a large scale factor before calling .backward() (which scales every gradient by the same factor via the chain rule, pushing small gradients back into fp16's representable range), then divide the gradients back down by that same factor before the optimizer step.

Scaling up can also push some gradients to overflow into inf/nan, so a dynamic scaler checks for that after every step: skip the optimizer step and shrink the scale if it finds inf/nan, otherwise grow the scale so training spends as much time as possible at a scale large enough to protect small gradients.

Implement grad_scaler_step: given already-backward()'d gradients that were computed from a loss scaled by scale, unscale them, check whether any are non-finite, and return either the correctly unscaled gradients (and a grown scale) or None (and a shrunk scale) if the step must be skipped.

Source: src/21_manual_loss_scaling.py

def grad_scaler_step(
    grads: list[torch.Tensor], scale: float, growth_factor: float = 2.0,
    backoff_factor: float = 0.5,
) -> tuple[list[torch.Tensor] | None, float]: ...

>>> grads = [torch.tensor([2048.0, -4096.0])]
>>> unscaled, new_scale = grad_scaler_step(grads, scale=1024.0)
>>> unscaled[0]
tensor([ 2., -4.])
>>> new_scale
2048.0

Step-by-Step Approach

1. Unscale every gradient first: g / scale for each tensor in grads.
2. Check finiteness after unscaling, not before — a real overflow can still look like a large-but-finite number at the scaled magnitude, and only becomes visibly inf/nan once divided back down (or was already inf/nan going in, which unscaling preserves).
3. If any unscaled gradient contains inf or nan, return (None, scale * backoff_factor) — None signals the caller to skip the optimizer step entirely for this iteration.
4. If every gradient is finite, return (unscaled_grads, scale * growth_factor).
5. Note the asymmetry: a skipped step still discards the wasted forward/backward's gradients but shrinks scale for next time; a successful step both uses its gradients and grows scale, so well-behaved training spends most iterations at (or climbing toward) the largest safe scale.

The key insight is the order of operations: checking for inf/nan on the scaled gradients would miss real overflows that only manifest once the artificially-large scale factor is divided back out, so unscaling always has to happen before the finiteness check, not after.

Reference solution

import torch


def grad_scaler_step(
    grads: list[torch.Tensor], scale: float, growth_factor: float = 2.0,
    backoff_factor: float = 0.5,
) -> tuple[list[torch.Tensor] | None, float]:
    # unscale first: inf/nan must be judged on the *true* gradient
    # magnitude, not the scaled one
    unscaled = [g / scale for g in grads]
    found_inf = any(not torch.isfinite(g).all() for g in unscaled)

    if found_inf:
        # skip this optimizer step entirely and shrink scale so the next
        # attempt is less likely to overflow again
        return None, scale * backoff_factor

    # every step that doesn't overflow grows the scale, so training spends
    # as much time as possible at a scale large enough to protect small
    # fp16 gradients from flushing to zero
    return unscaled, scale * growth_factor

Key Functions & Tricks

• torch.isfinite(g).all() — the standard finiteness check; catches both inf (overflow) and nan (typically inf - inf or 0/0 arising downstream of an earlier overflow) in one call.
• Unscale-then-check ordering — the one-line detail that makes the difference between a scaler that actually catches overflow and one that misses it.
• None as a skip signal — matches GradScaler.step(optimizer)'s real behavior, which silently no-ops the wrapped optimizer.step() call when it detects inf/nan, rather than raising.
• Growth/backoff as plain multiplicative factors — this problem implements the single-step core of GradScaler's algorithm; the real implementation additionally only grows the scale every growth_interval consecutive successful steps rather than every single one.

How to Recognize This Pattern