One unusual batch, with an outlier or a long sequence, can produce a gradient thousands of times larger than usual. A single step along it throws the weights far away, and the loss can turn into nan.

Clipping caps the size of the step. All the gradients are treated as one long vector; if its length, the norm, is above a limit, every gradient is scaled down by the same factor. The direction stays, only the length shrinks.

import torch
from torch import nn

model = nn.Linear(2, 1)
x = torch.tensor([[100.0, -50.0]])
model(x).pow(2).sum().backward()
norm = nn.utils.clip_grad_norm_(
    model.parameters(), max_norm=1.0)
print(norm)  # the norm before clipping
grads = [p.grad.flatten()
         for p in model.parameters()]
print(torch.cat(grads).norm())  # 1.0

Clip between loss.backward() and optimizer.step(). clip_grad_norm_ returns the norm before clipping, which is worth logging: a norm that keeps hitting the limit suggests the learning rate is too high.

Write clip_gradients(parameters, max_norm) without clip_grad_norm_: treat the gradients of all parameters that have one as a single vector, scale them all in place by max_norm / norm when its norm is above max_norm, and return the norm before clipping as a float. Write train_step(model, optimizer, x, y, max_norm) doing one clipped step on the mean squared error and returning the loss as a float.

The explanation opens here once every check passes.