Why fp16 Training Stalls Without a GradScaler
A teammate switches a vision model from fp32 to fp16 autocast on a GPU
without bf16 support to double the batch size. The loss goes down for
the first epoch, then plateaus; the pre-clip gradient norm is exactly
0.0 on many steps and some layers' .grad tensors are all zeros. They
propose "just raise the learning rate". A second teammate says "add
GradScaler and the problem disappears; also we should call
clip_grad_norm_ right after loss.backward() as before".
- Explain, with the actual number ranges of fp16, why small gradients vanish and what loss scaling does about it. Use a true gradient of 10^{-8} and an initial scale of 65{,}536.
- A gradient whose true value is 2.0 appears in the same step. What happens at scale 65{,}536, and what does the scaler do next? Trace the scale until the step succeeds.
- Write the correct step (autocast, scaler, clipping, scheduler) and
say what is wrong with clipping right after
backward(). - Would you rather switch to bf16 if the hardware supported it? What do you give up?
- Does fp16 autocast reduce the memory of weights and optimizer state for a 100M-parameter AdamW model? What does it reduce?
1. Underflow
fp16 has a smallest normal value of about 6.1 \times 10^{-5} and a
smallest subnormal near 6 \times 10^{-8}. A gradient of 10^{-8} is
below even the subnormal range and rounds to exactly 0 (you can
check torch.tensor(1e-8).half()). Deep layers and small-loss regimes
produce many such gradients, so whole layers stop learning, which is
why the norm is exactly zero; raising the learning rate multiplies
zeros by a bigger number and does nothing.
Loss scaling multiplies the loss by S before backward(), so by the
chain rule every gradient is multiplied by S. At S = 65{,}536 the
true 10^{-8} becomes 6.6 \times 10^{-4}, well inside the normal
range, and the scaler divides the gradients by S again before the
optimizer step. Weights, gradients used by the optimizer and optimizer
state stay in float32, so the unscaled update is exact enough.
2. Overflow and the backoff
A true gradient of 2.0 becomes 2 \times 65{,}536 = 131{,}072,
above the fp16 maximum of 65{,}504, so it becomes inf. The scaler
detects non-finite gradients, skips optimizer.step() for this
iteration and multiplies the scale by 0.5. At S = 32{,}768 the value
is 65{,}536, still above 65{,}504, so it overflows again and the
step is skipped again; at S = 16{,}384 it is 32{,}768 and fits.
After a long run of clean steps (2,000 by default) the scaler tries a
doubled scale, so it hovers just below the overflow boundary. A scale
that keeps collapsing means real instability (log scaler.get_scale()).
3. The correct step
scaler = torch.amp.GradScaler("cuda")
opt.zero_grad(set_to_none=True)
with torch.autocast(device_type="cuda", dtype=torch.float16):
loss = criterion(model(x), y) # forward + loss only
scaler.scale(loss).backward() # backward outside autocast
scaler.unscale_(opt) # gradients to true scale
nn.utils.clip_grad_norm_(model.parameters(), 1.0)
scaler.step(opt) # skipped if inf/nan present
scaler.update()
sched.step()
Clipping right after backward() would clip scaled gradients: the
norm is about S times the real norm, so with S = 65{,}536 and a
threshold of 1.0 almost every step is rescaled to a tiny vector, and
the effective learning rate collapses. unscale_ first, then clip.
Also note the scheduler may warn when the scaler skipped the first
step(s); that is benign.
4. bf16
Yes, if available. bf16 has the same exponent range as float32 (max about 3.4 \times 10^{38}), so neither underflow of 10^{-8} nor overflow of a few thousand happens and no scaler is needed. The cost is precision: 8 significant bits versus fp16's 11, so 1 + 2^{-9} rounds to 1.0 in bf16 but not in fp16. That is rarely a problem in practice: the weights stay in fp32 and matmul accumulation is done in fp32.
5. Memory
No. Autocast keeps fp32 master weights (400 MB), fp32 gradients (400 MB) and fp32 AdamW moments (800 MB), 1.6 GB in total, exactly as in fp32 training. What shrinks are the activations saved for backward, which are stored in half precision and typically dominate at large batch or sequence length; that is where the extra batch size comes from. Casting the whole model to half precision would halve weights and gradients, but removes the fp32 master copy and small updates can round away.
Share this question