Defer the engine's live-object publish out of the step path - #5138
Defer the engine's live-object publish out of the step path#5138NuojCheng wants to merge 2 commits into
Conversation
`MaxTextTrainingEngine.fwd_bwd` and `update` each call `nnx.update` on the live NNX objects on every step. `nnx.update` walks the module graph, so its cost scales with the graph's size rather than the state written: 6.2 ms to publish 180 leaves of non-parameter state into an unrolled 28-layer qwen3-0.6b, once per micro-batch, and 16.7 ms for the parameters once per update. That was the last per-step graph walk left after `_refresh_pure_state` removed the two `nnx.split` calls. Nothing inside the step path reads those objects back -- `_read_model_pure` and `_read_state_pure` answer from the pure mirror, and the live objects serve only as containers -- so the walk is pure publication, for readers outside the engine. Defer it: `_publish_to_live` marks the mirror ahead of the live objects, and `_sync_live_objects` flushes before anything reads real values off them (the `model`/`optimizer`/`state` properties, `_refresh_pure_state`, checkpoint save and restore, weight sync, and the throttler's fallback when there is no gradient norm to wait on). `_invalidate_pure_state` flushes before forgetting the mirror, so dropping the cache never drops a step with it. The sync is not merely cosmetic: `update()` donates the state it hands the kernel, so a deferred reader would find deleted arrays, not stale ones. Publication stays eager whenever the pure state is disabled. Also stop computing the global gradient norm twice. `_update_kernel` computes `l2norm_pytree(grads)` for the metric and then calls `maxtext_utils.apply_gradient_clipping`, which recomputes the same norm inside `optax.clip_by_global_norm` -- a second full pass over every gradient plus a second cross-replica reduction. `_clip_by_grad_norm` applies optax's own scale to the norm already in hand. It is also the more accurate of the two: the engine's norm is computed in float32 because a sum of squares over bf16 overflows on production-size models, while optax's runs in the gradients' dtype, so under `grad_dtype=bfloat16` the engine was reporting a guarded norm and clipping with an unguarded one. At the default `grad_dtype=float32` the two are the same number. The fp8 path keeps the optax call, which holds `OVERWRITE_WITH_GRADIENT` out of both the norm and the scaling. qwen3-0.6b on 4x v6e, dp=4, per-device batch 2, seq 1024, adamw, bf16 compute and fp32 weights; median steady-state step over 9 post-warmup steps: GA=8 unrolled 878.4 ms -> 826.1 ms (-6.0%) GA=8 unrolled, no clip 878.4 ms -> 836.7 ms (-4.7%) GA=1 unrolled 144.0 ms -> 143.8 ms (no change) GA=8 scanned 852.2 ms -> 850.4 ms (no change) Losses are bit-identical across all 12 steps of every pair.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request optimizes the MaxText training engine by deferring the publication of pure state changes to live NNX objects, thereby avoiding expensive module graph walks on the step path. It also introduces a custom gradient clipping method (_clip_by_grad_norm) that reuses the pre-computed gradient norm to eliminate redundant passes and cross-replica reductions. The feedback suggests addressing a potential division-by-zero or NaN propagation issue in JAX when the gradient norm is exactly zero by using a safe division idiom, and adding a corresponding unit test case to verify robustness.
| threshold = self._config.gradient_clipping_threshold | ||
| if maxtext_utils.OVERWRITE_WITH_GRADIENT in grads: | ||
| return maxtext_utils.apply_gradient_clipping(grads, None, threshold) | ||
| scale = jnp.where(grad_norm < threshold, 1.0, threshold / grad_norm) | ||
| return jax.tree.map(lambda g: g * scale.astype(g.dtype), grads) |
There was a problem hiding this comment.
To prevent potential division-by-zero warnings or undefined behavior in JAX when grad_norm is 0.0 (e.g., when gradients are all zero), it is safer to use a standard JAX idiom that replaces the denominator with a dummy non-zero value in the unselected branch of jnp.where. Even though jnp.where selects the 1.0 branch when grad_norm < threshold, both branches are evaluated, which can trigger division-by-zero or NaN/inf propagation under certain compiler configurations.
| threshold = self._config.gradient_clipping_threshold | |
| if maxtext_utils.OVERWRITE_WITH_GRADIENT in grads: | |
| return maxtext_utils.apply_gradient_clipping(grads, None, threshold) | |
| scale = jnp.where(grad_norm < threshold, 1.0, threshold / grad_norm) | |
| return jax.tree.map(lambda g: g * scale.astype(g.dtype), grads) | |
| threshold = self._config.gradient_clipping_threshold | |
| if maxtext_utils.OVERWRITE_WITH_GRADIENT in grads: | |
| return maxtext_utils.apply_gradient_clipping(grads, None, threshold) | |
| safe_grad_norm = jnp.where(grad_norm < threshold, 1.0, grad_norm) | |
| scale = jnp.where(grad_norm < threshold, 1.0, threshold / safe_grad_norm) | |
| return jax.tree.map(lambda g: g * scale.astype(g.dtype), grads) |
| exercised by a given gradient tree. | ||
| """ | ||
| t = maxtext_engine.MaxTextTrainingEngine(self.mock_config) | ||
| for scale, case in ((100.0, "above the threshold"), (1e-3, "below the threshold")): |
There was a problem hiding this comment.
Add a test case for 0.0 (zero gradient) to verify that the gradient clipping logic handles zero gradients robustly without division-by-zero or NaN propagation.
| for scale, case in ((100.0, "above the threshold"), (1e-3, "below the threshold")): | |
| for scale, case in ((100.0, "above the threshold"), (1e-3, "below the threshold"), (0.0, "zero gradient")): |
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
…ange `eval_step` is the one step path that reads real values off the live model rather than the pure mirror, so it is the one place a deferred publish is observable -- and it does not observe stale weights, it raises: `update()` donates the state it hands the update kernel, so the parameters `eval_step` splits out of the live model are deleted arrays. `_compile_eval_for_batch` splits them too, to take its shardings, so it syncs as well rather than leaning on its only caller. The regression test goes after an `update()` deliberately. The pre-existing eval tests call `eval_step` on a freshly compiled engine, which has never deferred anything, which is why this got through them. Also documents why one write to `self._state` discharges both deferral sites, and why the three mirror fields have to keep being set and cleared as a group -- the deferral is gated on `_params_pure` and the flush on `_state_pure`, so they are only equivalent for as long as that holds. `_clip_by_grad_norm` moves to its own PR. It is unrelated to deferring the publish, and it is a numerical change rather than a pure optimization -- under `grad_dtype=bfloat16` it clips with the float32 norm where optax clipped with a bf16 one -- so it should be reviewed on its own terms.
|
Pushed 8fac116 with two changes. 1.
|
End-to-end against tunix HEAD, on qwen3.5-35b-a3bRan this branch against a checkout of tunix HEAD Config:
(n=19 after dropping the first 3 steps; H1 mean 4043.9 / min 4040.9 / max 4049.0. Eval Engine vs. PeftTrainer is a dead heat on step time — 0.04% apart in auto, 0.06% in Against the same sweep on the pinned tunix: engine/auto was 4042.0 ms / 61.39 G and The counterfactual, on this modelThis PR's
One thing this surfaced that is not in this PRH3 initially failed outright. Under That is pre-existing on Repro |
1 participant
Description
Removes the last two pieces of redundant per-step work from
MaxTextTrainingEngine.1. The step path no longer walks the live NNX module graph.
fwd_bwdandupdateeach callednnx.updateon the live NNX objects on every step.nnx.updatewalks the module graph, so its cost scales with the graph's size rather than with the state actually written: 6.2 ms to publish 180 leaves of non-parameter state into an unrolled 28-layer qwen3-0.6b, once per micro-batch, and 16.7 ms for the parameters once per update. That was the last per-step graph walk left after_refresh_pure_statemoved the twonnx.splitcalls to once-per-compile.Nothing inside the step path reads those objects back —
_read_model_pureand_read_state_pureanswer from the pure mirror, and the live objects serve only as containers — so the walk is pure publication, for readers outside the engine. So defer it._publish_to_livemarks the mirror as ahead of the live objects, and_sync_live_objectsflushes before anything reads real values off them: themodel/optimizer/stateproperties,_refresh_pure_state, checkpoint save and restore,_get_trainable_params_state, and the throttler's fallback path when there is no gradient norm to wait on._invalidate_pure_stateflushes before forgetting the mirror, so dropping the cache never drops a step's results with it.The sync is not merely cosmetic.
update()donates the state it hands to the update kernel, so a reader that saw a deferred publish would find arrays that are deleted, not merely stale. Publication stays eager whenever the pure state is disabled, since there would be nothing to publish from later.2. The global gradient norm is computed once instead of twice.
_update_kernelcomputedl2norm_pytree(grads)for the metric and then calledmaxtext_utils.apply_gradient_clipping, which recomputes the same norm insideoptax.clip_by_global_norm— a second full pass over every gradient plus a second cross-replica reduction of a scalar._clip_by_grad_normapplies optax's own scale (1under the threshold, elsethreshold / norm) to the norm already in hand.Reusing it is also the more accurate of the two. The engine's norm is computed in float32 for the reason stated at its call site — a sum of squares over bf16 overflows on production-size models — while optax's runs in the gradients' own dtype. Under
grad_dtype=bfloat16the engine was therefore reporting a guarded norm and then clipping with an unguarded one. At the defaultgrad_dtype=float32the two are the same number and this changes nothing numerically. The fp8 path keeps the optax call, which holdsOVERWRITE_WITH_GRADIENTout of both the norm and the scaling; that carve-out is not worth reproducing for the saving.Benchmark
qwen3-0.6b on 4× v6e,
ici_data_parallelism=4,per_device_batch_size=2,max_target_length=1024, adamw, bf16 compute / fp32 weights, synthetic data. Median steady-state step over 9 post-warmup steps, this branch against its merge base onmain:Losses are bit-identical across all 12 steps of every pair.
Two honest null results, and what I think is behind them:
scan_layers=false.InflightThrottler(max_inflight_computations: 2) forces the host to synchronize every couple of micro-batches and the graph walk stops being free. I did not instrument this, so treat it as the probable mechanism rather than a measured one.Change 2 in isolation does not move wall time at this model size, but it is not a no-op. Lowering and compiling the
updatekernel on its own and readingcost_analysis():grad_dtype=float32grad_dtype=float32grad_dtype=bfloat16grad_dtype=bfloat16Real work is removed; it is simply invisible next to eight forward/backward passes on a 0.6B model. It should matter more at low gradient accumulation on larger models, where the update kernel is a bigger share of the step.
Tests
JAX_PLATFORMS=cpumatters: most of these are markedcpu_onlyand silently skip on an accelerator testbed.55 passed on this branch (50 pre-existing, unmodified, plus 5 new); the same 50 pass on
main. New coverage:test_live_objects_are_not_walked_on_the_step_path— a steady-statefwd_bwd+updateissues nonnx.updateagainst the engine's own live objects. This is the whole point of the change, so a regression that quietly reintroduces the walk should fail a test rather than only show up in a profile.test_reading_the_model_flushes_a_deferred_publish— a deferred publish is invisible to.model.test_dropping_the_pure_state_flushes_rather_than_drops_the_step—_invalidate_pure_stateflushes instead of discarding.test_publication_stays_eager_without_a_pure_mirror— the fallback still writes through.test_clip_by_grad_norm_matches_optax— the hand-rolled scale matchesoptax.clip_by_global_normon both sides of the threshold.The benchmark above is the end-to-end evidence; no long-running convergence run was done, on the grounds that the losses are bit-identical.
Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.