Clip gradients by their global norm

Easy10 min11 hidden tests

A team training a language model sees the loss spike whenever one bad batch produces huge gradients. They clip by global norm: treat the gradients of every parameter tensor as one long vector, and if its length (the L2 norm) is above max_norm, scale every gradient down by the same factor so the length becomes exactly max_norm. Write clip_by_global_norm(grads, max_norm) that returns the clipped gradients and the norm measured before clipping. Gradients [3.0, 0.0] and [[0.0, 4.0]] have a global norm of 5.0, so with max_norm = 1.0 they become [0.6, 0.0] and [[0.0, 0.8]].

Drag the corner, set max_norm

scale = max_norm / ‖g‖ when ‖g‖ > max_norm
‖g₁‖, first tensor‖g₂‖, second tensorGlobal normmax_norm

Global norm √(3.00² + 4.00²) = 5.00 is above max_norm 1.0, so every gradient is multiplied by 1.0 / 5.00 = 0.20: ‖g₁‖ becomes 0.60 and ‖g₂‖ 0.80. Both shrink by the same factor, so the update keeps its direction.

Norm 5.00, clipped by 0.20.

Examples

Norm 5.0 is above 1.0, so everything shrinks by a factor of 5

clip_by_global_norm([[3.0, 0.0], [[0.0, 4.0]]], 1.0)

returns ([[0.6, 0], [[0, 0.8]]], 5)

Norm 0.5 is already small enough

clip_by_global_norm([[0.3, 0.4], [0.0]], 1.0)

returns ([[0.3, 0.4], [0]], 0.5)

Requirements

  • grads is a list of one or more arrays (or nested lists) of any shapes. max_norm is a positive float.
  • The global norm is sqrt of the sum of the squares of every entry in every array.
  • If the global norm is above max_norm, multiply every array by max_norm / norm. Otherwise, including a norm of 0, leave the values as they are.
  • Return a tuple (clipped, norm): a list of float arrays with the same shapes as the inputs, and the global norm before clipping as a float.
The idea behind it

Gradient clipping limits how far one step can move the weights. In deep networks and long sequences, a rare batch can produce a gradient hundreds of times larger than usual, and one such step can wreck a training run. Clipping by global norm measures the size of the whole update as one vector, ||g|| = sqrt(sum of every entry squared), across all parameter tensors. If that norm is above a threshold, every gradient is multiplied by max_norm / ||g||. The update keeps its direction and only its length shrinks. Clipping each tensor on its own, or capping each value, would bend the direction instead. Most LLM recipes clip at a norm of 1.0, and the norm before clipping is logged too: a run whose gradient norm keeps climbing is about to diverge.

pythonCtrl+Enter runs
import numpy as np

def clip_by_global_norm(grads, max_norm):
"""Scale a list of gradient arrays so their combined L2 norm is at most max_norm.

grads: list of arrays (any shapes), one per parameter tensor
max_norm: positive float, the largest allowed global norm
Returns (clipped, norm): a list of float arrays shaped like grads,
and the global norm measured before clipping, as a float.
"""
raise NotImplementedError

Stuck? Reveal one rung at a time

Worked solution

Give it a real try first. Then read the solution and why each line is there.