Gradient Clipping

One line of insurance

Gradient clipping caps the global norm (or value) of the gradient before the optimizer step. If the unclipped norm exceeds max_norm, scale the gradient down so the norm equals max_norm. The optimizer then sees a bounded update and can't take a wild step.

Standard default: max_norm=1.0 for transformers and RNNs, max_norm=0.5 for very unstable training, no clipping for stable CNN training. The cost is a tiny bit of compute per step — basically free.

Tip: If your training ever spits out a NaN loss without warning, the first patch is to add gradient clipping. The second is to log gradient norms before clipping so you can see how often it's saving you.

Norm clipping vs value clipping

Norm clipping (clip_grad_norm_) — preserve direction, scale magnitude. Almost always what you want.

Value clipping (clip_grad_value_) — cap each coordinate independently. Distorts the gradient direction; rarely what you want unless you're doing something specific (e.g. PPO in RL).

Where to put the call

Between loss.backward() and opt.step(). With mixed precision, between scaler.unscale_(opt) and scaler.step(opt). The order matters because clipping operates on the actual gradient values.

Principle: Gradient clipping is cheap insurance. Default to having it in every training loop, and only remove it if you have a specific reason. The cost of a NaN at step 50,000 is much higher than the cost of one extra line.

Code

Gradient clipping with mixed precision·python

import torch
from torch.cuda.amp import autocast, GradScaler
from torch.nn.utils import clip_grad_norm_

scaler = GradScaler()

opt.zero_grad()
with autocast():
    loss = loss_fn(model(xb), yb)

scaler.scale(loss).backward()
scaler.unscale_(opt)                # bring grads back to FP32 scale
clip_grad_norm_(model.parameters(), max_norm=1.0)
scaler.step(opt)
scaler.update()

One line of insurance

Norm clipping vs value clipping

Where to put the call

Code

External links

Exercise

Progress

Comments 0