Study Plan

Global Norm Gradient Clipping

JAX

Medium

Clip a gradient pytree by its global L2 norm, which is the length of all its leaves taken together as one long vector.

\lVert g \rVert = \sqrt{\sum_{\ell} \sum_{i} g_{\ell,i}^{2}}, \qquad s = \begin{cases} 1 & \text{if } \lVert g \rVert \le c \\ c \,/\, \lVert g \rVert & \text{otherwise} \end{cases}

Here $g_{\ell,i}$ is entry $i$ of leaf $\ell$ of grads, $c$ is max_norm, and every leaf is multiplied by the same scale $s$ . Gradients whose global norm is already at most max_norm come back unchanged, and larger ones shrink until their global norm equals max_norm.

Your function is called inside jax.jit with grads and max_norm as traced arguments.

Return a dict with two keys. The value under clipped is a pytree with the structure of grads whose leaves are JAX arrays with the shape and dtype of the matching gradient leaf, and the value under global_norm is the norm before clipping as a float32 JAX array of shape ().

Examples

Input: grads = {"w": [3.0, 0.0], "b": 4.0}, max_norm = 1.0

Output: {"clipped": {"b": 0.8, "w": [0.6, 0.0]}, "global_norm": 5.0}

Explanation: The squares of all entries add up to 9 + 0 + 16 = 25, so the global norm is 5. It is larger than max_norm, so every leaf is multiplied by 1 / 5, which brings the global norm down to exactly 1.

Input: grads = {"dense": {"w": [[0.1, -0.2], [0.2, 0.0]], "b": [0.4, 0.0]}, "head": [-0.4]}, max_norm = 2.0

Output: {"clipped": {"dense": {"b": [0.4, 0.0], "w": [[0.1, -0.2], [0.2, 0.0]]}, "head": [-0.4]}, "global_norm": 0.640312}

Input: grads = {"layers": [{"w": [0.6, 0.6]}, {"w": [0.6, 0.6]}, {"w": [0.6, 0.6]}]}, max_norm = 1.0

Output: {"clipped": {"layers": [{"w": [0.408248, 0.408248]}, {"w": [0.408248, 0.408248]}, {"w": [0.408248, 0.408248]}]}, "global_norm": 1.469694}

Hint 1

jax.tree.leaves(tree) returns every leaf as a flat list.

Hint 2

jnp.sqrt(sum(jnp.sum(l.astype(jnp.float32) ** 2) for l in leaves)) is the global norm.

Hint 3

jnp.minimum(1.0, max_norm / norm) gives the scale without a Python if.

Requirements

Choose between the two cases with array operations such as jnp.where or jnp.minimum, because a Python if on a traced value fails under jax.jit.
Compute the squared sums and the norm in float32 even when the leaves are bfloat16, then cast each clipped leaf back to the dtype of its gradient leaf.

Constraints

grads is a pytree of dicts and lists with at least one leaf, and each leaf has between zero and three dimensions.
All leaves share one dtype, float32 or bfloat16. Public cases and custom Run inputs use float32.
max_norm is a positive finite Python float.
Every gradient value is finite, and all of them may be zero.
In the Run panel, a JSON list whose items are all objects is read as a list node of the tree, and any other list or number is read as one array leaf.

Try Similar Problems

AdamW from ScratchMedium SGD with Momentum over a PytreeEasy Gradient Accumulation with ScanHard Pad and Mask to Bucket LengthMedium Hessian-Vector ProductMedium Per-Example Gradient ClippingMedium Straight-Through Estimator with Custom VJPMedium Traced Learning Rate ScheduleMedium Exponential Moving Average with ScanMedium Spectral Norm by Power IterationMedium Batched Attention with vmapMedium Pairwise Distances with Nested vmapMedium