Gradient Clipping by Global Norm - Optimization & Training

Gradient Clipping by Global Norm

What you'll build

clip_grad_norm(grads, max_norm): treat the gradients of all parameters as one long vector, compute its L2 norm, and if it exceeds max_norm, rescale every gradient by the same factor. This is what torch.nn.utils.clip_grad_norm_ does.

Background

A single bad batch can produce a huge gradient that throws the weights far away: a loss spike, sometimes followed by NaNs. Clipping caps the size of the step. Clipping by the global norm (rather than per tensor or per element) keeps the update's direction and the relative sizes of all gradients; only the step length changes. Most LLM training runs clip with max_norm=1.0.

Formula

total_norm = sqrt( sum_i sum(g_i ** 2) )        # over every array and every element
clip_coef  = max_norm / (total_norm + eps)       # eps = 1e-6, as in PyTorch
if clip_coef < 1:
    g_i = g_i * clip_coef                        # same factor for every array

• g_i: the gradient array of parameter i (any shape)
• max_norm: the largest allowed global norm
• eps: guards against division by zero when every gradient is 0

Inputs

• grads: list of NumPy float arrays of any shapes
• max_norm: positive float
• eps: small float, default 1e-6

Expected return format

Return a tuple (clipped_grads, total_norm):

• clipped_grads: a list of new arrays with the same shapes as grads (don't modify the inputs)
• total_norm: the global norm before clipping, as a Python float

Edge cases and conventions

• Gradients whose global norm is at most max_norm come back unchanged (copies).
• After clipping, the global norm equals max_norm (up to the tiny eps).
• All-zero gradients: total_norm = 0.0 and nothing changes.

Constraints

NumPy only; one pass over the arrays (no per-element Python loops).

Follow-up

• Where does clipping go in a mixed-precision loop? (After scaler.unscale_(optimizer), before scaler.step.)
• How does clipping by value (clip_grad_value_) differ, and why does it change the update direction?

Input: clip_grad_norm([np.array([3.0, 4.0])], max_norm=1.0)
Output: ([array([0.6, 0.8])], 5.0)

Input: clip_grad_norm([np.array([1.0, 2.0]), np.array([2.0])], max_norm=5.0)
Output: ([array([1., 2.]), array([2.])], 3.0)

import numpy as np

def clip_grad_norm(grads: list, max_norm: float, eps: float = 1e-6):
    """
    Clip a list of gradient arrays by their global L2 norm.

    Args:
        grads: List of gradient arrays (any shapes)
        max_norm: Maximum allowed global norm
        eps: Small constant added to the norm (PyTorch uses 1e-6)

    Returns:
        (clipped_grads, total_norm): list of arrays, and the norm before clipping
    """
    # Your code here
    pass