Make shard_optimizer_over_data (Zero-1) work in MaxTextTrainingEngine - #5104
Conversation
![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 implements the shard_optimizer_over_data (Zero-1) optimization in MaxTextTrainingEngine. It introduces logic to shard the optimizer's parameter-shaped state over the data axis, performs the optimizer update on these slices, and gathers the updated parameters back. It also fixes a bug in add_data_to_sharding where PartitionSpec was treated as a pytree leaf, preventing duplicate sharding. Additionally, a comprehensive suite of unit and integration tests has been added to verify the correctness, performance characteristics, and checkpoint compatibility of the Zero-1 implementation. There are no review comments, so I have no feedback to provide.
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
The flag was read only by `gradient_accumulation.py`, which the engine does not go through, so setting it here allocated a fully replicated optimizer and said nothing about it. `nnx.Optimizer` builds its moments eagerly as `zeros_like` of each parameter, so they inherit the parameter layout. Moving them onto the data axis once, before `_compile_for_batch` reads the state's shardings back off the arrays, is enough for the rest of the engine to follow -- the update kernel's in/out shardings are derived from exactly those arrays. `_update_kernel` then reshards the gradients and the parameters onto that layout, and gathers the new parameters back on the way out. Nothing else changes: no kernel signature moves, and the parameters cross every jit boundary replicated as before. The gradient reshard is the same one that discharges the deferral's `unreduced` tag, so the two compose -- one cross-replica reduction per optimizer step, on 1/N of the optimizer. `add_data_to_sharding`'s "already present" guard never fired: a PartitionSpec is a pytree leaf, so `jax.tree.leaves(pspec)` gives back the spec. A leaf already sharded over "data" got a second one and `NamedSharding` rejected it (`DuplicateSpecError`), which is what a recompile under Zero-1 would hit. qwen3-0.6b on 4x v6e, dp=4, adamw, micro-batch 8x1024, median steady-state step over ~19 post-warmup steps: GA arm step peak HBM live HBM 8 baseline 593.1ms 12.39 GiB 7.93 GiB 8 defer 437.6ms 12.34 GiB 7.88 GiB 8 zero1 597.3ms 9.08 GiB 4.62 GiB 8 defer + zero1 437.5ms 9.01 GiB 4.55 GiB Zero-1 is time-neutral under accumulation (the all-gather is once per step, amortized over the micro-batches) and returns 3.3 GiB per device. At GA=1 it costs 1.5-5ms, where there is nothing to amortize it over. Losses with Zero-1 alone are bit-identical to the baseline across all 9 steps at every GA -- the optimizer's arithmetic is elementwise, so splitting the tensor changes nothing. The deferral's float32 reassociation accounts for all the divergence in the combined arm (<= 9.3e-5 relative).
7f90995 to
6605358
Compare
| ) | ||
|
|
||
|
|
||
| _ZERO1_DECLINED_WARNING = ( |
There was a problem hiding this comment.
We should error instead of warn when a user setting cannot be applied
| return | ||
| self._params_pure, self._rest_pure, self._state_pure = params_pure, rest_pure, new_state_pure | ||
|
|
||
| def _note_zero1_declined(self, reason: str) -> None: |
|
|
||
| return jax.tree.map(target, params_pure, params_shardings) | ||
|
|
||
| def _shard_optimizer_state_over_data(self) -> None: |
There was a problem hiding this comment.
this function is surprising, I think the optimizer should always be sharded, we only need to construct such a sharding spec? Are we expecting this is a no-op?
There was a problem hiding this comment.
optimizer states never gets sharded by "data". This function enforce sharding optimizer states using "data" axis for zero-1. We have similar function for maxtext pre-train.
| local_state = nnx.merge(self._state_graphdef, state_pure, copy=True) | ||
| if hasattr(local_state, "apply_gradients"): | ||
| if self._config.skip_step_on_spikes: | ||
| local_state.apply_gradients(grads, loss=mean_loss, grad_norm=grad_norm) |
There was a problem hiding this comment.
woah this is very surprising to me, I guess this is existing code but I thought fwd_bwd only was meant to compute gradients, and applying them was dedicated to separate API like update
There was a problem hiding this comment.
indeed fwd_bwd only computes/accumulates gradients in tunix http://google3/third_party/py/tunix/experimental/train/peft_trainer_v2.py;l=586;rcl=974880264
| self._reduced_params_shardings: Any = None | ||
| self._unreduced_grad_shardings: Any = None | ||
| self._plain_grad_shardings: Any = None | ||
| # Set together by `_compile_for_batch` when Zero-1 is on: the parameters as |
There was a problem hiding this comment.
I am very surprised this can be implemented without touching the _compile_for_batch function http://google3/third_party/py/maxtext/src/maxtext/training_engine/maxtext_engine.py;l=738;rcl=974852334
I would think we need two separate functions fwd_bwd and fwd_bwd_unreduced, or at least we need to modify the call signature of fwd_bwd to return unreduced gradients. However I may have the signature of fwd_bwd incorrect, does it not return gradients?
2 participants
Stacked on #5099 (base
engine-ga-unreduced).What
shard_optimizer_over_data(Zero-1) is a silent no-op inMaxTextTrainingEngine. It is read only bygradient_accumulation.py, which the engine does not go through, so setting it allocates a fully replicated optimizer and says nothing about it.This makes it work. It matters here specifically because Zero-1 is what makes #5099 usable: the deferred all-reduce needs
datato be the sole batch axis, which rules out FSDP — and Zero-1 is already mutually exclusive with FSDP by config validation (types.py:4727). So Zero-1 is the memory story you reach for when the whole point of the mesh is data parallelism, and the two compose into one cross-replica reduction per optimizer step, on 1/N of the optimizer.How
nnx.Optimizerallocates its moments eagerly, aszeros_likeof each parameter, so they inherit the parameter layout. Resharding them once — before_compile_for_batchreads the state's layout back off the arrays — is all it takes for the rest of the engine to follow, since the update kernel's in/out shardings are derived from exactly those arrays._update_kernelthen does two reshards in opposite directions:The gradient reshard is the same one that discharges #5099's
unreducedtag, which is why the two cost one collective between them rather than two.Nothing else moves. No kernel signature changes, the parameters cross every jit boundary replicated exactly as before, and the checkpoint format is untouched —
_reduced_accumulated_gradsstill writes a plain replicated total.Placement goes through the pre-train path's own
sharding.add_data_to_sharding, one function of(shape, base sharding), applied to the parameters, the gradients and the moments alike. That is what makes the three agree without matching up two pytrees: a moment mirrors its parameter's shape and starts from its layout, so it lands on the same spec. A leaf with no dimension the data axis divides — a scalarcount, an odd-sized bias — comes back unchanged and stays replicated, on all three trees.Results — qwen3-0.6b on 4× v6e
data=4, fsdp=1, micro-batch8×1024,shard_mode=explicit, adamw (sgd carries no parameter-shaped state, so there would be nothing to shard), no clipping,remat=none. Median steady-state step over ~19 post-warmup steps, untraced; HBM is the max over the 4 devices.Zero-1 is time-neutral under accumulation — its all-gather is once per step and the micro-batches amortize it — and returns 3.3 GiB per device, which is 3/4 of the two f32 moments over 596M parameters, as expected. At GA=1 it costs 1.5–5 ms, where there is nothing to amortize it over.
The two are independent and stack: the deferral buys the time, Zero-1 buys the memory, and neither takes back what the other gave.
Losses with Zero-1 alone are bit-identical to the baseline at every GA — every operation
adamwapplies is elementwise in the parameter, so splitting the tensor across replicas changes nothing. All the divergence in the combined arm is #5099's float32 reassociation (≤ 9.3e-5 relative, max over GA ∈ {1, 4, 8}).Collectives, from the optimized HLO on the CPU mesh:
first_kernel/accum_kernel(per micro-batch)_update_kernel(per step)Results — Zero-1 + explicit + GA=8 vs. pure tensor parallelism
Same rig, same 4× v6e, GA=8 throughout, global batch
64×1024.per_device_batch_sizeis8/4 = 2in every arm, so the micro-batch is8×1024no matter how the mesh is cut and all seven arms do identical work.…-zero1-ga8-dp4fsdp1tp1-explicit/…/2026_09_02_16_36_51/…-ga8-dp4fsdp1tp1-explicit/…/2026_09_02_16_39_09/…-ga8-dp4fsdp1tp1-explicit-nodefer/…/2026_09_02_16_41_37/…-ga8-dp1fsdp1tp4-explicit/…/2026_09_02_16_44_20/…-ga8-dp1fsdp1tp4/…/2026_09_02_16_46_46/…-zero1-ga8-dp2fsdp1tp2-explicit/…/2026_09_02_16_34_07/…-ga8-dp2fsdp1tp2-explicit/…/2026_09_02_16_31_11/Trace paths are relative to
gs://chengnuojin-xprof/engine-zero1-vs-tp/, and each ends inplugins/profile/<timestamp>/t1v-n-76d392d5-w-0.xplane.pb(the…above elidesqwen3-0.6b-engine-adamwat the front andplugins/profilein the middle). Point xprof at the arm directory.Neither strategy dominates: Zero-1 + DP is 2.0× faster, pure TP holds 1.7× less live HBM. The per-kernel medians say where it comes from — TP's
fwd_bwdis 97.4 ms per micro-batch against Zero-1's 43.3 ms, because TP all-reduces activations inside every layer of every micro-batch and accumulation cannot amortize that;update()is much closer (139.6 vs 80.9 ms). At this size TP is buying memory nobody needed: 596M f32 parameters is 2.2 GiB, and the whole Zero-1 arm fits in 4.55 GiB live against v6e's 32 GiB.Hybrid dp2 × tp2 is the worst of the three. It pays TP's in-layer collectives and collects only half of Zero-1's saving (1.09 GiB live, exactly the 2-way shard of the moments), landing at 85k tok/s.
Numerics hold across the two strategies: pure TP and defer + zero1 track to 8.7e-5 max relative difference over 23 steps (247.0099 → 169.1138 vs. 247.0115 → 169.1284), which is float32 reassociation between two completely different shardings.
Read
livefor the memory comparison, notpeak:peak_bytes_in_useis a high-water mark since process start and so includes compilation transients that vary run to run (the traced re-runs above report the sameliveto the centibyte and 1.2–2.2 GiB lesspeak). The Zero-1 saving is exactly 3.33 GiB oflivein both, which is 3/4 of the two f32 moments over 596M parameters.Timings are from untraced runs; the traces are separate runs of the same arms, 5–8% slower for being watched.
Dependency: #5099 needed a fix to survive this
--dp 2 --tp 2did not run at all at first. The deferral's gate refused a second mesh axis on the batch dimension, which caughtfsdpand missedtensorreaching the same contradiction through the feature dimension:Not a Zero-1 bug — it reproduced with
--zero1off — so the fix is a commit on the base branch (#5099), which now declines on any non-datamesh axis above size 1 rather than enumerating the axes known to break. Zero-1 itself is fine alongside TP, as the two dp2 × tp2 rows show.A real bug in
add_data_to_shardingA
PartitionSpecis a pytree leaf, so flattening one gives back the spec itself and this guard never fired. A leaf already sharded overdatagot a second one, andNamedShardingrejects the result outright:That is exactly what a recompile under Zero-1 hits —
_compile_for_batchre-places the moments, which are already sharded by then. Fixed to walktuple(pspec). The pre-train path has not hit it because its parameters are sharded overfsdp/tensorrather thandata.When it engages
_zero1_activereturns a reason, orNonewhen Zero-1 can run. It declines unlessshard_mode=explicit, the mesh's axis types are allExplicit, and the mesh has adataaxis of size > 1 — underautothe reshards are hints GSPMD may ignore, which would give a silently replicated optimizer again, i.e. the bug this fixes.When it declines something that was asked for, it now says so, once per engine instance. The failure mode this replaces was silence.
Tests
tests/post_training/unit/maxtext_engine_zero1_test.py, 23 cases on a 4-device CPU mesh, reusing #5099's tiny-real-decoder rig so both features are exercised on one config:_zero1_shardingplaces one leaf: adds the axis to the first dimension that divides, skips one that does not, leaves scalars and unshardable shapes alone, and leaves a leaf already sharded overdataalone (this one fails without theadd_data_to_shardingfix);update()and none to either micro-batch kernel, measured as a difference against the same kernel compiled without it;update();l2norm_pytreesums squares over elements that are now spread across replicas, and a replica-local sum would come out low by a factor of N;in_shardingsmismatch.Existing
maxtext_engine_deferred_all_reduce_test.py+maxtext_engine_test.py+maxtext_engine_constructor_test.py+sharding_nnx_test.py: 92 passed.Follow-ups, not in scope here
op_name="jit(train_step)/while/body/…"). Thereduced/unreducedscaffolding ingradient_accumulation.py:73-86is dead — both arms of thedata_parallel_activebranch assign the same thing, because scan carries reject those specs. So the engine with Defer the data-parallel gradient all-reduce to update() under gradient accumulation #5099 + this PR is currently the only place in MaxText where that reduction is actually deferred.--zero1and--optadded to theperf_parityrig from Match Tunix peft_trainer_v2 performance in MaxTextTrainingEngine #5060, which is not on this branch. The command waspython qwen3_engine_profile.py --ga 8 --dp 4 --fsdp 1 --tp 1 --shard-mode explicit --opt adamw --no-trace [--zero1] [--no-defer], varying--dp/--fsdp/--tpfor the mesh comparison and dropping--no-tracefor the xplane files.