From e92abaad61c777fe582fe8c7b426bdc03e9ee65a Mon Sep 17 00:00:00 2001 From: NuojCheng Date: Wed, 2 Sep 2026 06:47:00 +0000 Subject: [PATCH 1/4] Defer the data-parallel gradient all-reduce to update() under GA MaxTextTrainingEngine all-reduced the whole gradient tree across data replicas once per *micro*-batch. Only the sum matters, so at gradient_accumulation N it paid N times for a reduction that is correct once. Tag the parameters `reduced` over the data axis where `value_and_grad` differentiates them, and their cotangents come out `unreduced`: a per-replica partial that accumulates locally across micro-batches. `update()` reshards back to a plain spec, which is what emits the single all-reduce -- before the division, the norm and the optimizer, so nothing downstream has to know about the tag. This is what gradient_accumulation.py already does for the pre-train path, applied across the engine's separate jax.jit dispatches rather than inside one lax.scan. Gated to explicit sharding on an all-Explicit mesh where "data" is the only batch axis of size > 1. The last condition is not conservatism: with fsdp on the batch too, JAX rejects the backward pass, because the unreduced set has to be exactly the contracted axes and widening it to fsdp collides with the parameters being sharded there. Two places had to learn that gradients can be tagged. RMSNorm's scale alignment indexed `spec[...]`, which a tagged spec refuses -- it reads `spec.partitions` now, as the pre-train path's does. And an accumulator can outlive the shardings it was built under: a checkpoint holds the reduced total (Orbax cannot serialize an unreduced array at all), and a recompile can flip the deferral, so both restore and recompile move it back onto whatever the kernels now expect. qwen3-0.6b, 4x v6e, data=4 fsdp=1, micro-batch 8x1024, median steady-state step: ga=8 584.3ms -> 428.3ms (1.36x) ga=4 295.4ms -> 229.8ms (1.29x) ga=1 80.2ms -> 80.1ms (unchanged, as it should be) The optimized HLO says the same thing exactly: 596M f32 elements all-reduced per micro-batch became 596M once per optimizer step, and the micro-batch kernels are left with two scalars, the loss and its denominator. --- src/maxtext/layers/normalizations.py | 12 +- src/maxtext/training_engine/maxtext_engine.py | 154 ++++++- src/maxtext/utils/sharding.py | 8 + ...maxtext_engine_deferred_all_reduce_test.py | 390 ++++++++++++++++++ 4 files changed, 556 insertions(+), 8 deletions(-) create mode 100644 tests/post_training/unit/maxtext_engine_deferred_all_reduce_test.py diff --git a/src/maxtext/layers/normalizations.py b/src/maxtext/layers/normalizations.py index 987c06fa42..f2d4dfe5cb 100644 --- a/src/maxtext/layers/normalizations.py +++ b/src/maxtext/layers/normalizations.py @@ -44,11 +44,17 @@ def _align_scale_with_normalized_axis(scale: jnp.ndarray, y: jnp.ndarray) -> jnp This is a no-op whenever the two already agree, which includes every auto-sharding-equivalent layout for the ordinary layer norms. """ - activation_spec = jax.typeof(y).sharding.spec + # Read through `.partitions`: a scale carrying a reduced/unreduced tag (the deferred + # data-parallel all-reduce under gradient accumulation) rejects direct indexing. The tags + # are carried over to the new spec so the scale keeps its place on the deferred path. + activation_axis = jax.typeof(y).sharding.spec.partitions[-1] scale_spec = jax.typeof(scale).sharding.spec - if scale_spec[-1] == activation_spec[-1]: + if scale_spec.partitions[-1] == activation_axis: return scale - return jax.sharding.reshard(scale, jax.sharding.PartitionSpec(activation_spec[-1])) + return jax.sharding.reshard( + scale, + jax.sharding.PartitionSpec(activation_axis, unreduced=scale_spec.unreduced, reduced=scale_spec.reduced), + ) class RMSNorm(nnx.Module): diff --git a/src/maxtext/training_engine/maxtext_engine.py b/src/maxtext/training_engine/maxtext_engine.py index e397b3f4a5..815be224f4 100644 --- a/src/maxtext/training_engine/maxtext_engine.py +++ b/src/maxtext/training_engine/maxtext_engine.py @@ -157,6 +157,81 @@ def _batch_signature(dynamic_batch: Any, static_batch: dict[str, Any]) -> Any: "every step, so this is worth reporting rather than living with." ) +# The mesh axis a data-parallel gradient all-reduce runs over. Only this one is ever tagged +# `reduced`/`unreduced`; see `_deferred_all_reduce_shardings`. +_DATA_AXIS = "data" + + +def _tag_sharding(named_sharding: jax.sharding.NamedSharding, field: str) -> jax.sharding.NamedSharding: + """Marks a sharding `reduced` or `unreduced` over the data axis. + + A tensor already sharded over that axis is returned untouched: it holds no cross-replica + partial to defer, and JAX rejects a spec that both shards and reduces over one axis. + """ + if _DATA_AXIS in sharding.mesh_axes_for_dim(named_sharding.spec.partitions): + return named_sharding + return named_sharding.update(spec=named_sharding.spec.update(**{field: {_DATA_AXIS}})) + + +def _deferred_all_reduce_shardings(config: Any, mesh: Any, params_shardings: Any) -> tuple[Any, Any]: + """Returns `(reduced params, unreduced gradients)` sharding trees, or `(None, None)`. + + Tagging the parameters that are differentiated `reduced` over the data axis makes their + cotangents come out `unreduced`: each replica then holds a partial sum that accumulates + locally across micro-batches, and the cross-replica all-reduce runs once per optimizer + step instead of once per micro-batch. This is the same trick + `gradient_accumulation.py` plays for the pre-train path, applied across the engine's + separate `jax.jit` dispatches rather than inside one `jax.lax.scan`. + + `(None, None)` -- the untagged status quo -- whenever the tag would be unsound: + + - not explicit sharding, where reduced/unreduced specs do not exist; + - a mesh with any non-Explicit axis, which those specs are also rejected on. A caller can + hand the engine an all-Auto mesh regardless of `config.shard_mode`; + - "data" is not the only mesh axis of size > 1 that the activation batch dimension is + sharded over. A gradient contracts over the batch, and JAX requires the unreduced set + to be exactly the contracted axes -- with `data` and `fsdp` both on the batch it + rejects the backward pass outright ("unreduced axes should be equal to the contracting + specs. Got unreduced axes=frozenset({'data'}) and contracting spec=(('data', 'fsdp'), + None)"). Widening the tag to both is not the fix: a parameter sharded over `fsdp` + cannot also be unreduced over it. Read the batch axes off the resolved mesh rather + than `config.ici_data_parallelism`, which may still be -1 (auto-fill). + """ + if getattr(config, "shard_mode", None) != common_types.ShardMode.EXPLICIT: + return None, None + if mesh is None or mesh.shape.get(_DATA_AXIS, 1) <= 1: + return None, None + if any(axis_type != jax.sharding.AxisType.Explicit for axis_type in mesh.axis_types): + return None, None + try: + batch_axes = sharding.batch_mesh_axes(mesh, rules=config.logical_axis_rules) + except (KeyError, ValueError, IndexError): + # No usable "activation_batch" rule for this mesh: leave the gradients untagged. + return None, None + if batch_axes != frozenset({_DATA_AXIS}): + return None, None + return ( + jax.tree.map(lambda s: _tag_sharding(s, "reduced"), params_shardings), + jax.tree.map(lambda s: _tag_sharding(s, "unreduced"), params_shardings), + ) + + +def _conform_accumulator(value: Any, target: jax.sharding.NamedSharding) -> Any: + """Moves one accumulated-gradient leaf onto `target`, preserving the value it represents. + + Only ever needed when the accumulator outlives the shardings it was produced under: a + checkpoint restore hands back the summed gradient, and a recompile may turn the deferred + all-reduce on or off. Both directions are exact -- resharding away from `unreduced` runs + the all-reduce, and `device_put` onto it keeps the value on one data replica and zeroes + the others, so the pending all-reduce reproduces it. + """ + current = getattr(value, "sharding", None) + if current == target: + return value + if getattr(current, "spec", None) is not None and current.spec.unreduced: + return jax.sharding.reshard(value, target) + return jax.device_put(value, target) + @dataclasses.dataclass(kw_only=True) class RouterReplayTrainerPayload(abstract_engine.TrainerPayload): @@ -381,6 +456,13 @@ def __init__( # Summed loss denominators behind `_accumulated_grads`, which are unreduced: this is the # divisor `update()` applies once. self._accumulated_denominator: Any = None + # Set together by `_compile_for_batch` when the data-parallel all-reduce can be deferred + # to the optimizer step; all three `None` selects the untagged path. The parameters as + # `_fwd_bwd_kernel` differentiates them, the gradients as they cross every kernel + # boundary, and the gradients once reduced. See `_deferred_all_reduce_shardings`. + self._reduced_params_shardings: Any = None + self._unreduced_grad_shardings: Any = None + self._plain_grad_shardings: Any = None self._micro_step_count = 0 # Set when this run resumed from an intra-step checkpoint, cleared once the step it # resumed into completes and its finished state has been checkpointed. @@ -721,6 +803,14 @@ def diff_wrapper(p, r, b): "or a 2-element tuple/list: (loss, aux_metrics)." ) + if self._reduced_params_shardings is not None: + # Tag the differentiated parameters `reduced` over the data axis, so their cotangents + # come out `unreduced` and the accumulation below stays replica-local -- the + # cross-replica all-reduce then runs once, in `_update_kernel`. Deliberately outside + # `diff_wrapper`: autodiff transposes a reshard, so the same call one line further in + # would put an all-reduce back into every micro-batch. + params = jax.tree.map(jax.sharding.reshard, params, self._reduced_params_shardings) + grad_func = jax.value_and_grad(diff_wrapper, argnums=0, has_aux=True) # Every non-raising branch of `diff_wrapper` builds a LossOutput, so `loss_out` is always # one. The value returned by `value_and_grad` is the unreduced sum that was @@ -756,6 +846,13 @@ def _update_kernel(self, state_pure, accumulated_grads, accumulated_denominator, grad_norm = None is_skipped_val = None if state_pure is not None: + if self._plain_grad_shardings is not None: + # The gradients arrive `unreduced`: a per-replica partial sum over this step's + # micro-batches. Resharding them back is what emits the single cross-replica + # all-reduce that replaces the one every micro-batch used to pay. First, so that + # everything below -- the division, the norm, clipping, the optimizer -- sees + # ordinary gradients and needs no tag handling of its own. + accumulated_grads = jax.tree.map(jax.sharding.reshard, accumulated_grads, self._plain_grad_shardings) # This one division is the whole normalization. A zero total means every micro-batch # was empty; yield zeros rather than a NaN, as `gradient_accumulation.py` does. has_weights = accumulated_denominator > 0 @@ -919,17 +1016,41 @@ def accum_kernel(params, rest, dynamic, acc_grads, acc_denom): params_shardings = jax.tree.map(self._mesh_sharding, params_pure) rest_shardings = jax.tree.map(self._mesh_sharding, rest_pure) batch_shardings = self._batch_data_shardings(dynamic_batch) + # When the data-parallel all-reduce can be deferred, the gradients live on their own + # shardings -- `params_shardings` plus an `unreduced` tag -- everywhere they cross a + # jit boundary: out of both fwd/bwd kernels, back into the accumulating one, and into + # the update. `params_shardings` stays untagged, so the weights themselves are + # unaffected; `_fwd_bwd_kernel` applies the matching `reduced` tag inside. + self._reduced_params_shardings, self._unreduced_grad_shardings = _deferred_all_reduce_shardings( + self._config, self._mesh, params_shardings + ) + grad_shardings = self._unreduced_grad_shardings + if grad_shardings is None: + grad_shardings = params_shardings + self._plain_grad_shardings = None + else: + self._plain_grad_shardings = params_shardings first_in_shardings = (params_shardings, rest_shardings, batch_shardings) - accum_in_shardings = first_in_shardings + (params_shardings, replicated) - fwd_bwd_out_shardings = (None, None, rest_shardings, params_shardings, replicated) - update_in_shardings = (state_mesh_shardings, params_shardings, replicated, None) + accum_in_shardings = first_in_shardings + (grad_shardings, replicated) + fwd_bwd_out_shardings = (None, None, rest_shardings, grad_shardings, replicated) + update_in_shardings = (state_mesh_shardings, grad_shardings, replicated, None) update_out_shardings = (state_mesh_shardings, None, None) + # A live accumulator predates this compile -- a checkpoint restore hands one back, and + # a recompile can flip the deferral on or off -- so it may not be on the shardings the + # kernels were just built for. `jax.jit` matches `in_shardings` exactly and would + # reject it. + if self._accumulated_grads is not None: + with self._sharding_ctx(): + self._accumulated_grads = jax.tree.map(_conform_accumulator, self._accumulated_grads, grad_shardings) else: first_in_shardings = None accum_in_shardings = None fwd_bwd_out_shardings = None update_in_shardings = None update_out_shardings = None + self._reduced_params_shardings = None + self._unreduced_grad_shardings = None + self._plain_grad_shardings = None # 1. JIT Compile Micro FWD/BWD Pass. # @@ -1168,6 +1289,21 @@ def eval_step(self, payload: abstract_engine.TrainerPayload, **kwargs: Any) -> N "Logged once per engine instance." ) + def _reduced_accumulated_grads(self) -> Any: + """Returns the accumulated gradients in the form Orbax can serialize. + + While the data-parallel all-reduce is deferred they are held `unreduced` -- a + per-replica partial sum -- which Orbax cannot write (`device_indices_map` is undefined + for one) and which would not be a meaningful thing to write anyway. Resharding runs the + all-reduce the pending `update()` would have run, so the checkpoint holds exactly the + total that step will apply. `_compile_for_batch` puts a restored total back on the + accumulator's shardings. + """ + if self._accumulated_grads is None or self._plain_grad_shardings is None: + return self._accumulated_grads + with self._sharding_ctx(): + return jax.tree.map(jax.sharding.reshard, self._accumulated_grads, self._plain_grad_shardings) + def save_checkpoint(self, metadata: Any, **kwargs: Any) -> None: """Forces asynchronous Orbax checkpoint serialization. @@ -1212,7 +1348,7 @@ def save_checkpoint(self, metadata: Any, **kwargs: Any) -> None: # The full history, not `get_metrics()`: CheckpointState.accumulated_metrics is # a list, and restore_checkpoint iterates it back into the recorder's buffer. accumulated_metrics=self._metrics_recorder.get_metrics_history(clear_cache=False), - accumulated_grads=self._accumulated_grads, + accumulated_grads=self._reduced_accumulated_grads(), # Recorded by the CheckpointManager into custom_metadata, so that a later save # at this same step can tell it supersedes this one. micro_step_count=self._micro_step_count, @@ -1236,7 +1372,7 @@ def restore_checkpoint(self, **kwargs: Any) -> Any: checkpoint_state = checkpointing.CheckpointState( model=self.model, optimizer=self.optimizer, - accumulated_grads=self._accumulated_grads, + accumulated_grads=self._reduced_accumulated_grads(), ) restored_step, restored_checkpoint_state, restored_metadata = self._checkpoint_manager.restore_checkpoint( @@ -1305,6 +1441,14 @@ def restore_checkpoint(self, **kwargs: Any) -> Any: # above, which the branch above has already discarded. if self._micro_step_count > 0 and restored_checkpoint_state.accumulated_grads: self._accumulated_grads = restored_checkpoint_state.accumulated_grads + if self._unreduced_grad_shardings is not None: + # What was saved is the reduced total; what the already-compiled kernels take is an + # unreduced partial. Without this the resumed step dies on an `in_shardings` + # mismatch, since restoring does not recompile -- the batch shape has not changed. + with self._sharding_ctx(): + self._accumulated_grads = jax.tree.map( + _conform_accumulator, self._accumulated_grads, self._unreduced_grad_shardings + ) self._accumulated_denominator = jnp.float32(restored_denominator if restored_denominator else 0.0) rebuilt_losses = None diff --git a/src/maxtext/utils/sharding.py b/src/maxtext/utils/sharding.py index a1e320603f..9574a5a4b2 100644 --- a/src/maxtext/utils/sharding.py +++ b/src/maxtext/utils/sharding.py @@ -209,6 +209,14 @@ def mesh_axes_for_dim(axis_names): return tuple(axis for axis in axis_names if axis is not None) +def batch_mesh_axes(mesh, rules=None): + """Returns the mesh axes of size > 1 that the activation batch dimension is sharded over.""" + spec = logical_to_mesh_axes(("activation_batch",), mesh, rules=rules) + if spec is None: + return frozenset() + return frozenset(axis for axis in mesh_axes_for_dim(spec.partitions[0]) if mesh.shape.get(axis, 1) > 1) + + def mesh_axes_size(mesh, axes, *, label): """Returns the product of mesh sizes for a set of axes.""" size = 1 diff --git a/tests/post_training/unit/maxtext_engine_deferred_all_reduce_test.py b/tests/post_training/unit/maxtext_engine_deferred_all_reduce_test.py new file mode 100644 index 0000000000..5c66e28b82 --- /dev/null +++ b/tests/post_training/unit/maxtext_engine_deferred_all_reduce_test.py @@ -0,0 +1,390 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""The engine's deferred data-parallel gradient all-reduce. + +Under gradient accumulation the engine used to pay one cross-replica all-reduce of the +whole parameter tree per *micro*-batch. Tagging the differentiated parameters `reduced` +over the data axis makes their cotangents `unreduced`, so the accumulation stays +replica-local and the all-reduce happens once, in `update()`. + +That is a pure performance change, which makes it exactly the kind that can rot into a +no-op without anything failing. So the tests here assert on the compiled HLO -- where the +collectives actually are -- and every such assertion carries a vacuity guard proving the +same probe finds the all-reduce it is looking for when the deferral is off. +""" + +import os + +# Must precede the first JAX import: a data-parallel mesh needs more than one device, and +# the CPU backend reads this only at initialization. +os.environ.setdefault("XLA_FLAGS", "--xla_force_host_platform_device_count=4") + +import re # pylint: disable=wrong-import-position +import unittest # pylint: disable=wrong-import-position +from unittest import mock # pylint: disable=wrong-import-position + +from absl.testing import absltest # pylint: disable=wrong-import-position +from flax import nnx # pylint: disable=wrong-import-position +import jax # pylint: disable=wrong-import-position +from maxtext.configs import pyconfig # pylint: disable=wrong-import-position +from maxtext.training_engine import maxtext_engine # pylint: disable=wrong-import-position +from maxtext.utils import maxtext_utils # pylint: disable=wrong-import-position +import numpy as np # pylint: disable=wrong-import-position +import pytest # pylint: disable=wrong-import-position +from tests.utils.test_helpers import get_test_config_path # pylint: disable=wrong-import-position + +# training_engine imports tunix, so these tests need the post-training dependency bundle. +pytestmark = [pytest.mark.post_training] + +_REQUIRED_DEVICES = 4 + + +def _config(**overrides) -> pyconfig.HyperParameters: + """A tiny real model on an explicit, purely data-parallel mesh. + + Small enough to compile in seconds, but a *real* MaxText decoder rather than a stub: the + reduced/unreduced tags have to survive every layer that touches a parameter, and it was a + norm layer indexing `spec[...]` directly that broke first. + """ + argv = [ + "maxtext_engine_deferred_all_reduce_test.py", + get_test_config_path("base.yml"), + "model_name=default", + "run_name=engine_deferred_all_reduce_test", + "enable_checkpointing=False", + "convert_checkpoint_if_possible=False", + "skip_jax_distributed_system=True", + "enable_tensorboard=False", + "record_internal_nn_metrics=False", + "enable_dropout=False", + "init_weights_seed=0", + "dtype=float32", + "weight_dtype=float32", + "grad_dtype=float32", + "remat_policy=none", + "scan_layers=False", + "attention=dot_product", + # The tag only goes on when "data" is the sole batch axis of size > 1. + "shard_mode=explicit", + f"ici_data_parallelism={_REQUIRED_DEVICES}", + "ici_fsdp_parallelism=1", + "ici_tensor_parallelism=1", + "per_device_batch_size=1", + # Tiny model. These must be the `base_*` names: emb_dim and mlp_dim are derived. + "vocab_size=128", + "base_emb_dim=64", + "base_mlp_dim=128", + "base_num_decoder_layers=2", + "base_num_query_heads=4", + "base_num_kv_heads=4", + "head_dim=16", + "max_target_length=32", + # A constant schedule with no clipping, so `update()` is the optimizer and nothing else. + "opt_type=sgd", + "learning_rate=1e-2", + "gradient_clipping_threshold=0.0", + "warmup_steps_fraction=0.0", + "learning_rate_final_fraction=1.0", + "gradient_accumulation_steps=1", + ] + argv.extend(f"{k}={v}" for k, v in overrides.items()) + return pyconfig.initialize(argv) + + +def _batch(cfg: pyconfig.HyperParameters, seed: int) -> dict[str, np.ndarray]: + """Minimal batch shaped for `maxtext.trainers.pre_train.train.loss_fn`. + + NumPy, not `jnp`: the compiled kernel takes its batch on `P("data", None)`, and a + committed device array built here would arrive on `P()` and be rejected by `jax.jit`'s + exact `in_shardings` match. Host arrays are placed by the kernel itself. + """ + batch, seq = int(cfg.micro_batch_size_to_train_on), cfg.max_target_length + rng = np.random.default_rng(seed) + tokens = rng.integers(1, cfg.vocab_size, size=(batch, seq)).astype(np.int32) + positions = np.tile(np.arange(seq, dtype=np.int32), (batch, 1)) + segmentation = np.ones((batch, seq), dtype=np.int32) + return { + "inputs": tokens, + "targets": np.roll(tokens, -1, axis=-1), + "inputs_position": positions, + "inputs_segmentation": segmentation, + "targets_segmentation": segmentation, + } + + +def _no_deferral(): + """Patches the gate shut, leaving everything else about the engine identical.""" + return mock.patch.object(maxtext_engine, "_deferred_all_reduce_shardings", return_value=(None, None)) + + +# Matches the result shape of an all-reduce instruction in optimized HLO, covering both the +# fused `%x = f32[64]{0} all-reduce(...)` form and the tupled `ROOT %y = (f32[], f32[64]{0}) +# all-reduce-start(...)` one. +_ALL_REDUCE = re.compile(r"=\s*(.+?)\s+all-reduce(?:-start|-done)?\(") +# A shape with any dimension at all, i.e. not the `f32[]` of a scalar loss term. +_NON_SCALAR = re.compile(r"\[\s*\d") + + +def _array_all_reduces(hlo: str) -> list[str]: + """The result shapes of every all-reduce in `hlo` that moves more than a scalar. + + Scalars are ignored on purpose: the loss and its denominator are reduced across replicas + every micro-batch and always will be. What the deferral is about is the parameter-sized + traffic, which is four orders of magnitude larger even on the toy model here. + """ + shapes = [m.group(1) for line in hlo.splitlines() if (m := _ALL_REDUCE.search(line))] + return [s for s in shapes if _NON_SCALAR.search(s)] + + +class _KernelHlo: + """Captures the arguments the engine passes its jitted kernels, to re-lower them. + + `jax.jit` keeps no handle on the executable it cached, so the only way to read a kernel's + optimized HLO is to lower it again. Recording `ShapeDtypeStruct`s rather than the arrays + themselves keeps that independent of donation -- `_compiled_update` donates its state, so + by the time a test asks for the HLO those buffers are gone. + """ + + def __init__(self, engine, attr: str): + self._jitted = getattr(engine, attr) + self._avals = None + setattr(engine, attr, self._spy) + + def _spy(self, *args): + if self._avals is None: + # `sharding=None` for the host-side batch arrays, which have none; `jax.jit` places + # those from its own `in_shardings` either way. + self._avals = jax.tree.map( + lambda x: jax.ShapeDtypeStruct(x.shape, x.dtype, sharding=getattr(x, "sharding", None)), + args, + ) + return self._jitted(*args) + + def text(self) -> str: + if self._avals is None: + raise AssertionError("kernel was never called, so there is no HLO to read") + return self._jitted.lower(*self._avals).compile().as_text() + + +@unittest.skipIf( + jax.device_count() < _REQUIRED_DEVICES, + f"needs {_REQUIRED_DEVICES} devices; set XLA_FLAGS=--xla_force_host_platform_device_count={_REQUIRED_DEVICES}", +) +class DeferredAllReduceGateTest(absltest.TestCase): + """`_deferred_all_reduce_shardings` decides when the tag is legal. It must decline widely.""" + + def _params_shardings(self, mesh): + return {"w": jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec())} + + def test_tags_on_a_purely_data_parallel_explicit_mesh(self): + cfg = _config() + mesh = maxtext_utils.get_mesh_from_config(cfg) + reduced, unreduced = maxtext_engine._deferred_all_reduce_shardings( # pylint: disable=protected-access + cfg, mesh, self._params_shardings(mesh) + ) + + self.assertEqual(reduced["w"].spec.reduced, {"data"}) + self.assertEqual(unreduced["w"].spec.unreduced, {"data"}) + + def test_declines_under_auto_shard_mode(self): + cfg = _config(shard_mode="auto") + mesh = maxtext_utils.get_mesh_from_config(cfg) + + self.assertEqual( + (None, None), + maxtext_engine._deferred_all_reduce_shardings(cfg, mesh, self._params_shardings(mesh)), # pylint: disable=protected-access + ) + + def test_declines_on_an_auto_axis_mesh_even_in_explicit_mode(self): + """A caller can hand the engine a mesh built by bare `jax.sharding.Mesh(...)`. + + `shard_mode=explicit` then says one thing and the mesh another, and the reduced/unreduced + specs are rejected on Auto axes. The mesh wins. + """ + cfg = _config() + auto_mesh = jax.sharding.Mesh(maxtext_utils.create_device_mesh(cfg), cfg.mesh_axes) + + self.assertEqual( + (None, None), + maxtext_engine._deferred_all_reduce_shardings(cfg, auto_mesh, self._params_shardings(auto_mesh)), # pylint: disable=protected-access + ) + + def test_declines_when_there_are_no_data_replicas(self): + cfg = _config(ici_data_parallelism=1, ici_fsdp_parallelism=_REQUIRED_DEVICES) + mesh = maxtext_utils.get_mesh_from_config(cfg) + + self.assertEqual( + (None, None), + maxtext_engine._deferred_all_reduce_shardings(cfg, mesh, self._params_shardings(mesh)), # pylint: disable=protected-access + ) + + def test_declines_when_fsdp_also_shards_the_batch(self): + """Not conservatism -- JAX rejects the backward pass outright. + + A gradient contracts over the batch, and the unreduced set has to be exactly the + contracted axes: "unreduced axes should be equal to the contracting specs. Got unreduced + axes=frozenset({'data'}) and contracting spec=(('data', 'fsdp'), None)". Widening the tag + to `fsdp` is not available either, since the parameters are sharded over it. + """ + cfg = _config(ici_data_parallelism=2, ici_fsdp_parallelism=2) + mesh = maxtext_utils.get_mesh_from_config(cfg) + + self.assertEqual( + (None, None), + maxtext_engine._deferred_all_reduce_shardings(cfg, mesh, self._params_shardings(mesh)), # pylint: disable=protected-access + ) + + def test_a_parameter_already_sharded_over_data_is_left_alone(self): + """`reduced` and a shard over the same axis are contradictory, and JAX says so.""" + cfg = _config() + mesh = maxtext_utils.get_mesh_from_config(cfg) + shardings = {"w": jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec("data"))} + + reduced, unreduced = maxtext_engine._deferred_all_reduce_shardings(cfg, mesh, shardings) # pylint: disable=protected-access + + self.assertEmpty(reduced["w"].spec.reduced) + self.assertEmpty(unreduced["w"].spec.unreduced) + + +@pytest.mark.integration_test +@unittest.skipIf( + jax.device_count() < _REQUIRED_DEVICES, + f"needs {_REQUIRED_DEVICES} devices; set XLA_FLAGS=--xla_force_host_platform_device_count={_REQUIRED_DEVICES}", +) +class DeferredAllReduceTest(absltest.TestCase): + """End to end on a real decoder: where the collectives land, and what the weights do.""" + + def _run(self, micro_batches: int, steps: int = 2): + """Runs `steps` optimizer steps of `micro_batches` each; returns (engine, hlo probes).""" + cfg = _config() + mesh = maxtext_utils.get_mesh_from_config(cfg) + with jax.set_mesh(mesh): + engine = maxtext_engine.MaxTextTrainingEngine(cfg, mesh=mesh) + engine.compile(_batch(cfg, 0)) + probes = { + "fwd_bwd": _KernelHlo(engine, "_compiled_fwd_bwd"), + "accum": _KernelHlo(engine, "_compiled_fwd_bwd_accum"), + "update": _KernelHlo(engine, "_compiled_update"), + } + for step in range(steps): + for micro in range(micro_batches): + engine.fwd_bwd(_batch(cfg, step * micro_batches + micro)) + engine.update() + return engine, probes + + def test_the_gate_opens_on_this_configuration(self): + """Guards every other test in this class: without this they would all pass vacuously.""" + engine, _ = self._run(micro_batches=1, steps=1) + + self.assertIsNotNone( + engine._plain_grad_shardings, # pylint: disable=protected-access + "the deferral never engaged, so nothing below is testing it", + ) + + def test_micro_batch_kernels_move_only_scalars_across_replicas(self): + """The whole point: no parameter-sized all-reduce per micro-batch, only per step.""" + with jax.set_mesh(maxtext_utils.get_mesh_from_config(_config())): + _, probes = self._run(micro_batches=2, steps=1) + first = _array_all_reduces(probes["fwd_bwd"].text()) + accum = _array_all_reduces(probes["accum"].text()) + update = _array_all_reduces(probes["update"].text()) + + self.assertEmpty(first, f"first micro-batch still all-reduces arrays: {first}") + self.assertEmpty(accum, f"accumulating micro-batches still all-reduce arrays: {accum}") + # Vacuity guard from the other side: the traffic did not vanish, it moved. + self.assertNotEmpty(update, "no array all-reduce in update() either -- the gradients are never reduced") + + def test_without_the_deferral_every_micro_batch_pays(self): + """Proves the probe above can fail. Same model, same probe, tag withheld.""" + with _no_deferral(), jax.set_mesh(maxtext_utils.get_mesh_from_config(_config())): + engine, probes = self._run(micro_batches=2, steps=1) + first = _array_all_reduces(probes["fwd_bwd"].text()) + accum = _array_all_reduces(probes["accum"].text()) + update = _array_all_reduces(probes["update"].text()) + + self.assertIsNone(engine._plain_grad_shardings) # pylint: disable=protected-access + self.assertNotEmpty(first, "baseline should all-reduce the gradients in the first micro-batch") + self.assertNotEmpty(accum, "baseline should all-reduce the gradients in every micro-batch") + self.assertEmpty(update, f"baseline should have nothing left to reduce in update(): {update}") + + def test_deferring_does_not_change_the_weights(self): + """Same sum, different association order. + + Deferring does not drop or duplicate a term -- it moves the cross-replica addition from + before the micro-batch sum to after it. In exact arithmetic the two agree identically, + and on this CPU mesh they do; in float32 on a real accelerator the reassociation shows up + around 1e-6 relative, which is what the tolerance here allows for. + """ + with jax.set_mesh(maxtext_utils.get_mesh_from_config(_config())): + deferred, _ = self._run(micro_batches=3, steps=2) + with _no_deferral(), jax.set_mesh(maxtext_utils.get_mesh_from_config(_config())): + baseline, _ = self._run(micro_batches=3, steps=2) + + want_tree = nnx.to_pure_dict(nnx.state(baseline.model, nnx.Param)) + got_tree = nnx.to_pure_dict(nnx.state(deferred.model, nnx.Param)) + for (path, want), got in zip(jax.tree.flatten_with_path(want_tree)[0], jax.tree.leaves(got_tree)): + np.testing.assert_allclose( + np.asarray(got), np.asarray(want), rtol=1e-6, atol=1e-6, err_msg=f"parameter {jax.tree_util.keystr(path)}" + ) + + def test_an_unreduced_accumulator_survives_a_checkpoint_round_trip(self): + """Saving mid-step must write the gradient total, not one replica's partial. + + Orbax cannot serialize an unreduced array at all -- `device_indices_map` is undefined for + one -- so `save_checkpoint` reduces first. This pins that the value written is the same + total `update()` would have applied, by finishing the step from the restored state and + checking the weights match an uninterrupted run. + """ + output_dir = self.create_tempdir().full_path + cfg = _config( + enable_checkpointing=True, + base_output_directory=output_dir, + async_checkpointing=False, + checkpoint_period=1, + ) + mesh = maxtext_utils.get_mesh_from_config(cfg) + with jax.set_mesh(mesh): + engine = maxtext_engine.MaxTextTrainingEngine(cfg, mesh=mesh) + engine.compile(_batch(cfg, 0)) + engine.fwd_bwd(_batch(cfg, 0)) + engine.fwd_bwd(_batch(cfg, 1)) + # Mid-step: two micro-batches in, no update() yet, so the accumulator is unreduced. + self.assertTrue( + jax.tree.leaves(engine._accumulated_grads)[0].sharding.spec.unreduced, # pylint: disable=protected-access + "the accumulator is not unreduced, so this test is not exercising the reduce-on-save", + ) + engine.save_checkpoint(metadata={"marker": 1}, force=True) + engine._checkpoint_manager.wait_until_finished() # pylint: disable=protected-access + + engine.restore_checkpoint() + engine.update() + resumed = jax.tree.leaves(nnx.to_pure_dict(nnx.state(engine.model, nnx.Param))) + + uninterrupted_cfg = _config() + uninterrupted_mesh = maxtext_utils.get_mesh_from_config(uninterrupted_cfg) + with jax.set_mesh(uninterrupted_mesh): + straight = maxtext_engine.MaxTextTrainingEngine(uninterrupted_cfg, mesh=uninterrupted_mesh) + straight.compile(_batch(uninterrupted_cfg, 0)) + straight.fwd_bwd(_batch(uninterrupted_cfg, 0)) + straight.fwd_bwd(_batch(uninterrupted_cfg, 1)) + straight.update() + expected = jax.tree.leaves(nnx.to_pure_dict(nnx.state(straight.model, nnx.Param))) + + for i, (got, want) in enumerate(zip(resumed, expected)): + np.testing.assert_allclose(np.asarray(got), np.asarray(want), rtol=1e-6, atol=1e-6, err_msg=f"leaf {i}") + + +if __name__ == "__main__": + absltest.main() From 6667f27ffacfe092dea5b19c3b3c906199598344 Mon Sep 17 00:00:00 2001 From: NuojCheng Date: Wed, 2 Sep 2026 16:26:51 +0000 Subject: [PATCH 2/4] Decline the deferred all-reduce on any non-data mesh axis The gate only refused meshes where a second axis shared the *batch* dimension, which caught `fsdp` and missed `tensor`. Tensor parallelism reaches the same contradiction through the feature dimension instead: qwen3-0.6b at dp2 x tp2 on 4x v6e dies on the first micro-batch with ShardingTypeError: out_sharding's unreduced axes should be equal to the contracting specs. Got unreduced axes=frozenset({'data'}) and contracting spec=('data', None, 'tensor') and would have kept dying for every other axis that shards something contracted -- `expert`, `context`, `tensor_sequence`. Enumerating them is how `tensor` was missed in the first place, so the rule is now the blunt one: "data" alone above size 1, or no deferral. Nothing is lost that worked before, since none of those meshes ran. --- src/maxtext/training_engine/maxtext_engine.py | 21 +++++--- ...maxtext_engine_deferred_all_reduce_test.py | 49 +++++++++++++++++++ 2 files changed, 63 insertions(+), 7 deletions(-) diff --git a/src/maxtext/training_engine/maxtext_engine.py b/src/maxtext/training_engine/maxtext_engine.py index 815be224f4..68de3734cd 100644 --- a/src/maxtext/training_engine/maxtext_engine.py +++ b/src/maxtext/training_engine/maxtext_engine.py @@ -188,14 +188,19 @@ def _deferred_all_reduce_shardings(config: Any, mesh: Any, params_shardings: Any - not explicit sharding, where reduced/unreduced specs do not exist; - a mesh with any non-Explicit axis, which those specs are also rejected on. A caller can hand the engine an all-Auto mesh regardless of `config.shard_mode`; - - "data" is not the only mesh axis of size > 1 that the activation batch dimension is - sharded over. A gradient contracts over the batch, and JAX requires the unreduced set - to be exactly the contracted axes -- with `data` and `fsdp` both on the batch it - rejects the backward pass outright ("unreduced axes should be equal to the contracting + - any mesh axis other than "data" has size > 1. JAX requires the unreduced set to be + exactly the axes the gradient contracts over, and "data" is the only one the tag ever + names, so a second axis over any contracted dimension makes the backward pass illegal. + `fsdp` gets there through the batch ("unreduced axes should be equal to the contracting specs. Got unreduced axes=frozenset({'data'}) and contracting spec=(('data', 'fsdp'), - None)"). Widening the tag to both is not the fix: a parameter sharded over `fsdp` - cannot also be unreduced over it. Read the batch axes off the resolved mesh rather - than `config.ici_data_parallelism`, which may still be -1 (auto-fill). + None)") and `tensor` through the feature dimension ("... and contracting spec=('data', + None, 'tensor')"). Widening the tag is not the fix in either case: a parameter sharded + over `fsdp` or `tensor` cannot also be unreduced over it. So the rule is the blunt one + -- pure data parallelism or no deferral -- rather than a list of the axes known to + break, which is how `tensor` was missed. Read the mesh that resolved rather than + `config.ici_*_parallelism`, which may still be -1 (auto-fill). + - the batch dimension is not sharded over "data" after all, leaving no cross-replica + partial to defer and nothing for the tag to describe. """ if getattr(config, "shard_mode", None) != common_types.ShardMode.EXPLICIT: return None, None @@ -203,6 +208,8 @@ def _deferred_all_reduce_shardings(config: Any, mesh: Any, params_shardings: Any return None, None if any(axis_type != jax.sharding.AxisType.Explicit for axis_type in mesh.axis_types): return None, None + if any(size > 1 for axis, size in mesh.shape.items() if axis != _DATA_AXIS): + return None, None try: batch_axes = sharding.batch_mesh_axes(mesh, rules=config.logical_axis_rules) except (KeyError, ValueError, IndexError): diff --git a/tests/post_training/unit/maxtext_engine_deferred_all_reduce_test.py b/tests/post_training/unit/maxtext_engine_deferred_all_reduce_test.py index 5c66e28b82..14f1f8c1a5 100644 --- a/tests/post_training/unit/maxtext_engine_deferred_all_reduce_test.py +++ b/tests/post_training/unit/maxtext_engine_deferred_all_reduce_test.py @@ -246,6 +246,24 @@ def test_declines_when_fsdp_also_shards_the_batch(self): maxtext_engine._deferred_all_reduce_shardings(cfg, mesh, self._params_shardings(mesh)), # pylint: disable=protected-access ) + def test_declines_when_tensor_parallelism_shards_the_features(self): + """The batch is on `data` alone here, and the backward pass is still rejected. + + `fsdp` breaks the tag through the batch dimension; `tensor` breaks it through the + feature dimension of the same activation, which a batch-axis check cannot see: + "unreduced axes should be equal to the contracting specs. Got unreduced + axes=frozenset({'data'}) and contracting spec=('data', None, 'tensor')". Measured on + 4x v6e with qwen3-0.6b at dp2 x tp2, where it crashed `fwd_bwd` on the first + micro-batch. + """ + cfg = _config(ici_data_parallelism=2, ici_tensor_parallelism=2) + mesh = maxtext_utils.get_mesh_from_config(cfg) + + self.assertEqual( + (None, None), + maxtext_engine._deferred_all_reduce_shardings(cfg, mesh, self._params_shardings(mesh)), # pylint: disable=protected-access + ) + def test_a_parameter_already_sharded_over_data_is_left_alone(self): """`reduced` and a shard over the same axis are contradictory, and JAX says so.""" cfg = _config() @@ -385,6 +403,37 @@ def test_an_unreduced_accumulator_survives_a_checkpoint_round_trip(self): for i, (got, want) in enumerate(zip(resumed, expected)): np.testing.assert_allclose(np.asarray(got), np.asarray(want), rtol=1e-6, atol=1e-6, err_msg=f"leaf {i}") + def test_a_tensor_parallel_mesh_still_trains(self): + """The gate declining has to leave a working engine behind, not a broken one. + + This does not reproduce the crash. The toy model here shards plenty over `tensor` and + traces clean on CPU anyway; what raised `ShardingTypeError` was qwen3-0.6b at dp2 x tp2 + on 4x v6e, whose attention kernels contract over a tensor-sharded dimension the + `dot_product` path does not. So the assertion that pins the fix is the gate one above, + and this pins the other half: with the tag off, a tensor-parallel mesh completes a step + and moves the weights. + """ + cfg = _config(ici_data_parallelism=2, ici_tensor_parallelism=2) + mesh = maxtext_utils.get_mesh_from_config(cfg) + with jax.set_mesh(mesh): + engine = maxtext_engine.MaxTextTrainingEngine(cfg, mesh=mesh) + engine.compile(_batch(cfg, 0)) + before = jax.tree.leaves(nnx.to_pure_dict(nnx.state(engine.model, nnx.Param))) + before = [np.asarray(leaf) for leaf in before] + engine.fwd_bwd(_batch(cfg, 0)) + engine.fwd_bwd(_batch(cfg, 1)) + engine.update() + after = jax.tree.leaves(nnx.to_pure_dict(nnx.state(engine.model, nnx.Param))) + + self.assertIsNone( + engine._plain_grad_shardings, # pylint: disable=protected-access + "the deferral engaged on a tensor-parallel mesh, which JAX rejects", + ) + self.assertTrue( + any(not np.array_equal(np.asarray(got), want) for got, want in zip(after, before)), + "the step ran but no parameter moved, so nothing was trained", + ) + if __name__ == "__main__": absltest.main() From e2ea34c1034459192c1790ec455c1d0336c875c9 Mon Sep 17 00:00:00 2001 From: NuojCheng Date: Wed, 2 Sep 2026 15:11:36 +0000 Subject: [PATCH 3/4] Make shard_optimizer_over_data (Zero-1) work in MaxTextTrainingEngine 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). --- src/maxtext/training_engine/maxtext_engine.py | 179 +++++++- src/maxtext/utils/sharding.py | 7 +- .../unit/maxtext_engine_zero1_test.py | 434 ++++++++++++++++++ 3 files changed, 612 insertions(+), 8 deletions(-) create mode 100644 tests/post_training/unit/maxtext_engine_zero1_test.py diff --git a/src/maxtext/training_engine/maxtext_engine.py b/src/maxtext/training_engine/maxtext_engine.py index 68de3734cd..534a808ab6 100644 --- a/src/maxtext/training_engine/maxtext_engine.py +++ b/src/maxtext/training_engine/maxtext_engine.py @@ -59,6 +59,9 @@ # by `_check_pure_state_reusable`. _MODEL_STATE_KEY = "model" +# Where the same split puts the `nnx.Optimizer`, including the optax state Zero-1 shards. +_OPTIMIZER_STATE_KEY = "optimizer" + _PURE_STATE_FALLBACK_WARNING = ( "Cannot keep the train state as a pure pytree across steps (%s), so every fwd_bwd and " "update will re-walk the NNX module graph. That is correct but slow -- the two " @@ -223,6 +226,58 @@ def _deferred_all_reduce_shardings(config: Any, mesh: Any, params_shardings: Any ) +_ZERO1_DECLINED_WARNING = ( + "`shard_optimizer_over_data` (Zero-1) is set, but this engine cannot honour it (%s), so the " + "optimizer state stays replicated over the data axis. Logged once per engine instance." +) + + +def _zero1_active(config: Any, mesh: Any) -> str | None: + """Returns why Zero-1 cannot run here, or None when it can. + + Zero-1 shards the optimizer's parameter-shaped state over the data axis, so each replica + keeps and updates 1/N of the moments. The engine implements it by resharding the + gradients and the parameters onto that same layout inside `_update_kernel` and gathering + the new parameters back on the way out, which needs the reshards to be real ops on a + mesh whose axes are `Explicit` -- under `auto` the layout is GSPMD's to choose and these + would be hints it may ignore, giving a silently replicated optimizer again. + + The flag being off is a reason like any other, so a single call answers "should this run" + rather than leaving the caller to test the flag as well. + """ + if not getattr(config, "shard_optimizer_over_data", False): + return "it is not enabled" + if getattr(config, "shard_mode", None) != common_types.ShardMode.EXPLICIT: + return "it needs shard_mode=explicit" + if mesh is None: + return "the engine has no mesh" + if mesh.shape.get(_DATA_AXIS, 1) <= 1: + return f"the mesh has no {_DATA_AXIS!r} axis to shard the optimizer over" + if any(axis_type != jax.sharding.AxisType.Explicit for axis_type in mesh.axis_types): + return "the mesh has non-Explicit axes" + return None + + +def _zero1_sharding(mesh: Any, aval: Any, base: jax.sharding.NamedSharding | None) -> jax.sharding.NamedSharding | None: + """Returns `base` with the data axis added, or None to leave the value where it is. + + Thin wrapper over the pre-train path's `add_data_to_sharding` so the parameters, the + gradients and the optimizer moments are placed by one function of `(shape, base + sharding)`. That is what makes them 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 value with no dimension the data axis divides -- a scalar `count`, an odd-sized bias -- + comes back unchanged and stays replicated, on all three trees alike. + """ + if base is None or not hasattr(aval, "shape"): + return None + try: + target = sharding.add_data_to_sharding(mesh, (), aval, base) + except AssertionError: + # add_data_to_sharding rejects a shape it cannot shard; leave the value replicated. + return None + return None if target == base else target + + def _conform_accumulator(value: Any, target: jax.sharding.NamedSharding) -> Any: """Moves one accumulated-gradient leaf onto `target`, preserving the value it represents. @@ -470,6 +525,12 @@ def __init__( 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 + # `_update_kernel` shards them to meet the optimizer state, and as it hands them back. + # `None` keeps the whole update on the replicated layout. See `_zero1_active`. + self._zero1_params_shardings: Any = None + self._gathered_params_shardings: Any = None + self._zero1_warned = False self._micro_step_count = 0 # Set when this run resumed from an intra-step checkpoint, cleared once the step it # resumed into completes and its finished state has been checkpointed. @@ -736,6 +797,82 @@ def _publish_state(self, new_state_pure: Any) -> None: 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: + """Says once that Zero-1 was asked for and not done, which used to happen in silence. + + Nothing here is wrong when Zero-1 declines -- the run is correct, just without the + saving -- so this is a warning rather than an error, and once rather than per compile. + """ + if getattr(self._config, "shard_optimizer_over_data", False) and not self._zero1_warned: + self._zero1_warned = True + logging.warning(_ZERO1_DECLINED_WARNING, reason) + + def _zero1_shardings_for(self, params_pure: Any, params_shardings: Any) -> Any: + """Returns the Zero-1 sharding tree for the parameters, or None to keep them replicated.""" + declined = _zero1_active(self._config, self._mesh) + if declined is not None: + self._note_zero1_declined(declined) + return None + + def target(leaf, base): + sharded = _zero1_sharding(self._mesh, leaf, base) + return base if sharded is None else sharded + + return jax.tree.map(target, params_pure, params_shardings) + + def _shard_optimizer_state_over_data(self) -> None: + """Moves the optimizer's parameter-shaped state onto the Zero-1 layout, in place. + + `nnx.Optimizer` allocates the moments eagerly, as `zeros_like` of each parameter, so + they arrive replicated over the data axis however `shard_optimizer_over_data` is set -- + which is why the flag has so far been a silent no-op here. Resharding them once, before + `_compile_for_batch` reads the state's layout back off the arrays, is all it takes for + the rest of the engine to follow: the update kernel's in/out shardings are derived from + exactly these arrays. + + Every leaf is placed by `_zero1_sharding`, moments and bookkeeping alike, rather than + by walking for the ones named `mu`/`nu`. A scalar `count` has no dimension to shard and + comes back untouched, and a partitioned optimizer (Muon's `muon`/`adam` branches) needs + no special case. Idempotent: a leaf already carrying the data axis is left alone, so a + recompile or a restored checkpoint re-runs this for free. + """ + if _zero1_active(self._config, self._mesh) is not None: + return + state_pure = self._read_state_pure() + if _OPTIMIZER_STATE_KEY not in state_pure: + return + + moved = False + + def place(leaf): + nonlocal moved + target = _zero1_sharding(self._mesh, leaf, self._mesh_sharding(leaf)) + if target is None: + return leaf + moved = True + return jax.device_put(leaf, target) + + optimizer_pure = jax.tree.map(place, state_pure[_OPTIMIZER_STATE_KEY]) + if not moved: + return + with self._sharding_ctx(): + nnx.update(self._state, nnx.State({_OPTIMIZER_STATE_KEY: optimizer_pure.raw_mapping})) + self._invalidate_pure_state() + self._refresh_pure_state() + + def _reshard_model_params(self, state_pure: Any, params_shardings: Any) -> Any: + """Returns `state_pure` with its `nnx.Param` leaves moved onto `params_shardings`. + + Used twice inside `_update_kernel`, in opposite directions: down to the Zero-1 layout + the optimizer state lives on, then back up to the replicated one the forward pass and + the kernel's `out_shardings` expect. Only parameters move -- the optimizer state is + already where it belongs, and the rngs and batch statistics alongside it have no + Zero-1 layout to speak of. + """ + params_pure, rest_pure = nnx.split_state(state_pure[_MODEL_STATE_KEY], nnx.Param, ...) + params_pure = jax.tree.map(jax.sharding.reshard, params_pure, params_shardings) + return self._with_model_state(state_pure, nnx.merge_state(params_pure, rest_pure)) + def _fwd_bwd_kernel(self, params, rest, batch, acc_grads=None, acc_denom=None): """Executes a single forward and backward pass and folds the result into the accumulator. @@ -853,13 +990,24 @@ def _update_kernel(self, state_pure, accumulated_grads, accumulated_denominator, grad_norm = None is_skipped_val = None if state_pure is not None: - if self._plain_grad_shardings is not None: - # The gradients arrive `unreduced`: a per-replica partial sum over this step's - # micro-batches. Resharding them back is what emits the single cross-replica - # all-reduce that replaces the one every micro-batch used to pay. First, so that - # everything below -- the division, the norm, clipping, the optimizer -- sees - # ordinary gradients and needs no tag handling of its own. - accumulated_grads = jax.tree.map(jax.sharding.reshard, accumulated_grads, self._plain_grad_shardings) + # Where the gradients have to land before the optimizer can use them. Under Zero-1 + # that is the sharded layout the moments live on; otherwise the plain parameter one. + grad_target = self._zero1_params_shardings + if grad_target is None: + grad_target = self._plain_grad_shardings + if grad_target is not None: + # Resharding away from `unreduced` -- a per-replica partial sum over this step's + # micro-batches -- is what emits the single cross-replica reduction that replaces + # the one every micro-batch used to pay. First, so that everything below (the + # division, the norm, clipping, the optimizer) sees ordinary gradients and needs + # no tag handling of its own. When the deferral is off but Zero-1 is on, the same + # line is just the local slice onto the optimizer's layout. + accumulated_grads = jax.tree.map(jax.sharding.reshard, accumulated_grads, grad_target) + if self._zero1_params_shardings is not None: + # Meet the gradients and the moments on the sharded layout. Free -- slicing a + # replicated array is local -- and it is what makes the optimizer's arithmetic, + # and the memory traffic under it, run on 1/N of every parameter. + state_pure = self._reshard_model_params(state_pure, self._zero1_params_shardings) # This one division is the whole normalization. A zero total means every micro-batch # was empty; yield zeros rather than a NaN, as `gradient_accumulation.py` does. has_weights = accumulated_denominator > 0 @@ -887,6 +1035,11 @@ def _update_kernel(self, state_pure, accumulated_grads, accumulated_denominator, else: local_state.apply_gradients(grads) _, new_state_pure = nnx.split(local_state) + if self._zero1_params_shardings is not None: + # The one all-gather Zero-1 costs: each replica updated its own slice of every + # parameter, and the forward pass needs all of them. The moments stay behind, + # sharded, which is the whole point. + new_state_pure = self._reshard_model_params(new_state_pure, self._gathered_params_shardings) return new_state_pure, grad_norm, is_skipped_val return state_pure, grad_norm, is_skipped_val @@ -1006,6 +1159,9 @@ def _compile_for_batch(self, dynamic_batch: Any, static_batch: dict[str, Any]) - # The only place the graphs are walked: a recompile is when they may legitimately have # changed shape, and everything after is maintained as plain pytrees. self._refresh_pure_state() + # Before the shardings below are read off the state: this is what puts the optimizer + # moments on the Zero-1 layout, and `state_mesh_shardings` has to see them there. + self._shard_optimizer_state_over_data() state_pure = self._read_state_pure() params_pure, rest_pure = self._read_model_pure(getattr(self._state, _MODEL_STATE_KEY, self._model)) @@ -1031,6 +1187,11 @@ def accum_kernel(params, rest, dynamic, acc_grads, acc_denom): self._reduced_params_shardings, self._unreduced_grad_shardings = _deferred_all_reduce_shardings( self._config, self._mesh, params_shardings ) + # Zero-1 lives entirely inside `_update_kernel`, so it changes no kernel signature: + # the parameters cross every jit boundary replicated exactly as before, and only the + # optimizer state -- already moved above -- is stored sharded. + self._zero1_params_shardings = self._zero1_shardings_for(params_pure, params_shardings) + self._gathered_params_shardings = params_shardings if self._zero1_params_shardings is not None else None grad_shardings = self._unreduced_grad_shardings if grad_shardings is None: grad_shardings = params_shardings @@ -1058,6 +1219,10 @@ def accum_kernel(params, rest, dynamic, acc_grads, acc_denom): self._reduced_params_shardings = None self._unreduced_grad_shardings = None self._plain_grad_shardings = None + # `_zero1_shardings_for` is not reached on this branch, so the request is declined here. + self._note_zero1_declined(_zero1_active(self._config, self._mesh)) + self._zero1_params_shardings = None + self._gathered_params_shardings = None # 1. JIT Compile Micro FWD/BWD Pass. # diff --git a/src/maxtext/utils/sharding.py b/src/maxtext/utils/sharding.py index 9574a5a4b2..73eeab7d28 100644 --- a/src/maxtext/utils/sharding.py +++ b/src/maxtext/utils/sharding.py @@ -684,7 +684,12 @@ def add_data_to_sharding(mesh, path, aval, sharding): raise AssertionError(f"Could not shard {jax.tree_util.keystr(path)} of shape={aval.shape} with {sharding=}") from e pspec = sharding.spec - if "data" in jax.tree.leaves(pspec): + # `tuple(pspec)`, not `pspec`: a PartitionSpec is a pytree *leaf*, so flattening one gives + # back the spec itself and this guard never fired. Its entries are what have to be walked, + # and they nest -- a dimension sharded over two axes is a tuple. Without this, a leaf + # already sharded over "data" gets a second one and `NamedSharding` rejects the result + # outright (`DuplicateSpecError: P(('data', 'data'), None)`). + if "data" in jax.tree.leaves(tuple(pspec)): return sharding for idx, (size, partition) in enumerate(zip(sharded_shape, pspec)): diff --git a/tests/post_training/unit/maxtext_engine_zero1_test.py b/tests/post_training/unit/maxtext_engine_zero1_test.py new file mode 100644 index 0000000000..9bb609d3e8 --- /dev/null +++ b/tests/post_training/unit/maxtext_engine_zero1_test.py @@ -0,0 +1,434 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""`shard_optimizer_over_data` (Zero-1) in `MaxTextTrainingEngine`. + +The flag used to be 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. It now +shards the parameter-shaped optimizer state over the data axis and does the update on those +slices, gathering the new parameters back at the end. + +Like the deferral it pairs with, this is a change nothing functional depends on -- get it +wrong and the model still trains, just without the saving. So the tests assert on where the +arrays actually are (`sharding.spec`, and the shard each device holds) and on the compiled +HLO, and each such assertion is mirrored by the same probe run with the flag off. +""" + +import os + +# Must precede the first JAX import: a data-parallel mesh needs more than one device, and +# the CPU backend reads this only at initialization. +os.environ.setdefault("XLA_FLAGS", "--xla_force_host_platform_device_count=4") + +import re # pylint: disable=wrong-import-position +import unittest # pylint: disable=wrong-import-position + +from absl.testing import absltest # pylint: disable=wrong-import-position +from flax import nnx # pylint: disable=wrong-import-position +import jax # pylint: disable=wrong-import-position +from maxtext.training_engine import maxtext_engine # pylint: disable=wrong-import-position +from maxtext.utils import maxtext_utils # pylint: disable=wrong-import-position +import numpy as np # pylint: disable=wrong-import-position +import pytest # pylint: disable=wrong-import-position + +# The tiny-real-decoder rig this shares with the deferral it composes with: same model, same +# mesh, same batch. Reusing it is the point -- the two features have to hold on one config. +from tests.post_training.unit.maxtext_engine_deferred_all_reduce_test import ( # pylint: disable=wrong-import-position + _REQUIRED_DEVICES, + _KernelHlo, + _array_all_reduces, + _batch, + _config, + _no_deferral, +) + +# training_engine imports tunix, so these tests need the post-training dependency bundle. +pytestmark = [pytest.mark.post_training] + +_DATA = maxtext_engine._DATA_AXIS # pylint: disable=protected-access + +# The result shape of an all-gather in optimized HLO, in both the fused and the async form. +_ALL_GATHER = re.compile(r"=\s*(.+?)\s+all-gather(?:-start|-done)?\(") +# The dimensions inside one `f32[128,64]{1,0}`; a tupled result yields one match per element. +_SHAPE_DIMS = re.compile(r"\[([\d,]*)\]") + + +def _gathered_elements(hlo: str) -> int: + """Total result size of every all-gather in `hlo`, as a stand-in for gathered volume. + + Counting instructions would be brittle -- XLA fuses and splits them freely -- and the + absolute number here means little, since an async gather's start and done both count. Only + the difference against the same kernel compiled without Zero-1 is ever asserted on, and + both sides of that are counted the same way. + """ + total = 0 + for line in hlo.splitlines(): + if match := _ALL_GATHER.search(line): + for dims in _SHAPE_DIMS.findall(match.group(1)): + if dims: + total += int(np.prod([int(dim) for dim in dims.split(",")])) + return total + + +def _axes(spec) -> list[str]: + """The mesh axes a `PartitionSpec` names. + + A `PartitionSpec` is a pytree *leaf*, so flattening one gives back the spec itself; its + entries have to be opened first, and they nest -- a dimension sharded over two axes is a + tuple. The same trap the guard in `add_data_to_sharding` was written with. + """ + return jax.tree.leaves(tuple(spec)) + + +def _zero1_config(**overrides): + """The shared config with Zero-1 on and a stateful optimizer to shard.""" + # SGD, the shared default, carries no parameter-shaped state at all, so Zero-1 would have + # nothing to move and every assertion below would pass vacuously. + overrides.setdefault("opt_type", "adamw") + return _config(shard_optimizer_over_data=True, **overrides) + + +def _moments(engine): + """`{path: array}` for every parameter-shaped optimizer moment in the engine's state.""" + _, state_pure = nnx.split(engine.state) + return { + jax.tree_util.keystr(path): leaf + for path, leaf in jax.tree_util.tree_leaves_with_path(state_pure) + if "['mu']" in jax.tree_util.keystr(path) or "['nu']" in jax.tree_util.keystr(path) + } + + +def _params(engine): + """`{path: array}` for the model's parameters.""" + return { + jax.tree_util.keystr(path): leaf + for path, leaf in jax.tree.flatten_with_path(nnx.to_pure_dict(nnx.state(engine.model, nnx.Param)))[0] + } + + +@unittest.skipIf( + jax.device_count() < _REQUIRED_DEVICES, + f"needs {_REQUIRED_DEVICES} devices; set XLA_FLAGS=--xla_force_host_platform_device_count={_REQUIRED_DEVICES}", +) +class Zero1GateTest(absltest.TestCase): + """`_zero1_active` decides whether the engine can honour the flag. It must decline widely.""" + + def test_declines_when_the_flag_is_off(self): + cfg = _config() + mesh = maxtext_utils.get_mesh_from_config(cfg) + + self.assertIsNotNone(maxtext_engine._zero1_active(cfg, mesh)) # pylint: disable=protected-access + + def test_opens_on_an_explicit_data_parallel_mesh(self): + cfg = _zero1_config() + mesh = maxtext_utils.get_mesh_from_config(cfg) + + self.assertIsNone(maxtext_engine._zero1_active(cfg, mesh)) # pylint: disable=protected-access + + def test_declines_under_auto_shard_mode(self): + """Under `auto` the reshards are hints GSPMD may ignore, which would replicate silently.""" + cfg = _zero1_config(shard_mode="auto") + mesh = maxtext_utils.get_mesh_from_config(cfg) + + self.assertIn("explicit", maxtext_engine._zero1_active(cfg, mesh)) # pylint: disable=protected-access + + def test_declines_on_an_auto_axis_mesh_even_in_explicit_mode(self): + """A caller can hand the engine a bare `jax.sharding.Mesh` whatever `shard_mode` says.""" + cfg = _zero1_config() + explicit_mesh = maxtext_utils.get_mesh_from_config(cfg) + auto_mesh = jax.sharding.Mesh(explicit_mesh.devices, explicit_mesh.axis_names) + + self.assertIn("Explicit", maxtext_engine._zero1_active(cfg, auto_mesh)) # pylint: disable=protected-access + + def test_declines_when_there_are_no_data_replicas(self): + cfg = _zero1_config(ici_data_parallelism=1, ici_tensor_parallelism=_REQUIRED_DEVICES) + mesh = maxtext_utils.get_mesh_from_config(cfg) + + self.assertIn(_DATA, maxtext_engine._zero1_active(cfg, mesh)) # pylint: disable=protected-access + + def test_declines_without_a_mesh(self): + self.assertIn("mesh", maxtext_engine._zero1_active(_zero1_config(), None)) # pylint: disable=protected-access + + +@unittest.skipIf( + jax.device_count() < _REQUIRED_DEVICES, + f"needs {_REQUIRED_DEVICES} devices; set XLA_FLAGS=--xla_force_host_platform_device_count={_REQUIRED_DEVICES}", +) +class Zero1ShardingTest(absltest.TestCase): + """`_zero1_sharding` places one leaf. Everything Zero-1 moves goes through it.""" + + def setUp(self): + super().setUp() + self.mesh = maxtext_utils.get_mesh_from_config(_zero1_config()) + + def _replicated(self, rank): + return jax.sharding.NamedSharding(self.mesh, jax.sharding.PartitionSpec(*(None,) * rank)) + + def _place(self, shape): + return maxtext_engine._zero1_sharding( # pylint: disable=protected-access + self.mesh, jax.ShapeDtypeStruct(shape, jax.numpy.float32), self._replicated(len(shape)) + ) + + def test_adds_the_data_axis_to_the_first_dimension_that_divides(self): + self.assertEqual(self._place((128, 64)).spec, jax.sharding.PartitionSpec(_DATA, None)) + + def test_skips_a_dimension_the_data_axis_does_not_divide(self): + self.assertEqual(self._place((3, 64)).spec, jax.sharding.PartitionSpec(None, _DATA)) + + def test_leaves_a_scalar_alone(self): + """`adamw`'s step `count`, and every rng counter beside it. Nothing to slice.""" + self.assertIsNone(self._place(())) + + def test_leaves_a_shape_no_dimension_of_which_divides_alone(self): + self.assertIsNone(self._place((3, 5))) + + def test_leaves_a_leaf_already_sharded_over_data_alone(self): + already = jax.sharding.NamedSharding(self.mesh, jax.sharding.PartitionSpec(_DATA, None)) + + self.assertIsNone( + maxtext_engine._zero1_sharding( # pylint: disable=protected-access + self.mesh, jax.ShapeDtypeStruct((128, 64), jax.numpy.float32), already + ) + ) + + +@pytest.mark.integration_test +@unittest.skipIf( + jax.device_count() < _REQUIRED_DEVICES, + f"needs {_REQUIRED_DEVICES} devices; set XLA_FLAGS=--xla_force_host_platform_device_count={_REQUIRED_DEVICES}", +) +class Zero1Test(absltest.TestCase): + """End to end on a real decoder: where the optimizer state sits, and what the weights do.""" + + def _run(self, micro_batches: int = 2, steps: int = 2, cfg=None, probe: bool = False): + """Runs `steps` optimizer steps of `micro_batches` each. + + Returns `(engine, {kernel: optimized hlo})`, the HLO empty unless `probe`. Reading it + means lowering the kernel again, and the `reduced` tag the fwd/bwd kernels apply needs + the mesh to be set for that -- so it is read here, inside the context, not by the caller. + """ + cfg = cfg if cfg is not None else _zero1_config() + mesh = maxtext_utils.get_mesh_from_config(cfg) + kernels = {"first": "_compiled_fwd_bwd", "accum": "_compiled_fwd_bwd_accum", "update": "_compiled_update"} + with jax.set_mesh(mesh): + engine = maxtext_engine.MaxTextTrainingEngine(cfg, mesh=mesh) + engine.compile(_batch(cfg, 0)) + probes = {name: _KernelHlo(engine, attr) for name, attr in kernels.items()} if probe else {} + for step in range(steps): + for micro in range(micro_batches): + engine.fwd_bwd(_batch(cfg, step * micro_batches + micro)) + engine.update() + return engine, {name: probe_for.text() for name, probe_for in probes.items()} + + def test_the_gate_opens_on_this_configuration(self): + """Guards every other test in this class: without this they would all pass vacuously.""" + engine, _ = self._run(micro_batches=1, steps=1) + + self.assertIsNotNone( + engine._zero1_params_shardings, # pylint: disable=protected-access + "Zero-1 never engaged, so nothing below is testing it", + ) + + def test_the_optimizer_moments_are_sharded_over_the_data_axis(self): + """The saving itself: each replica holds and updates 1/N of every moment.""" + engine, _ = self._run() + + moments = _moments(engine) + self.assertNotEmpty(moments, "adamw kept no parameter-shaped state, so there is nothing to shard") + for path, leaf in moments.items(): + self.assertIn(_DATA, _axes(leaf.sharding.spec), f"{path} is not sharded over {_DATA!r}") + shard = leaf.addressable_shards[0].data.shape + self.assertEqual( + np.prod(shard) * _REQUIRED_DEVICES, + np.prod(leaf.shape), + f"{path} claims to be sharded but each device still holds {shard} of {leaf.shape}", + ) + + def test_the_parameters_themselves_stay_replicated(self): + """Zero-1, not Zero-2/3: only the optimizer state is stored sharded. + + The forward pass wants whole parameters and every kernel signature is unchanged, so the + slicing lives entirely inside `update()`. + """ + engine, _ = self._run() + + for path, leaf in _params(engine).items(): + self.assertNotIn(_DATA, _axes(leaf.sharding.spec), f"parameter {path} came back sharded") + + def test_without_the_flag_the_moments_stay_replicated(self): + """Proves the probes above can fail. Same model, same optimizer, flag off.""" + engine, _ = self._run(cfg=_config(opt_type="adamw")) + + self.assertIsNone(engine._zero1_params_shardings) # pylint: disable=protected-access + for path, leaf in _moments(engine).items(): + self.assertNotIn(_DATA, _axes(leaf.sharding.spec), f"{path} is sharded with the flag off") + + def test_zero1_costs_one_all_gather_in_update_and_nothing_per_micro_batch(self): + """Where the traffic Zero-1 adds is, and where it must not be. + + Each replica updates its own slice, so the new parameters have to be gathered before the + next forward pass -- once per optimizer step, in `update()`. If that gather ever appears + in a micro-batch kernel instead, Zero-1 has become a per-micro-batch cost. + """ + baseline, baseline_probes = self._run(cfg=_config(opt_type="adamw"), probe=True) + engine, probes = self._run(probe=True) + + added = {k: _gathered_elements(probes[k]) - _gathered_elements(baseline_probes[k]) for k in probes} + self.assertGreater(added["update"], 0, "update() gathers nothing, so the parameters were never sharded") + self.assertEqual(added["first"], 0, "Zero-1 added an all-gather to the first micro-batch") + self.assertEqual(added["accum"], 0, "Zero-1 added an all-gather to the accumulating micro-batches") + self.assertIsNotNone(engine._zero1_params_shardings) # pylint: disable=protected-access + self.assertIsNone(baseline._zero1_params_shardings) # pylint: disable=protected-access + + def test_zero1_composes_with_the_deferred_all_reduce(self): + """The pair is the point: one reduction per step, on 1/N of the optimizer. + + Zero-1 reshards the gradients onto the moments' layout inside `update()`, which is the + same reshard that discharges the deferral's `unreduced` tag. So turning it on must not + put parameter-sized traffic back into the micro-batches. + """ + engine, probes = self._run(probe=True) + + self.assertIsNotNone(engine._plain_grad_shardings) # pylint: disable=protected-access + self.assertIsNotNone(engine._zero1_params_shardings) # pylint: disable=protected-access + self.assertEmpty(_array_all_reduces(probes["first"])) + self.assertEmpty(_array_all_reduces(probes["accum"])) + self.assertNotEmpty(_array_all_reduces(probes["update"]), "the gradients are never reduced at all") + + def test_zero1_does_not_change_the_weights(self): + """Same optimizer arithmetic, run elementwise on disjoint slices instead of on all of it. + + Every operation `adamw` applies is elementwise in the parameter, so splitting the tensor + across replicas changes nothing about the result -- on this CPU mesh, not even the last + bit. The tolerance is for accelerators, where the gather is not exact. + """ + zero1, _ = self._run(micro_batches=3, steps=3) + baseline, _ = self._run(micro_batches=3, steps=3, cfg=_config(opt_type="adamw")) + + want, got = _params(baseline), _params(zero1) + self.assertEqual(sorted(want), sorted(got)) + for path, expected in want.items(): + np.testing.assert_allclose( + np.asarray(got[path]), np.asarray(expected), rtol=1e-6, atol=1e-6, err_msg=f"parameter {path}" + ) + + def test_gradient_clipping_still_sees_the_whole_gradient(self): + """The one part of `update()` that is not elementwise. + + `l2norm_pytree` sums squares over every element, and under Zero-1 those elements are + spread across replicas. If the sum stayed replica-local the norm would come out too + small by a factor of N and clipping would barely bite; the weights would then differ. + """ + clipped = {"gradient_clipping_threshold": 1e-4, "opt_type": "adamw"} + zero1, _ = self._run(micro_batches=2, steps=2, cfg=_zero1_config(**clipped)) + baseline, _ = self._run(micro_batches=2, steps=2, cfg=_config(**clipped)) + + want, got = _params(baseline), _params(zero1) + for path, expected in want.items(): + np.testing.assert_allclose( + np.asarray(got[path]), np.asarray(expected), rtol=1e-6, atol=1e-6, err_msg=f"parameter {path}" + ) + + def test_zero1_alone_is_enough_without_the_deferral(self): + """The two are independent. With the deferral withheld, Zero-1 still shards the moments.""" + with _no_deferral(): + engine, _ = self._run() + + self.assertIsNone(engine._plain_grad_shardings) # pylint: disable=protected-access + self.assertIsNotNone(engine._zero1_params_shardings) # pylint: disable=protected-access + for path, leaf in _moments(engine).items(): + self.assertIn(_DATA, _axes(leaf.sharding.spec), f"{path} is not sharded over {_DATA!r}") + + def test_a_recompile_leaves_the_already_sharded_moments_where_they_are(self): + """A second batch shape re-enters `_compile_for_batch`, which re-places the moments. + + They are already on the Zero-1 layout by then, so the placement has to be a no-op -- + adding the data axis a second time produces `P(('data', 'data'), ...)`, which + `NamedSharding` rejects outright and which would take the whole engine down. + """ + cfg = _zero1_config() + mesh = maxtext_utils.get_mesh_from_config(cfg) + with jax.set_mesh(mesh): + engine = maxtext_engine.MaxTextTrainingEngine(cfg, mesh=mesh) + engine.compile(_batch(cfg, 0)) + engine.fwd_bwd(_batch(cfg, 0)) + engine.update() + # A shorter sequence: a different dynamic batch shape, so `fwd_bwd` recompiles. + short = {name: value[:, : cfg.max_target_length // 2] for name, value in _batch(cfg, 1).items()} + engine.fwd_bwd(short) + engine.update() + + for path, leaf in _moments(engine).items(): + self.assertEqual(_axes(leaf.sharding.spec).count(_DATA), 1, f"{path} was sharded over {_DATA!r} twice") + + def test_the_moments_survive_a_checkpoint_round_trip_still_sharded(self): + """Restoring does not recompile, so what Orbax hands back has to land where it left. + + `_compiled_update` was built against sharded moments. If they came back replicated the + resumed step would die on an `in_shardings` mismatch -- and if it somehow did not, the + optimizer would silently be replicated again for the rest of the run. + """ + output_dir = self.create_tempdir().full_path + cfg = _zero1_config( + enable_checkpointing=True, + base_output_directory=output_dir, + async_checkpointing=False, + checkpoint_period=1, + ) + mesh = maxtext_utils.get_mesh_from_config(cfg) + with jax.set_mesh(mesh): + engine = maxtext_engine.MaxTextTrainingEngine(cfg, mesh=mesh) + engine.compile(_batch(cfg, 0)) + # One full step first, so the moments are non-zero and actually carry information. + engine.fwd_bwd(_batch(cfg, 0)) + engine.update() + # Then mid-step: one micro-batch in, no update() yet. + engine.fwd_bwd(_batch(cfg, 1)) + engine.save_checkpoint(metadata={"marker": 1}, force=True) + engine._checkpoint_manager.wait_until_finished() # pylint: disable=protected-access + + engine.restore_checkpoint() + restored = _moments(engine) + engine.update() + resumed = _params(engine) + + self.assertNotEmpty(restored, "no moments came back, so nothing below is checked") + for path, leaf in restored.items(): + self.assertIn(_DATA, _axes(leaf.sharding.spec), f"{path} came back replicated from the checkpoint") + + uninterrupted, _ = self._run(micro_batches=1, steps=2) + for path, expected in _params(uninterrupted).items(): + np.testing.assert_allclose( + np.asarray(resumed[path]), np.asarray(expected), rtol=1e-6, atol=1e-6, err_msg=f"parameter {path}" + ) + + def test_a_request_the_engine_cannot_honour_is_reported_once(self): + """The failure mode this replaces was silence: the flag set, and nothing done about it.""" + cfg = _zero1_config(shard_mode="auto") + mesh = maxtext_utils.get_mesh_from_config(cfg) + with jax.set_mesh(mesh): + engine = maxtext_engine.MaxTextTrainingEngine(cfg, mesh=mesh) + with self.assertLogs(level="WARNING") as logs: + engine.compile(_batch(cfg, 0)) + engine.fwd_bwd(_batch(cfg, 0)) + engine.update() + + declined = [line for line in logs.output if "Zero-1" in line] + self.assertLen(declined, 1, f"expected exactly one Zero-1 warning, got {declined}") + self.assertIn("shard_mode=explicit", declined[0]) + self.assertIsNone(engine._zero1_params_shardings) # pylint: disable=protected-access + + +if __name__ == "__main__": + absltest.main() From fcec4fca1f078ed02b3467af1515ee2419706a4a Mon Sep 17 00:00:00 2001 From: chengnuojin Date: Wed, 2 Sep 2026 23:04:40 +0000 Subject: [PATCH 4/4] Let TrainerPayload subclasses follow whichever TrainerPayload tunix ships Tunix turned TrainerPayload into a frozen flax.struct.dataclass. Python rejects both directions of frozen/non-frozen dataclass inheritance, so every subclass has to follow the installed tunix rather than pick a side: against tunix head, RouterReplayTrainerPayload and the test suite's DummyPayload both raise TypeError at module load, and hard-coding frozen=True would break the pinned revision instead. abstract_engine.payload_dataclass picks the decorator once, next to the TrainerPayload re-export it has to track. flax.struct.dataclass on the frozen path rather than a bare frozen=True, because only the former registers the subclass as a pytree node -- a plain frozen subclass of a registered parent is absent from the registry, so jax.tree.leaves() silently yields the payload itself as one leaf instead of its fields. --- .../training_engine/abstract_engine.py | 22 +++++++++++++++++++ src/maxtext/training_engine/maxtext_engine.py | 2 +- .../post_training/unit/maxtext_engine_test.py | 9 +++++--- 3 files changed, 29 insertions(+), 4 deletions(-) diff --git a/src/maxtext/training_engine/abstract_engine.py b/src/maxtext/training_engine/abstract_engine.py index b333f17160..0d7468167b 100644 --- a/src/maxtext/training_engine/abstract_engine.py +++ b/src/maxtext/training_engine/abstract_engine.py @@ -71,9 +71,31 @@ "TrainerPayload", "TrainingConfig", "WeightedMetric", + "payload_dataclass", ] +# Decorate every `TrainerPayload` subclass with this rather than with `dataclasses.dataclass`. +# +# Tunix turned `TrainerPayload` into a frozen `flax.struct.dataclass`; at the revision +# `src/dependencies/extra_deps/post_train_github_deps.txt` pins, it is still a plain mutable +# dataclass. Python rejects *both* directions of frozen/non-frozen dataclass inheritance +# ("cannot inherit non-frozen dataclass from a frozen one", and the converse), so a subclass has +# to follow whichever tunix is installed rather than pick a side; hard-coding either one turns a +# dependency bump into a TypeError at module load. +# +# `flax.struct.dataclass` on the frozen path, not a bare `frozen=True`: only the former registers +# the subclass as a pytree node. A plain frozen subclass of a registered parent is absent from the +# registry, so `jax.tree.leaves(payload)` yields the payload itself as one leaf instead of its +# fields -- which type-checks, runs, and is wrong for any caller putting a payload across a +# `jax.jit` or `jax.device_put` boundary. +payload_dataclass = ( + flax.struct.dataclass(frozen=True, kw_only=True) + if TrainerPayload.__dataclass_params__.frozen + else dataclasses.dataclass(kw_only=True) +) + + @flax.struct.dataclass class MetricsBuffer: """A buffer for storing and aggregating unreduced metrics on-device. diff --git a/src/maxtext/training_engine/maxtext_engine.py b/src/maxtext/training_engine/maxtext_engine.py index 534a808ab6..a4ee9c8045 100644 --- a/src/maxtext/training_engine/maxtext_engine.py +++ b/src/maxtext/training_engine/maxtext_engine.py @@ -295,7 +295,7 @@ def _conform_accumulator(value: Any, target: jax.sharding.NamedSharding) -> Any: return jax.device_put(value, target) -@dataclasses.dataclass(kw_only=True) +@abstract_engine.payload_dataclass class RouterReplayTrainerPayload(abstract_engine.TrainerPayload): """A TrainerPayload extension carrying forced router-replay expert decisions. diff --git a/tests/post_training/unit/maxtext_engine_test.py b/tests/post_training/unit/maxtext_engine_test.py index c3821b9b8c..64c16b1a0a 100644 --- a/tests/post_training/unit/maxtext_engine_test.py +++ b/tests/post_training/unit/maxtext_engine_test.py @@ -61,10 +61,13 @@ def __init__(self): self.calls = nnx.BatchStat(jnp.array(0.0)) -@dataclasses.dataclass(kw_only=True) +# `payload_dataclass` resolves to a dataclass decorator at import time, but pylint only +# recognises the literal `dataclasses.dataclass`/`flax.struct.dataclass` forms and so reports the +# `field()` calls below as being outside a dataclass. +@abstract_engine.payload_dataclass class DummyPayload(abstract_engine.TrainerPayload): - token_ids: Any = dataclasses.field(default_factory=lambda: jnp.ones((2, 2))) - token_mask: Any = dataclasses.field(default_factory=lambda: jnp.ones((2, 2))) + token_ids: Any = dataclasses.field(default_factory=lambda: jnp.ones((2, 2))) # pylint: disable=invalid-field-call + token_mask: Any = dataclasses.field(default_factory=lambda: jnp.ones((2, 2))) # pylint: disable=invalid-field-call class MaxTextTrainingEngineTest(absltest.TestCase):