Clip gradients with the norm the engine already computed

[NuojCheng]

Split out of #5138, which was a deferral change and had no business carrying this.

What

_update_kernel computes grad_norm for reporting, then one line later calls
maxtext_utils.apply_gradient_clipping, whose optax.clip_by_global_norm computes
global_norm(grads) again. This clips with the norm already in hand instead. The scale is
optax's exactly: 1 under the threshold, threshold / norm over it.

This is a numerical change, not only a saving

The engine's norm is computed in float32 -- deliberately, because a sum of squares over
bf16 overflows on production-size models -- while optax's runs in the gradients' own dtype.
So under grad_dtype=bfloat16 the engine has been reporting a guarded norm and clipping
with an unguarded one, and it now clips with the guarded one.

At the default grad_dtype=float32 the two are the same number and nothing changes.
Only configs that set grad_dtype: bfloat16 see different clipped gradients, and for those
the new behaviour is the correct one. That divergence is the reason this is a separate PR:
it should be reviewed on its own terms rather than approved as part of a deferral change.

Division-by-zero

Addresses the review left on #5138. Having the norm in hand is what makes the division
guardable, so the guard comes for free here:

under = grad_norm < threshold
safe_norm = jnp.where(under, 1.0, grad_norm)
scale = jnp.where(under, 1.0, threshold / safe_norm)

This is the same idiom the accumulated-denominator division 10 lines above already uses.

Worth noting the direction: optax is the side that fails this. clip_by_global_norm
divides by the norm unguarded and leans on lax.select to discard the result, so on an
all-zero gradient tree the value is right but a NaN is computed to get there, and
JAX_DEBUG_NANS stops on it. So this removes a hazard MaxText has today rather than
avoiding one it would have introduced. test_clip_by_grad_norm_does_not_divide_by_zero
asserts both halves of that.

fp8

Unchanged, still routed to optax: apply_gradient_clipping holds OVERWRITE_WITH_GRADIENT
out of both the norm and the scaling, and reproducing that carve-out is not worth the
saving.

Tests

tests/post_training/unit/maxtext_engine_test.py -- 52 passed on CPU.

• test_clip_by_grad_norm_matches_optax -- both sides of the threshold, against optax
• test_clip_by_grad_norm_does_not_divide_by_zero -- all-zero tree under jax.debug_nans

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[gemini-code-assist]

Code Review

This pull request introduces a custom _clip_by_grad_norm method in maxtext_engine.py to optimize gradient clipping by reusing the already computed gradient norm, avoiding redundant computations and cross-replica reductions. It also safely guards against division-by-zero errors on all-zero gradients, preventing issues under debug_nans. Unit tests are added to verify correctness and the division-by-zero guard. I have no feedback to provide as there are no review comments.

[codecov]

Codecov Report

❌ Patch coverage is 77.77778% with 2 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/training_engine/maxtext_engine.py	77.77%	1 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!