Gradient Clipping and Mixed Precision

A training run that ran clean yesterday produces a loss curve that goes vertical at step 8,217. The culprit is a single batch with a gradient norm of 4,200, twenty times the previous peak. Without clipping, the optimizer applies a step that resets everything the model learned over the previous hour. With a global L2 clip at norm 1.0, that same batch contributes a unit-norm update, the loss stays on its trend line, and the run survives.

Mixed-precision training raises throughput by 2-3x by computing the forward pass and most of the backward pass in FP16. The cost is FP16's narrow exponent range. A gradient that overflows in FP16 evaluates to Inf, which propagates through later layers as NaN, which sets every weight to NaN at the next optimizer step. PyTorch's GradScaler handles this by multiplying the loss by a large scaling factor before the backward pass and dividing the gradients by the same factor before the optimizer step. If any gradient is Inf or NaN at unscale time, the scaler skips the step and halves the factor. If the previous N steps were clean, it doubles the factor. Over a run, the factor settles at the highest value the FP16 range allows.

The build problem is wiring the two together. Clip before unscale and the threshold applies to scaled gradients. Clip after unscale and the order of GradScaler calls matters. The right order is scaler.scale(loss).backward(), then scaler.unscale_(optimizer), then the clip (clip_global_l2_norm in this lesson, clip_grad_norm_ in stock PyTorch), then scaler.step(optimizer), then scaler.update(). Any other order gives a loop that is silently broken.

Global L2 norm

The global L2 norm is the Euclidean norm of the concatenated gradient vector across all parameters, not a per-parameter norm. PyTorch ships it as torch.nn.utils.clip_grad_norm_(parameters, max_norm). This lesson computes the norm by hand instead (compute_global_l2_norm) and scales the gradients in place by max_norm / (pre_clip + 1e-12) when the norm exceeds max_norm. clip_global_l2_norm returns (pre_clip, post_clip), so you can log both the natural and the clipped value. You need both to diagnose "we are clipping at every step."

autocast and GradScaler

torch.amp.autocast(device_type) is a context manager that runs eligible operations (most matmul-class ones) in FP16. torch.amp.GradScaler(device_type) is the helper that scales the loss before backward and inverse-scales the gradients before the optimizer step. They are designed as a pair, so a GPU loop uses both.

This lesson uses CPU autocast because that is what runs in CI. The same pattern transfers verbatim to CUDA by changing device_type="cpu" to device_type="cuda". On CPU the GradScaler is a stub, since CPU autocast already operates in BF16 by default and does not need loss scaling. In this lesson's code the scaler is only enabled when device_type is "cuda". The call sites stay in anyway, so the wiring is identical to the GPU loop.

NaN and Inf detection

Detection happens in two places. First, the loss is checked with torch.isfinite before backward. An Inf or NaN loss does not produce useful gradients, so the step is skipped without entering the optimizer. Second, after scaler.unscale_(optimizer), has_non_finite_grad(...) scans the unscaled gradients and treats any Inf or NaN as a skip. Between them the two checks cover the forward-pass and backward-pass failure modes.

Scaling factor diagnostics

The scaling factor is the GradScaler's internal state. Each step, read scaler.get_scale() and log it next to the learning rate and gradient norm. A healthy run shows the factor climbing in powers of two until it saturates near 2^17 or 2^18. A misbehaving run shows it oscillating between high and low values, which means the gradients are sometimes in range and sometimes not. Without the log, you cannot see any of this. On CPU the scaler is disabled, so the logged value stays constant.

code/main.py implements:

• clip_global_l2_norm: computes the global L2 norm by hand (compute_global_l2_norm), scales gradients in place when the norm exceeds max_norm, and returns (pre_clip, post_clip). It returns early, leaving gradients untouched, when the norm is non-finite or already within max_norm.
• has_non_finite_grad: a helper that scans gradients for NaN and Inf.
• AmpTrainState: wraps a model, an AdamW optimizer, a GradScaler, and an autocast device. Its step(inputs, targets) runs the full clipping, scaling, and skip-on-NaN pipeline.
• StepLog and SkipLog: structured per-step records.
• A demo that trains a small Linear(16, 32) / GELU / Linear(32, 4) model with AdamW and MSE loss for 20 steps, injects an Inf into the gradient on step 5 to exercise the skip path, and prints the resulting log.

