Dynamic Loss Scaling for FP16

The problem

Training in FP16 (16-bit floating point) halves the memory and is much faster, but FP16 has a small range: its largest number is 65504, and numbers below about $6 \times 10^{-8}$ round to exactly 0. Many gradients are that small, and they would simply vanish.

Loss scaling fixes it: multiply the loss, and therefore every gradient, by a large factor $S$ before the backward pass, so that small gradients stay representable; then divide the gradients by $S$ , in FP32 (32-bit floating point), before the update. You will simulate this, and then write the dynamic scaler that chooses $S$ during training, as PyTorch does.

Part 1: a round trip through FP16

Simulate a backward pass in FP16: multiply by the scale, store in FP16, read back in FP32, divide by the scale.

Example. Three gradients, $10^{-8}$ , $10^{-3}$ and $1.0$ .

With scale = 1: $10^{-8}$ comes back as 0 (lost), $10^{-3}$ as $0.0010004$ (slightly rounded), $1.0$ as $1.0$ .

With scale = 65536 ( $2^{16}$ ): $10^{-8}$ comes back as $0.99997 \times 10^{-8}$ (saved), $10^{-3}$ as $0.0010004$ , and $1.0$ as inf: $1.0 \cdot 65536 = 65536$ is above 65504, so it overflows.

The scale has to be large enough for the small gradients and small enough for the large ones, and the right value changes during training. Hence the dynamic scaler.

Part 2: the dynamic scaler

The scaler keeps the current scale and a count of good steps in a row. At every step it receives the gradients of the scaled loss, in FP16:

If any value is inf or nan, the scale was too large: multiply scale by backoff_factor (0.5 by default), reset the count, and return None, meaning "skip this step".
Otherwise, return the gradients unscaled: cast to FP32 and divided by the scale used for this step. Count one good step; after growth_interval good steps in a row, multiply scale by growth_factor (2 by default) and reset the count, to try a larger scale again.

Example. init_scale = 8 and growth_interval = 2:

step gradients received returned scale after the step

1 [8, 16] [1, 2] 8 (1 good step)

2 [inf] None 4 (halved, count reset)

3 [4] [1] 4 (1 good step)

4 [4] [1] 8 (2 good steps: doubled)

In step 4 the gradient is divided by 4, the scale used to produce it, and only then is the scale doubled.

step	gradients received	returned	scale after the step
1	`[8, 16]`	`[1, 2]`	8 (1 good step)
2	`[inf]`	`None`	4 (halved, count reset)
3	`[4]`	`[1]`	4 (1 good step)
4	`[4]`	`[1]`	8 (2 good steps: doubled)

What to implement

fp16_roundtrip(g, scale): multiply the FP32 array g by scale, cast to np.float16, cast back to np.float32 and divide by scale.
class DynamicLossScaler(init_scale=2.0**16, growth_factor=2.0, backoff_factor=0.5, growth_interval=2000), with an attribute scale (starting at init_scale) and a method step(scaled_grads), where scaled_grads is a list of FP16 arrays. It returns None or the list of unscaled FP32 arrays, as in part 2.

Requirements

Unscale with the scale that was used for this step, before changing it.
The defaults are those of PyTorch's torch.amp.GradScaler.

Hints

np.isfinite(g).all() checks one array for inf and nan.
Do the division in FP32: dividing in FP16 can underflow to 0 again.