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()