Mixed Precision Training › Loss Scaling

Loss scaling shifts gradient values into a range where FP16 can represent them accurately, then undoes the shift before applying them to the weights.

Gradient values during training tend to be small, often clustering in the range $[10^{-8}, 10^{-2}]$. FP16 can represent values in this range, but values near the bottom of the range lose most of their significant bits. More critically, underflow to zero is common for the smallest gradients.

The trick is to multiply the loss by a large scalar $S$ (the scale factor) before the backward pass. By the chain rule, every gradient throughout the network is multiplied by $S$. This shift moves all gradient values upward in magnitude, reducing underflow and making them comfortably representable in FP16. Think of it as temporarily stretching the gradient magnitude range so that the values that would have been too small for FP16 are now safely within its representable range.

The chain rule argument is worth spelling out explicitly. If the loss is $L$ and the scale factor is $S$, the scaled loss is $L' = S \cdot L$. The gradient of $L'$ with respect to any parameter $\theta$ is:

$$\frac{\partial L'}{\partial \theta} = S \cdot \frac{\partial L}{\partial \theta}$$

So the backward pass on the scaled loss automatically produces gradients that are each multiplied by $S$, with no additional code required. The chain rule propagates the scaling factor through the entire computation graph.

Concretely, if the unscaled gradient is $g$ and the scale factor is $S$, then the scaled gradient stored in FP16 is $\tilde{g} = S \cdot g$. After the backward pass, before the optimizer update, we divide by $S$:

$$g = \frac{\tilde{g}}{S}$$

where:

• $g$: the true gradient that would have been computed in FP32
• $\tilde{g}$: the scaled gradient stored in FP16
• $S$: the loss scale factor, typically a power of 2 to avoid rounding

The division by $S$ is done in FP32 (on the master weight gradients), so it does not suffer from precision loss. The whole operation is mathematically a no-op: you multiply by $S$ before the backward pass and divide by $S$ after. The only effect is that the intermediate representation in FP16 uses a better portion of the available bit range.

Why use powers of 2 for $S$? Because powers of 2 can be divided exactly in floating point arithmetic. Dividing by $2^{15}$ is equivalent to subtracting 15 from the FP32 exponent, which introduces zero rounding error. Any other value of $S$ would introduce a small rounding error in the unscaling step, partially defeating the purpose.

Let us trace through the full loss scaling procedure with concrete numbers to see precisely what changes and what stays the same.

Suppose a gradient value computed during the backward pass, at some layer, is $g = 0.00000245$ (approximately $2.45 \times 10^{-6}$). In FP32, this is representable with full precision. In FP16, the nearest representable normal value is approximately $2.44 \times 10^{-6}$, a small rounding error of about 0.4%. So far, FP16 handles this reasonably well.

Now suppose the gradient update for a weight with value 1.0 is $\eta \cdot g = 10^{-4} \times 2.45 \times 10^{-6} = 2.45 \times 10^{-10}$. The ULP for the value 1.0 in FP16 is approximately $9.77 \times 10^{-4}$. A gradient update of $2.45 \times 10^{-10}$ is roughly 4 million times smaller than the ULP. In FP16, the weight simply does not change.

With loss scaling by $S = 2^{15} = 32768$:

The scaled gradient is $\tilde{g} = S \cdot g = 32768 \times 2.45 \times 10^{-6} = 0.0803$. This is in the middle of FP16's normal range and carries full precision. The FP16 representation is $\tilde{g}_{\text{fp16}} \approx 0.0803$, negligible rounding.

After the backward pass, the unscaling step divides in FP32: $g_{\text{recovered}} = 0.0803 / 32768 = 2.45 \times 10^{-6}$. This is the original gradient, recovered with essentially no error. The optimizer applies this gradient to the FP32 master weight, and the update proceeds correctly.

The effective change is that the gradient $2.45 \times 10^{-6}$ was temporarily multiplied by 32768 during the FP16 backward pass to move it from the subnormal/dangerous region of FP16's range into a well-represented middle region, then divided back to restore the original value in FP32 before the optimizer touches it.

The scale factor must be large enough to prevent underflow but small enough to prevent overflow. If $S$ is too large, the scaled gradients themselves overflow FP16 (exceeding 65504), producing Inf values. If $S$ is too small, gradients underflow to zero.

Static loss scaling uses a fixed $S$, typically around $2^{15} = 32768$ or $2^{16} = 65536$. This works reasonably well for many models but requires tuning. A scale that works for one architecture may fail for another.

Dynamic loss scaling addresses this by adjusting $S$ automatically during training. The algorithm is:

1. Start with a large scale factor (e.g., $2^{15}$).
2. After each backward pass, check whether any gradient is Inf or NaN.
3. If overflow is detected: skip the optimizer step (the corrupted gradients would harm the model), reduce $S$ by a constant factor (e.g., halve it), and continue.
4. If no overflow is detected: periodically increase $S$ (e.g., multiply by 2 every 2000 steps, or whenever overflow has not been seen for a stretch).

The adaptive behavior of dynamic loss scaling is important in practice. Early in training, gradients can be large and scale should be lower. As training stabilizes, gradients become smaller and scale can safely increase. A static scale chosen for late-stage training would overflow early; a scale chosen for early training would underflow late.

The skipped steps when overflow is detected are not wasted. The model state is preserved; only the corrupted update is discarded. The scale factor is immediately reduced, which prevents the next step from overflowing as well. In practice, overflow events are rare after the initial warmup period, and the skipped steps represent a tiny fraction of total training compute.

The growth interval (the number of consecutive clean steps before the scale increases) is a subtle but important tuning parameter. The default of 2000 in PyTorch reflects a conservative balance: you want the scale to grow quickly enough to keep gradients well-represented, but not so aggressively that it oscillates between overflow and recovery. If your training shows frequent overflow events throughout, consider reducing the growth factor or increasing the growth interval. If overflow events disappear quickly but you suspect underflow later in training, a shorter growth interval allows the scale to rise faster and capture fine-grained gradient information.

Gradient clipping (see the chapter on Gradient Clipping in Part X) caps gradient norms before they are applied to weights, preventing catastrophically large updates. When using loss scaling, gradient clipping must happen on the unscaled gradients. The GradScaler.unscale_() method in PyTorch handles this by dividing the scaled gradients by $S$ in-place before you call the gradient clip function.

The correct order of operations is:

1. Scaled backward pass: scaled_loss = scaler.scale(loss); scaled_loss.backward()
2. Unscale gradients: scaler.unscale_(optimizer)
3. Clip gradients: torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
4. Optimizer step: scaler.step(optimizer) (internally checks for Inf/NaN and skips if found)
5. Update scale: scaler.update()

Applying gradient clipping before unscaling would clip the scaled gradients at the wrong threshold, effectively tightening the clip by a factor of $S$. If your loss scale is 65536 and your intended clip norm is 1.0, clipping before unscaling would apply a clip of $1.0 / 65536 \approx 0.000015$, cutting off nearly all gradient signal.