← All Problems

47. Manual Mixed-Precision Loss Scaling (GradScaler-Style)

General Hard Training Mechanics & Engineering Tradeoffs
Grounding: General industry practice — this is the well-documented internal logic of PyTorch's own torch.cuda.amp.GradScaler (scale → check for inf/NaN after backward → skip-and-backoff or unscale-and-grow), reimplemented here by hand at the gradient level to test the underlying mechanics rather than just calling the library.

Problem

Training in fp16 saves memory and bandwidth, but fp16's tiny exponent range means small gradients can silently underflow to zero. The standard fix is loss scaling: multiply the loss by a large constant before .backward() so gradients are shifted into fp16's representable range, then divide ("unscale") them by that same constant before the optimizer step. The wrinkle is that scaling too aggressively can push gradients the other direction, into inf/NaN — a scaler has to detect that after the fact, skip the optimizer step for that batch, and back off the scale factor; when steps go cleanly it grows the scale factor to use more of fp16's range.

Implement the core decision logic: given a batch of gradients produced from a scaled loss, either unscale them and grow the scale, or discard them and shrink the scale.

Source: src/47_manual_grad_scaler.py

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

Examples:
>>> grads = [torch.tensor([2.0, 4.0]), torch.tensor([8.0])]
>>> unscaled, new_scale = scaler_step(grads, scale=2.0)
>>> unscaled
[tensor([1., 2.]), tensor([4.])]
>>> new_scale
4.0

Step-by-Step Approach

  1. Check every gradient tensor in the list for non-finite values with torch.isfinite(g).all() — a single overflowed/NaN'd gradient anywhere invalidates the whole step.
  2. If any tensor contains inf/NaN: return None in place of the grads (signals "skip the optimizer step") and shrink the scale by multiplying it by backoff_factor.
  3. Otherwise, the step is safe: divide every gradient by scale to undo the loss-scaling multiplication applied before .backward().
  4. Grow the scale by multiplying it by growth_factor, so the next iteration uses more of fp16's dynamic range if this one stayed finite.
  5. Return the list of unscaled gradients alongside the new scale.

The key insight is that skipping a step is not a bug workaround, it's the intended safety valve: an optimizer step built from partially-overflowed gradients would silently corrupt the model, so discarding the whole step and shrinking the scale is strictly safer than trying to salvage it.

Reference solution

def scaler_step(grads, scale, growth_factor=2.0, backoff_factor=0.5):
    # any single non-finite gradient invalidates the whole step
    found_inf = any(not torch.isfinite(g).all() for g in grads)

    if found_inf:
        # skip optimizer step, shrink scale for the next attempt
        return None, scale * backoff_factor

    # safe: undo the scaling multiplied into the loss before backward()
    unscaled = [g / scale for g in grads]
    # no overflow -> push scale up to use more of fp16's range next time
    new_scale = scale * growth_factor
    return unscaled, new_scale

Key Functions & Tricks

  • torch.isfinite — the standard check for both inf and NaN in one call (distinct from torch.isnan, which misses inf).
  • .all() reduction — collapses the per-element finite check down to one boolean per gradient tensor.
  • any(...) over the list — a single bad tensor anywhere in the parameter list must fail the whole step, not just that tensor's update.
  • Multiplicative growth/backoff — the scale is adjusted geometrically, not additively, matching how torch.cuda.amp.GradScaler tunes its internal scale factor.

How to Recognize This Pattern

The signal is "a numeric range problem where a scalar multiplier shifts values into a safe range, but the safe range itself is unknown and must be discovered adaptively." Loss scaling for fp16 is the canonical example; the same shape shows up in adaptive step-size or trust-region methods more generally (grow when safe, shrink when you overshoot). Real torch.cuda.amp.GradScaler only grows the scale after growth_interval consecutive clean steps (not every single one, as simplified here) to avoid oscillating right at the overflow boundary — worth mentioning if asked to extend this. Common pitfalls: checking for inf/NaN in the loss instead of the gradients (the loss itself can be totally finite while individual small gradients still underflow, or large ones still overflow), and calling optimizer.step() unconditionally instead of gating it on the "found_inf" result.

← Previous Next →