How do torch.amp.autocast and GradScaler fit into a PyTorch training loop?

## What autocast does `torch.amp.autocast` is a context manager that intercepts operator dispatch. Inside it, operations on a per-op allowlist — matrix multiplies, convolutions, linear layers, the heavy compute — are cast to the chosen 16-bit dtype, while operations that are sensitive to precision loss — softmax, layer norm, most reductions, loss functions — stay in float32. You do not cast the model or the inputs yourself; you wrap the region and let the dispatcher decide per operation. The parameters themselves stay in float32. This is the "mixed" in mixed precision: a float32 master copy of the weights, 16-bit arithmetic for the expensive kernels, float32 accumulation inside those kernels on hardware with tensor cores. The payoff is roughly halved activation memory and substantially faster matmuls on modern GPUs. Two placement rules follow directly: - **Only the forward pass and the loss go inside the context.** Backward is deliberately outside; autograd records the dtype each op used going forward and reuses it going backward, so wrapping backward gains nothing and can confuse the scaling logic. - **The optimizer step is outside too**, because it must operate on float32 master weights. ## Why float16 needs a loss scale float16 has 5 exponent bits, giving a smallest normal magnitude around 6e-5. Gradients in a deep network routinely sit below that, and anything smaller flushes to zero — the update simply disappears, and the model trains to a visibly worse result with no error anywhere. Loss scaling fixes it by exploiting linearity: multiply the loss by S before `backward()`, and every gradient in the graph comes out multiplied by S too, shifted up into the representable range. Divide by S before the optimizer sees them and you are back where you started, minus the underflow. bfloat16 has 8 exponent bits — the same range as float32, traded against fewer mantissa bits — so gradients do not underflow and no scaling is needed. That is the main practical reason bf16 is the default choice on hardware that supports it. ## The GradScaler protocol The four calls, in order: 1. `scaler.scale(loss).backward()` — multiplies the loss by the current scale factor, then runs the normal backward. `.grad` now holds scaled gradients. 2. `scaler.unscale_(optimizer)` — *optional*, and only needed if you want to touch the gradients yourself, most commonly to clip them. It divides the gradients in place by the scale factor. 3. `scaler.step(optimizer)` — unscales the gradients if step 2 did not, inspects them for inf or NaN, and calls `optimizer.step()` **only if they are all finite**. If any are not, the step is skipped entirely. 4. `scaler.update()` — adjusts the scale for the next iteration: multiply it up after a run of successful steps, cut it back sharply after a skipped one. That skipping behaviour is the part candidates most often miss. Early in training you should expect a handful of skipped steps while the scaler finds a workable factor — that is normal, not a bug. But it means the number of optimizer steps is not exactly the number of iterations, which slightly desynchronizes anything counting steps, and it means a persistently skipped step (`scaler.get_scale()` collapsing toward 1) is a signal of genuine numerical trouble in the model, not a scaling problem. ## What must stay outside - Backward and the optimizer step, as above. - Anything that manually inspects or edits `.grad` must come after `unscale_`, or it operates on scaled values. - Custom operations that are numerically fragile can be forced back to float32 by wrapping them in `autocast(device_type=..., enabled=False)` and casting the inputs explicitly. - Do not assume tensors coming out of the context are float32 — a value produced inside autocast may be float16, and combining it later with a float32 tensor in an in-place operation can raise a dtype error. ## Checkpoints and diagnostics `GradScaler` has its own `state_dict()`/`load_state_dict()` holding the current scale and growth tracker. Include it in checkpoints alongside the model, optimizer and scheduler, otherwise a resumed run rediscovers the scale from the initial value and may skip a few steps at the seam. When debugging suspected AMP problems, the standard sequence is: disable the scaler and see whether the NaNs persist (if they do, the model is unstable regardless of precision); try bfloat16 instead of float16 (if that fixes it, it was range, not precision); and inspect `scaler.get_scale()` over time to see whether steps are being skipped repeatedly. ## The namespace to use Write `torch.amp.autocast(device_type="cuda", ...)` and `torch.amp.GradScaler("cuda")`. The older device-specific spellings under `torch.cuda.amp` still exist but are the deprecated form; the `torch.amp` namespace is the current one and takes the device type as an argument, which is what makes the same code run on CPU or other backends.