Loss Scaling: Scale the Loss Before Backward, Unscale Gradients Before the Step

The problem, from first principles

01-fp16-bf16-representable-range showed that float16 values smaller than roughly 6e-8 vanish to exactly 0.0 — and real gradients, especially late in training or in deep networks, routinely take on values in exactly that danger zone. If a gradient silently underflows to zero, its parameter simply stops updating, with no error raised anywhere.

From theory to code

Implement scale_loss(loss, scale), unscale_gradients(gradients, scale), and has_inf_or_nan(gradients). The trick: multiply the loss by some large scale before backpropagation (every gradient scales up by the same factor, via the chain rule), keeping them safely away from the underflow floor throughout backward — then divide back down by scale after backward, before the optimizer actually uses them.

Constraints

scale_loss(loss, scale) returns loss * scale.
unscale_gradients(gradients, scale) returns each gradient array divided by scale, same shapes, as a new list.
has_inf_or_nan(gradients) returns True if any element in any gradient array is inf or nan; False (including for an empty list) otherwise.
None of the three mutate their inputs.

Hints

Open one at a time. Each gives away a little more than the last.

Hint 1

The chain rule means scaling the loss by a constant scales every gradient computed from it by that exact same constant — so a gradient that would have been 1e-6 (dangerously close to fp16's floor) becomes 1e-6 * scale, comfortably clear of it, if scale is large enough.

Hint 2

np.isfinite(g) is False for both inf and nan — checking not np.all(np.isfinite(g)) per array, then combining across the whole list with any(...), catches either failure mode in one pass.

Loss Scaling: Scale the Loss Before Backward, Unscale Gradients Before the Step

The problem, from first principles

From theory to code

Constraints

Hints

Keywords