Run it:

python3 code/main.py

The script exits zero and prints a per-step log with each row tagged STEP or SKIP. At least one row is a SKIP.

Four patterns turn the loop into a production training step.

Skip counter as an alert, not a log line. A handful of skipped steps per run is healthy. Hundreds of skips per epoch should raise a hard alert: the model is in a regime FP16 cannot hold, and the loop is failing silently. The lesson tracks a 1,000-step rolling skip rate, and in production you would page on a rate above 5 percent.

Clip threshold lives in the config. max_norm = 1.0 is the modern default for language-model training. Sweep it on a small model first. Larger thresholds let the model recover from genuinely difficult batches, while smaller thresholds bound the worst case at the cost of a noisier loss curve. Put the threshold in the same YAML or JSON config as the schedule from lesson 44.

Norm log goes to a CSV with the schedule. The columns are step, lr, grad_l2_pre_clip, grad_l2_post_clip, loss, skipped, skip_reason, scaler_scale. A reviewer who opens the file sees the schedule, the gradient story, the scaling factor, and the skip outcome with its reason in one row. Splitting the columns across files is a recipe for misaligned analyses.

scaler.update() runs on clean steps and on non-finite-gradient skips. On a clean step the scaler reads its no-inf counter, increments it, and possibly doubles the factor. On a non-finite-gradient skip it halves the factor and resets the counter. Forgetting update() on that path produces the bug where the scaling factor never changes. A non-finite-loss skip is the exception: scaler.scale(loss).backward() never ran, so calling update() would break GradScaler's call ordering, and the loop deliberately leaves it out.

Production patterns:

• Autocast device matches optimizer device. Use torch.amp.autocast(device_type="cuda") for GPU training and torch.amp.autocast(device_type="cpu") for CPU. A mismatch produces a silent type error: the loss curve looks fine while the model is not learning.
• Loss check before backward. torch.isfinite(loss).all() is one tensor reduction. The cost is negligible, and on a NaN loss it saves an entire training step. Run it on every step.
• set_to_none=True in zero_grad. This sets gradients to None instead of zero, which lets the optimizer skip computation for unaffected parameter groups. It is a free throughput improvement and shrinks the bug surface slightly.

On a real project, outputs/skill-clip-amp.md would describe which clip threshold and autocast device the training step uses, where the per-step CSV lives in version control, and what the production skip-rate alert threshold is. This lesson ships the engine.

1. Replace the synthetic Inf injection with a real loss spike (multiply one batch's target by 1e8) and verify the skip path triggers.
2. Add a --bf16 mode that switches autocast to BF16 instead of FP16. BF16 has a wider exponent range than FP16 and rarely needs loss scaling, so overflow-driven skips should largely disappear. The injected Inf still triggers its skip, so compare the skip counts with that step excluded.
3. Add a unit test that the gradient-clip wrapper returns the pre-clip and post-clip norm correctly when no clipping occurs.
4. Add a rolling-window skip-rate computation and a CLI flag that fails the run if the rate exceeds a configured threshold for 100 consecutive steps.
5. Wire the loop to write the canonical CSV (step, lr, grad_l2_pre_clip, grad_l2_post_clip, loss, skipped, skip_reason, scaler_scale) and confirm the file survives a Ctrl-C by flushing after every row.

• Micikevicius et al., Mixed Precision Training (arXiv 1710.03740): the original loss-scaling proposal
• Pascanu, Mikolov, Bengio, On the difficulty of training recurrent neural networks (arXiv 1211.5063): the gradient-clipping reference paper
• PyTorch torch.amp.GradScaler: the scaler API this lesson wraps
• PyTorch torch.nn.utils.clip_grad_norm_: the clipping primitive this lesson uses
• Phase 19 · 42: the downloader whose corpus feeds the loop
• Phase 19 · 43: the dataloader the loop consumes
• Phase 19 · 44: the schedule this loop composes with