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_normGlobal 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
gradsis a list of one or more arrays (or nested lists) of any shapes.max_normis a positive float.- The global norm is
sqrtof the sum of the squares of every entry in every array. - If the global norm is above
max_norm, multiply every array bymax_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.
Worked solution
Give it a real try first. Then read the solution and why each line is there.