From ebc0191442bc3e92f4e9279e0741dcf63b2dacee Mon Sep 17 00:00:00 2001
From: Param Bole <parambole@google.com>
Date: Sat, 10 Oct 2026 00:43:06 -0700
Subject: [PATCH 1/6] [MaxTextEngine] Prepare a micro-batch's inputs on the
 mesh ahead of its fwd_bwd

prepare_fwd_bwd_input(payload) maps a micro-batch to loss inputs and
device_puts the traced half onto the shardings fwd_bwd compiles them onto
(_batch_data_shardings), leaving the closed-over half untouched;
fwd_bwd(PreparedFwdBwdInput) then dispatches without moving data. Tunix's
TrainerWorker.stage_fwd_bwd calls it when the RL loop sends the next step's
first micro-batch while the device still computes the previous step's last
micro-batch and update. Costs one micro-batch of inputs in device memory until
it is trained (~20 MB per device for a [64 x 65,536] router-replay batch).

Tests: prepared micro-batches train bit-identically (weights and metrics,
compiled and eager, no recompile), also when staged through a real Tunix
TrainerWorker; a one-token change is seen; preparing changes no engine state.
---
 src/maxtext/training_engine/maxtext_engine.py |  41 ++-
 .../maxtext_engine_prepared_input_test.py     | 233 ++++++++++++++++++
 2 files changed, 271 insertions(+), 3 deletions(-)
 create mode 100644 tests/post_training/unit/maxtext_engine_prepared_input_test.py

diff --git a/src/maxtext/training_engine/maxtext_engine.py b/src/maxtext/training_engine/maxtext_engine.py
index 049da90b6b..61ece83dc6 100644
--- a/src/maxtext/training_engine/maxtext_engine.py
+++ b/src/maxtext/training_engine/maxtext_engine.py
@@ -470,6 +470,17 @@ def _to_aval(value: Any) -> Any:
   return jax.ShapeDtypeStruct(value.shape, value.dtype, sharding=getattr(value, "sharding", None))
 
 
+@dataclasses.dataclass(frozen=True)
+class PreparedFwdBwdInput:
+  """A micro-batch already mapped to loss inputs and placed on the mesh.
+
+  Made by `MaxTextTrainingEngine.prepare_fwd_bwd_input`; `fwd_bwd` takes it in place of
+  the payload it came from.
+  """
+
+  batch: Any
+
+
 def _row_and_mask(payload: abstract_engine.RLTrainerPayload) -> tuple[jax.Array, jax.Array]:
   """Returns the `[batch, seq]` model row and its mask, prompt part then completion part.
 
@@ -2404,17 +2415,41 @@ def _start_running_sum(self, grads: Any) -> Any:
     cast = iter(self._cast_uncast_grads([g for g, is_uncast in zip(leaves, uncast) if is_uncast]))
     return treedef.unflatten([next(cast) if is_uncast else g for g, is_uncast in zip(leaves, uncast)])
 
+  def prepare_fwd_bwd_input(self, payload: abstract_engine.TrainerPayload) -> PreparedFwdBwdInput:
+    """Maps a micro-batch to loss inputs and places them on the mesh, ahead of its `fwd_bwd`.
+
+    For a caller that receives the micro-batch early (Tunix's `TrainerWorker.stage_fwd_bwd`,
+    while the device still computes the previous step's last micro-batch and update): the
+    input mapping and the host-to-device copy run here, and `fwd_bwd(prepared)` later
+    dispatches without moving data. The traced inputs go to exactly the shardings `fwd_bwd`
+    compiles them onto (`_batch_data_shardings`), with the closed-over half untouched, so the
+    kernel, its inputs and the result are those of `fwd_bwd(payload)`. Holds one micro-batch
+    of inputs in device memory until that `fwd_bwd`.
+
+    Args:
+      payload: Packed micro-batch training input.
+
+    Returns:
+      The input to pass to `fwd_bwd` instead of `payload`.
+    """
+    batch = self._prepare_batch(payload)
+    if self._mesh is not None and isinstance(batch, dict):
+      dynamic_batch, static_batch = _split_static_and_dynamic(batch)
+      dynamic_batch = jax.device_put(dynamic_batch, self._batch_data_shardings(dynamic_batch))
+      batch = {**dynamic_batch, **static_batch}
+    return PreparedFwdBwdInput(batch=batch)
+
   @_profiled_step("fwd_bwd", control_profile=True)
-  def fwd_bwd(self, payload: abstract_engine.TrainerPayload, **kwargs: Any) -> None:
+  def fwd_bwd(self, payload: abstract_engine.TrainerPayload | PreparedFwdBwdInput, **kwargs: Any) -> None:
     """Executes a micro-batch forward-backward pass and accumulates gradients.
 
     Args:
-      payload: Packed micro-batch training input.
+      payload: Packed micro-batch training input, or what `prepare_fwd_bwd_input` made of it.
       **kwargs: Implementation-specific options. `apply_optimizer` (bool), which Tunix passes as the
         orchestrator's flag for the step's final micro-batch, selects the final micro-batch's dropless program
         under `retry_dropless_final_micro_batch`; everything else is ignored.
     """
-    batch = self._prepare_batch(payload)
+    batch = payload.batch if isinstance(payload, PreparedFwdBwdInput) else self._prepare_batch(payload)
     apply_optimizer = kwargs.get("apply_optimizer")
     final = False
     if self._dropless_replay and self._final_dropless:
diff --git a/tests/post_training/unit/maxtext_engine_prepared_input_test.py b/tests/post_training/unit/maxtext_engine_prepared_input_test.py
new file mode 100644
index 0000000000..c3401a9ad4
--- /dev/null
+++ b/tests/post_training/unit/maxtext_engine_prepared_input_test.py
@@ -0,0 +1,233 @@
+# Copyright 2023–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
+#
+#    https://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.
+
+"""Tests for MaxTextTrainingEngine.prepare_fwd_bwd_input.
+
+A micro-batch prepared ahead (mapped and placed on the mesh while earlier work
+still runs) must train exactly like the payload it came from: same compiled
+kernel, same inputs, bit-identical weights and metrics.
+"""
+
+import dataclasses
+from typing import Any
+from unittest import mock
+
+from absl.testing import absltest
+from flax import nnx
+from flax import struct
+import jax
+import jax.numpy as jnp
+from maxtext.configs import pyconfig
+from maxtext.training_engine import abstract_engine
+from maxtext.training_engine import maxtext_engine
+from maxtext.utils import maxtext_utils
+from tests.utils.test_helpers import get_test_config_path
+import numpy as np
+import optax
+import pytest
+
+# training_engine imports tunix, so these tests need the post-training dependency bundle.
+pytestmark = [pytest.mark.post_training]
+
+
+class _Model(nnx.Module):
+
+  def __init__(self):
+    self.weights = nnx.Param(jnp.array([1.0, 2.0]))
+
+
+@struct.dataclass(frozen=True, kw_only=True)
+class _Payload(abstract_engine.TrainerPayload):
+  token_ids: Any
+  token_mask: Any
+  metadata: dict[str, Any] = struct.field(pytree_node=False, default_factory=dict)
+
+
+def _payload(seed: int) -> _Payload:
+  rng = np.random.default_rng(seed)
+  return _Payload(
+      token_ids=rng.normal(size=(2, 4)).astype(np.float32),
+      token_mask=(rng.random((2, 4)) > 0.25).astype(np.float32),
+      metadata={"seed": seed},
+  )
+
+
+def _loss_fn(model, tokens, mask, scale):
+  # Depends on the data and on both weights, so any change to either shows.
+  pred = model.weights.value[0] * tokens + model.weights.value[1] * tokens**2
+  return (
+      abstract_engine.WeightedMetric(
+          unreduced_sum=jnp.sum(pred * mask) * scale,
+          denominator=jnp.sum(mask),
+      ),
+      {},
+  )
+
+
+def _model_input(payload: _Payload) -> dict[str, Any]:
+  # Tunix's shape: traced arrays plus a closed-over Python constant.
+  return {"tokens": payload.token_ids, "mask": payload.token_mask, "scale": 2.0}
+
+
+class PrepareFwdBwdInputTest(absltest.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    self.config = pyconfig.initialize(
+        [None, get_test_config_path()],
+        model_name="llama3.1-8b",
+        run_name="test_run",
+        base_output_directory=self.create_tempdir().full_path,
+        init_weights_seed=42,
+        micro_batch_size_to_train_on=2,
+        gradient_accumulation_steps=1,
+        enable_dropout=False,
+        record_internal_nn_metrics=False,
+        enable_tensorboard=False,
+        tensorboard_dir=self.create_tempdir().full_path,
+        skip_jax_distributed_system=True,
+        enable_checkpointing=False,
+        profiler_steps=0,
+    )
+    self.mesh = jax.sharding.Mesh(maxtext_utils.create_device_mesh(self.config), self.config.mesh_axes)
+    for patcher in (
+        mock.patch.object(
+            maxtext_engine.train_utils,
+            "create_training_optimizer",
+            return_value=(lambda step: jnp.array(0.001), optax.adam(0.05)),
+        ),
+        mock.patch.object(
+            maxtext_engine.model_creation_utils,
+            "from_pretrained",
+            side_effect=lambda *a, **k: (_Model(), self.mesh),
+        ),
+    ):
+      self.addCleanup(patcher.stop)
+      patcher.start()
+
+  def _engine(self, compiled: bool) -> maxtext_engine.MaxTextTrainingEngine:
+    engine = maxtext_engine.MaxTextTrainingEngine(self.config)
+    engine.with_loss_fn(_loss_fn)
+    engine.with_gen_model_input_fn(_model_input)
+    if compiled:
+      engine.compile(_payload(0))
+    return engine
+
+  def _train(self, engine, payloads, prepare_at: set[int]):
+    """Two steps of two micro-batches; micro-batches in `prepare_at` go in prepared.
+
+    A prepared micro-batch is prepared before the previous update is applied, as
+    `TrainerWorker.stage_fwd_bwd` does while that update is still on the device.
+    """
+    prepared = {i: None for i in prepare_at}
+    if 0 in prepared:
+      prepared[0] = engine.prepare_fwd_bwd_input(payloads[0])
+    metrics = []
+    for step in range(2):
+      for j in range(2):
+        i = 2 * step + j
+        engine.fwd_bwd(prepared[i] if i in prepared else payloads[i])
+        if j == 1 and i + 1 in prepared:
+          prepared[i + 1] = engine.prepare_fwd_bwd_input(payloads[i + 1])
+      engine.update()
+      buffer = engine.get_metrics(clear_cache=True)
+      metrics.append({k: np.asarray(v).tolist() for k, v in sorted(buffer.scalar_metrics.items())})
+      metrics[-1]["loss"] = np.asarray(buffer.weighted_metrics["loss"].compute()).tolist()
+    return np.asarray(engine.model.weights.value), metrics
+
+  def test_prepared_micro_batches_train_bit_identically(self):
+    payloads = [_payload(s) for s in (1, 2, 3, 4)]
+    for compiled in (True, False):
+      ref = self._engine(compiled)
+      new = self._engine(compiled)
+      with (
+          mock.patch.object(new, "_compile_for_batch", wraps=new._compile_for_batch) as new_compiles,
+          mock.patch.object(ref, "_compile_for_batch", wraps=ref._compile_for_batch) as ref_compiles,
+      ):
+        ref_w, ref_m = self._train(ref, payloads, prepare_at=set())
+        new_w, new_m = self._train(new, payloads, prepare_at={0, 2})
+      np.testing.assert_array_equal(ref_w, new_w)
+      self.assertEqual(ref_m, new_m)
+      # No kernel was rebuilt for a prepared input.
+      self.assertEqual(ref_compiles.call_count, new_compiles.call_count)
+      self.assertEqual(new.train_step, 2)
+
+  def test_the_comparison_sees_a_one_token_difference(self):
+    payloads = [_payload(s) for s in (1, 2, 3, 4)]
+    changed = list(payloads)
+    tokens = np.array(payloads[2].token_ids)
+    tokens[0, 0] += 1e-3
+    changed[2] = dataclasses.replace(payloads[2], token_ids=tokens)
+    ref_w, _ = self._train(self._engine(True), payloads, prepare_at=set())
+    new_w, _ = self._train(self._engine(True), changed, prepare_at={2})
+    self.assertFalse(np.array_equal(ref_w, new_w))
+
+  def test_staged_through_tunix_trainer_worker_trains_bit_identically(self):
+    # The path Tunix's pipelined step boundary takes: TrainerWorker.stage_fwd_bwd
+    # calls prepare_fwd_bwd_input, fwd_bwd_staged trains what it returned.
+    from tunix.experimental.common import datatypes  # pylint: disable=g-import-not-at-top
+    from tunix.experimental.worker import trainer_worker  # pylint: disable=g-import-not-at-top
+
+    payloads = [_payload(s) for s in (1, 2, 3, 4)]
+
+    def run(stage: bool):
+      engine = self._engine(True)
+      worker = trainer_worker.TrainerWorker(trainer_factory=lambda: engine)
+      worker.initialize()
+      prepare = mock.patch.object(engine, "prepare_fwd_bwd_input", wraps=engine.prepare_fwd_bwd_input)
+      with prepare as spy:
+        for i, payload in enumerate(payloads):
+          request = datatypes.TrainRequest(request_id=f"r{i}", payload=payload)
+          if stage and i % 2 == 0:
+            worker.stage_fwd_bwd(request)
+            worker.fwd_bwd_staged(f"r{i}")
+          else:
+            worker.fwd_bwd(request)
+          if i % 2 == 1:
+            worker.update()
+      return np.asarray(engine.model.weights.value), spy.call_count
+
+    ref_w, ref_prepared = run(stage=False)
+    new_w, new_prepared = run(stage=True)
+    np.testing.assert_array_equal(ref_w, new_w)
+    self.assertEqual((ref_prepared, new_prepared), (0, 2))
+
+  def test_preparing_places_the_traced_inputs_and_leaves_state_alone(self):
+    engine = self._engine(True)
+    engine.fwd_bwd(_payload(1))
+    before = (
+        engine.micro_step_count,
+        engine.train_step,
+        np.asarray(engine.model.weights.value),
+    )
+    prepared = engine.prepare_fwd_bwd_input(_payload(2))
+    self.assertIsInstance(prepared, maxtext_engine.PreparedFwdBwdInput)
+    self.assertEqual(prepared.batch["scale"], 2.0)
+    dynamic, _ = maxtext_engine._split_static_and_dynamic(prepared.batch)
+    shardings = engine._batch_data_shardings(dynamic)
+    for key in ("tokens", "mask"):
+      self.assertIsInstance(prepared.batch[key], jax.Array)
+      self.assertEqual(prepared.batch[key].sharding, shardings[key])
+    np.testing.assert_array_equal(np.asarray(prepared.batch["tokens"]), _payload(2).token_ids)
+    after = (
+        engine.micro_step_count,
+        engine.train_step,
+        np.asarray(engine.model.weights.value),
+    )
+    self.assertEqual(before[:2], after[:2])
+    np.testing.assert_array_equal(before[2], after[2])
+
+
+if __name__ == "__main__":
+  absltest.main()

From ef37a0b5f29a3dd0a3a6c408c54aa1e0fbba1a19 Mon Sep 17 00:00:00 2001
From: Param Bole <parambole@google.com>
Date: Sat, 10 Oct 2026 05:57:22 -0700
Subject: [PATCH 2/6] [MaxTextEngine] Report how long a caller waited for the
 device

wait_for_device() blocks until every computation dispatched so far has
finished and returns the seconds it waited. Tunix's TrainerWorker calls it
after staging the next step's first micro-batch (prepare_fwd_bwd_input): a
wait of ~0 means the device was already idle, so the transfer did not hide
behind the step's last micro-batch and update, and the RL loop stops staging.
It is a blocking wait rather than an is_ready() poll because under Pathways a
readiness check is an RPC to the IFRT proxy, so a poll can read "not ready"
for work that has finished.

InflightThrottler.block_until_ready() does the waiting without dequeuing and
without writing metrics, so the next wait_for_next / wait_for_all, and the
metrics it writes, are what they would have been.

Tests: training with a wait after each prepared micro-batch is bit-identical
to training without (weights, metrics, compile count, metric writes in the
same order); the wait leaves the in-flight queue as it was and returns with
every queued leaf ready; block_until_ready writes nothing and a drain after it
writes the same metrics. The review's 4-device test with Tunix's real model
input (a nested RLTrainerPayload, batch split four ways) is ported, with the
wait added.
---
 .../training_engine/inflight_throttler.py     |  11 +
 src/maxtext/training_engine/maxtext_engine.py |  14 ++
 ...text_engine_prepared_input_sharded_test.py | 225 ++++++++++++++++++
 .../maxtext_engine_prepared_input_test.py     | 109 ++++++++-
 4 files changed, 355 insertions(+), 4 deletions(-)
 create mode 100644 tests/post_training/unit/maxtext_engine_prepared_input_sharded_test.py

diff --git a/src/maxtext/training_engine/inflight_throttler.py b/src/maxtext/training_engine/inflight_throttler.py
index a809fe5e59..21fb91e012 100644
--- a/src/maxtext/training_engine/inflight_throttler.py
+++ b/src/maxtext/training_engine/inflight_throttler.py
@@ -65,6 +65,17 @@ def wait_for_next(self) -> None:
         self._flush_pending_metrics()
         self._pending_metrics = metrics
 
+  def block_until_ready(self) -> None:
+    """Waits for every inflight computation to finish, leaving the queue as it is.
+
+    Unlike `wait_for_all`, nothing is dequeued and no metrics are written: the next
+    `wait_for_next` / `wait_for_all` sees the same queue and finds it already finished.
+    """
+    with self._inflight_queue.mutex:
+      inflight = list(self._inflight_queue.queue)
+    for computation, _ in inflight:
+      jax.block_until_ready(computation)
+
   def wait_for_all(self) -> None:
     """Wait for all inflight computations to finish and log their metrics."""
     while not self._inflight_queue.empty():
diff --git a/src/maxtext/training_engine/maxtext_engine.py b/src/maxtext/training_engine/maxtext_engine.py
index 61ece83dc6..6a26a3401e 100644
--- a/src/maxtext/training_engine/maxtext_engine.py
+++ b/src/maxtext/training_engine/maxtext_engine.py
@@ -2439,6 +2439,20 @@ def prepare_fwd_bwd_input(self, payload: abstract_engine.TrainerPayload) -> Prep
       batch = {**dynamic_batch, **static_batch}
     return PreparedFwdBwdInput(batch=batch)
 
+  def wait_for_device(self) -> float:
+    """Blocks until the computations dispatched so far have finished; returns the seconds waited.
+
+    For a caller that did work after dispatching and wants to know whether it hid behind the
+    device (Tunix's `TrainerWorker.stage_fwd_bwd` reports it): ~0 means the device was already
+    idle. A blocking wait, not an `is_ready()` poll: under Pathways a readiness check is an RPC
+    to the proxy, so a poll can read "not ready" for work that has finished. Changes no state:
+    the in-flight queue and its pending metrics are left as they were (see
+    `InflightThrottler.block_until_ready`).
+    """
+    start = time.monotonic()
+    self._throttler.block_until_ready()
+    return time.monotonic() - start
+
   @_profiled_step("fwd_bwd", control_profile=True)
   def fwd_bwd(self, payload: abstract_engine.TrainerPayload | PreparedFwdBwdInput, **kwargs: Any) -> None:
     """Executes a micro-batch forward-backward pass and accumulates gradients.
diff --git a/tests/post_training/unit/maxtext_engine_prepared_input_sharded_test.py b/tests/post_training/unit/maxtext_engine_prepared_input_sharded_test.py
new file mode 100644
index 0000000000..7a3a9e001e
--- /dev/null
+++ b/tests/post_training/unit/maxtext_engine_prepared_input_sharded_test.py
@@ -0,0 +1,225 @@
+# Copyright 2023–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
+#
+#    https://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.
+
+"""prepare_fwd_bwd_input with Tunix's real model input, on a sharded mesh.
+
+`maxtext_engine_prepared_input_test` uses a one-device mesh and a flat toy
+input. Production feeds Tunix's `_algo_model_input`: {"train_example":
+RLTrainerPayload (a flax.struct pytree with static `metadata`), "algo_config":
+a closed-over object, "pad_id", "eos_id"}, on a mesh whose batch axis shards.
+Here the batch is split over 4 host devices, and a prepared micro-batch, with
+`wait_for_device` after it as Tunix's `TrainerWorker.stage_fwd_bwd` does, must
+train bit-identically to the payload it came from.
+"""
+
+import dataclasses
+import functools
+import os
+import types
+from unittest import mock
+
+# A single device shards nothing, so ask XLA for four host devices before JAX
+# initializes its backend.
+os.environ.setdefault("XLA_FLAGS", "--xla_force_host_platform_device_count=4")
+
+# pylint: disable=wrong-import-position
+from absl.testing import absltest
+from flax import nnx
+import jax
+import jax.numpy as jnp
+from maxtext.configs import pyconfig
+from maxtext.training_engine import abstract_engine
+from maxtext.training_engine import maxtext_engine
+from maxtext.utils import maxtext_utils
+from tests.utils.test_helpers import get_test_config_path
+import numpy as np
+import optax
+import pytest
+
+# training_engine imports tunix, so these tests need the post-training dependency bundle.
+pytestmark = [pytest.mark.post_training]
+
+B, P, C, L, K = 4, 8, 16, 2, 3
+
+
+class _Model(nnx.Module):
+
+  def __init__(self):
+    self.weights = nnx.Param(jnp.array([1.0, 2.0, 0.5]))
+
+
+def _payload(seed: int):
+  from tunix.experimental.common import datatypes  # pylint: disable=g-import-not-at-top
+
+  rng = np.random.default_rng(seed)
+  return datatypes.RLTrainerPayload(
+      prompt_ids=rng.integers(1, 50, (B, P)).astype(np.int32),
+      prompt_mask=np.ones((B, P), np.float32),
+      completion_ids=rng.integers(1, 50, (B, C)).astype(np.int32),
+      completion_mask=(rng.random((B, C)) > 0.2).astype(np.float32),
+      advantages=rng.normal(size=(B,)).astype(np.float32),
+      old_per_token_logps=-rng.random((B, C)).astype(np.float32),
+      routed_experts=rng.integers(0, 8, (B, P + C, L, K)).astype(np.int16),
+      metadata={"trajectory_ids": tuple(f"t{seed}_{i}" for i in range(B))},
+  )
+
+
+def _loss_fn(model, train_example, algo_config, pad_id, eos_id):
+  del eos_id
+  w = model.weights.value
+  x = train_example.completion_ids.astype(jnp.float32) / 50.0
+  m = train_example.completion_mask * (train_example.completion_ids != pad_id)
+  routed = jnp.mean(train_example.routed_experts.astype(jnp.float32), (2, 3))
+  pred = (
+      w[0] * x
+      + w[1] * x**2 * algo_config.scale
+      + w[2] * routed[:, P:] / 8.0
+      + train_example.old_per_token_logps
+  )
+  per_tok = pred * train_example.advantages[:, None] * m
+  return (
+      abstract_engine.WeightedMetric(unreduced_sum=jnp.sum(per_tok), denominator=jnp.sum(m)),
+      {},
+  )
+
+
+class NestedPayloadPreparedInputTest(absltest.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    if jax.device_count() < 4:
+      self.skipTest("needs 4 devices; set XLA_FLAGS=--xla_force_host_platform_device_count=4")
+    self.config = pyconfig.initialize(
+        [None, get_test_config_path()],
+        model_name="llama3.1-8b",
+        run_name="test_run",
+        base_output_directory=self.create_tempdir().full_path,
+        init_weights_seed=42,
+        micro_batch_size_to_train_on=B,
+        gradient_accumulation_steps=1,
+        enable_dropout=False,
+        record_internal_nn_metrics=False,
+        enable_tensorboard=False,
+        tensorboard_dir=self.create_tempdir().full_path,
+        skip_jax_distributed_system=True,
+        enable_checkpointing=False,
+        profiler_steps=0,
+    )
+    self.mesh = jax.sharding.Mesh(maxtext_utils.create_device_mesh(self.config), self.config.mesh_axes)
+    for patcher in (
+        mock.patch.object(
+            maxtext_engine.train_utils,
+            "create_training_optimizer",
+            return_value=(lambda step: jnp.array(0.001), optax.adam(0.05)),
+        ),
+        mock.patch.object(
+            maxtext_engine.model_creation_utils,
+            "from_pretrained",
+            side_effect=lambda *a, **k: (_Model(), self.mesh),
+        ),
+    ):
+      self.addCleanup(patcher.stop)
+      patcher.start()
+    self.algo_config = types.SimpleNamespace(scale=0.75)
+
+  def _engine(self):
+    from tunix.experimental.orchestrator import algorithm_adapter  # pylint: disable=g-import-not-at-top
+
+    engine = maxtext_engine.MaxTextTrainingEngine(self.config)
+    engine.with_loss_fn(_loss_fn, has_aux=True)
+    engine.with_gen_model_input_fn(
+        functools.partial(
+            algorithm_adapter._algo_model_input,  # pylint: disable=protected-access
+            algo_config=self.algo_config,
+            pad_id=0,
+            eos_id=1,
+        )
+    )
+    engine.compile(_payload(0))
+    return engine
+
+  def _train(self, engine, payloads, prepare_at):
+    """Two steps of two micro-batches; those in `prepare_at` prepared ahead."""
+    prepared = {}
+    if 0 in prepare_at:
+      prepared[0] = engine.prepare_fwd_bwd_input(payloads[0])
+      engine.wait_for_device()
+    metrics = []
+    for step in range(2):
+      for j in range(2):
+        i = 2 * step + j
+        engine.fwd_bwd(prepared.pop(i) if i in prepared else payloads[i])
+        if j == 1 and i + 1 in prepare_at:
+          # Prepared before the update is applied, as stage_fwd_bwd does, which
+          # then waits for the device to report whether that hid.
+          prepared[i + 1] = engine.prepare_fwd_bwd_input(payloads[i + 1])
+          engine.wait_for_device()
+      engine.update()
+      buf = engine.get_metrics(clear_cache=True)
+      metrics.append(np.asarray(buf.weighted_metrics["loss"].compute()).tolist())
+    return np.asarray(engine.model.weights.value), metrics
+
+  def test_bit_identical_on_a_sharded_mesh(self):
+    payloads = [_payload(s) for s in (1, 2, 3, 4)]
+    ref, new = self._engine(), self._engine()
+    with (
+        mock.patch.object(ref, "_compile_for_batch", wraps=ref._compile_for_batch) as ref_compiles,
+        mock.patch.object(new, "_compile_for_batch", wraps=new._compile_for_batch) as new_compiles,
+    ):
+      ref_w, ref_m = self._train(ref, payloads, prepare_at=set())
+      new_w, new_m = self._train(new, payloads, prepare_at={0, 2})
+    np.testing.assert_array_equal(ref_w, new_w)
+    self.assertEqual(ref_m, new_m)
+    self.assertEqual(ref_compiles.call_count, new_compiles.call_count)
+
+  def test_prepared_leaves_are_sharded_like_the_compile(self):
+    engine = self._engine()
+    prepared = engine.prepare_fwd_bwd_input(_payload(2))
+    dynamic, static = maxtext_engine._split_static_and_dynamic(prepared.batch)  # pylint: disable=protected-access
+    self.assertEqual(set(dynamic), {"train_example"})
+    self.assertEqual(set(static), {"algo_config", "pad_id", "eos_id"})
+    shardings = engine._batch_data_shardings(dynamic)  # pylint: disable=protected-access
+    example, sharding = dynamic["train_example"], shardings["train_example"]
+    self.assertEqual(example.metadata, {})  # stripped by _prepare_batch
+    split_2d = 0
+    for field in (
+        "completion_ids",
+        "completion_mask",
+        "old_per_token_logps",
+        "routed_experts",
+        "advantages",
+        "prompt_ids",
+    ):
+      leaf, want = getattr(example, field), getattr(sharding, field)
+      self.assertIsInstance(leaf, jax.Array, field)
+      self.assertEqual(leaf.sharding, want, field)
+      rows = {s.index[0] for s in leaf.addressable_shards}
+      if leaf.ndim >= 2 and len(rows) == 4:
+        split_2d += 1
+    # Not vacuous: the 2-D leaves really are split four ways.
+    self.assertGreaterEqual(split_2d, 3)
+
+  def test_one_token_change_is_seen(self):
+    payloads = [_payload(s) for s in (1, 2, 3, 4)]
+    changed = list(payloads)
+    ids = np.array(payloads[2].completion_ids)
+    ids[0, 0] = ids[0, 0] % 49 + 1
+    changed[2] = dataclasses.replace(payloads[2], completion_ids=ids)
+    ref_w, _ = self._train(self._engine(), payloads, prepare_at=set())
+    new_w, _ = self._train(self._engine(), changed, prepare_at={2})
+    self.assertFalse(np.array_equal(ref_w, new_w))
+
+
+if __name__ == "__main__":
+  absltest.main()
diff --git a/tests/post_training/unit/maxtext_engine_prepared_input_test.py b/tests/post_training/unit/maxtext_engine_prepared_input_test.py
index c3401a9ad4..5b2366c7ed 100644
--- a/tests/post_training/unit/maxtext_engine_prepared_input_test.py
+++ b/tests/post_training/unit/maxtext_engine_prepared_input_test.py
@@ -12,14 +12,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Tests for MaxTextTrainingEngine.prepare_fwd_bwd_input.
+"""Tests for MaxTextTrainingEngine.prepare_fwd_bwd_input and wait_for_device.
 
 A micro-batch prepared ahead (mapped and placed on the mesh while earlier work
 still runs) must train exactly like the payload it came from: same compiled
-kernel, same inputs, bit-identical weights and metrics.
+kernel, same inputs, bit-identical weights and metrics. Waiting for the device
+afterwards (what Tunix's `TrainerWorker.stage_fwd_bwd` does, to report whether
+the transfer hid behind the device) must change nothing either.
 """
 
 import dataclasses
+import types
 from typing import Any
 from unittest import mock
 
@@ -30,6 +33,7 @@
 import jax.numpy as jnp
 from maxtext.configs import pyconfig
 from maxtext.training_engine import abstract_engine
+from maxtext.training_engine import inflight_throttler
 from maxtext.training_engine import maxtext_engine
 from maxtext.utils import maxtext_utils
 from tests.utils.test_helpers import get_test_config_path
@@ -124,15 +128,18 @@ def _engine(self, compiled: bool) -> maxtext_engine.MaxTextTrainingEngine:
       engine.compile(_payload(0))
     return engine
 
-  def _train(self, engine, payloads, prepare_at: set[int]):
+  def _train(self, engine, payloads, prepare_at: set[int], wait: bool = False):
     """Two steps of two micro-batches; micro-batches in `prepare_at` go in prepared.
 
     A prepared micro-batch is prepared before the previous update is applied, as
     `TrainerWorker.stage_fwd_bwd` does while that update is still on the device.
+    With `wait`, `wait_for_device` follows each prepare, as it does there too.
     """
     prepared = {i: None for i in prepare_at}
     if 0 in prepared:
       prepared[0] = engine.prepare_fwd_bwd_input(payloads[0])
+      if wait:
+        engine.wait_for_device()
     metrics = []
     for step in range(2):
       for j in range(2):
@@ -140,6 +147,11 @@ def _train(self, engine, payloads, prepare_at: set[int]):
         engine.fwd_bwd(prepared[i] if i in prepared else payloads[i])
         if j == 1 and i + 1 in prepared:
           prepared[i + 1] = engine.prepare_fwd_bwd_input(payloads[i + 1])
+          if wait:
+            queued = engine._throttler._inflight_queue.qsize()
+            self.assertGreaterEqual(engine.wait_for_device(), 0.0)
+            # The in-flight queue is left for the update to drain as before.
+            self.assertEqual(engine._throttler._inflight_queue.qsize(), queued)
       engine.update()
       buffer = engine.get_metrics(clear_cache=True)
       metrics.append({k: np.asarray(v).tolist() for k, v in sorted(buffer.scalar_metrics.items())})
@@ -163,6 +175,46 @@ def test_prepared_micro_batches_train_bit_identically(self):
       self.assertEqual(ref_compiles.call_count, new_compiles.call_count)
       self.assertEqual(new.train_step, 2)
 
+  def test_waiting_for_the_device_changes_nothing(self):
+    payloads = [_payload(s) for s in (1, 2, 3, 4)]
+    for compiled in (True, False):
+      ref = self._engine(compiled)
+      new = self._engine(compiled)
+      with (
+          mock.patch.object(new, "_compile_for_batch", wraps=new._compile_for_batch) as new_compiles,
+          mock.patch.object(ref, "_compile_for_batch", wraps=ref._compile_for_batch) as ref_compiles,
+          mock.patch.object(new._metrics_logger, "write_metrics", wraps=new._metrics_logger.write_metrics) as new_w,
+          mock.patch.object(ref._metrics_logger, "write_metrics", wraps=ref._metrics_logger.write_metrics) as ref_w,
+      ):
+        ref_weights, ref_m = self._train(ref, payloads, prepare_at=set())
+        new_weights, new_m = self._train(new, payloads, prepare_at={0, 2}, wait=True)
+      np.testing.assert_array_equal(ref_weights, new_weights)
+      self.assertEqual(ref_m, new_m)
+      self.assertEqual(ref_compiles.call_count, new_compiles.call_count)
+      # The same metrics writes, in the same order (the wait writes none).
+      self.assertEqual(
+          [c.args[0].id for c in ref_w.call_args_list],
+          [c.args[0].id for c in new_w.call_args_list],
+      )
+
+  def test_wait_for_device_waits_for_the_update_and_keeps_the_queue(self):
+    engine = self._engine(True)
+    engine.fwd_bwd(_payload(1))
+    engine.fwd_bwd(_payload(2))
+    engine.update()
+    queue = engine._throttler._inflight_queue
+    queued = list(queue.queue)
+    self.assertNotEmpty(queued)
+    waited = engine.wait_for_device()
+    self.assertIsInstance(waited, float)
+    self.assertGreaterEqual(waited, 0.0)
+    self.assertEqual(list(queue.queue), queued)
+    # Everything dispatched before the wait, the update included, is done.
+    for leaves, _ in queued:
+      for leaf in leaves:
+        if isinstance(leaf, jax.Array):
+          self.assertTrue(leaf.is_ready())
+
   def test_the_comparison_sees_a_one_token_difference(self):
     payloads = [_payload(s) for s in (1, 2, 3, 4)]
     changed = list(payloads)
@@ -190,7 +242,9 @@ def run(stage: bool):
         for i, payload in enumerate(payloads):
           request = datatypes.TrainRequest(request_id=f"r{i}", payload=payload)
           if stage and i % 2 == 0:
-            worker.stage_fwd_bwd(request)
+            resp = worker.stage_fwd_bwd(request)
+            # The engine's wait_for_device answered how long the device ran on.
+            self.assertGreaterEqual(resp.metadata["device_wait_s"], 0.0)
             worker.fwd_bwd_staged(f"r{i}")
           else:
             worker.fwd_bwd(request)
@@ -229,5 +283,52 @@ def test_preparing_places_the_traced_inputs_and_leaves_state_alone(self):
     np.testing.assert_array_equal(before[2], after[2])
 
 
+class _Metrics:
+
+  def __init__(self):
+    self.written = []
+
+  def write_metrics(self, metrics, **kwargs):
+    del kwargs
+    self.written.append(metrics)
+
+
+class InflightThrottlerBlockUntilReadyTest(absltest.TestCase):
+
+  def _throttler(self):
+    logger = _Metrics()
+    config = types.SimpleNamespace(max_inflight_computations=2)
+    return inflight_throttler.InflightThrottler(config, logger), logger
+
+  def test_block_until_ready_neither_dequeues_nor_writes(self):
+    throttler, logger = self._throttler()
+    throttler.add_computation(jnp.ones(3) * 2, "m1")
+    throttler.add_computation(jnp.ones(3) * 3, None)
+    throttler.block_until_ready()
+    self.assertEqual(throttler._inflight_queue.qsize(), 2)
+    self.assertEqual(logger.written, [])
+    for leaves, _ in list(throttler._inflight_queue.queue):
+      self.assertTrue(all(leaf.is_ready() for leaf in leaves))
+
+  def test_a_drain_after_it_writes_what_it_would_have(self):
+    written = []
+    for block_first in (False, True):
+      throttler, logger = self._throttler()
+      throttler.add_computation(jnp.ones(3), "m1")
+      throttler.add_computation(jnp.ones(3), "m2")
+      if block_first:
+        throttler.block_until_ready()
+      throttler.wait_for_all()
+      written.append(logger.written)
+      self.assertTrue(throttler._inflight_queue.empty())
+    self.assertEqual(written[0], ["m1", "m2"])
+    self.assertEqual(written[0], written[1])
+
+  def test_block_until_ready_on_an_empty_queue_returns(self):
+    throttler, logger = self._throttler()
+    throttler.block_until_ready()
+    self.assertEqual(logger.written, [])
+
+
 if __name__ == "__main__":
   absltest.main()

From b72884c76ec59d66af7823f955eb006b9127e89f Mon Sep 17 00:00:00 2001
From: Param Bole <parambole@google.com>
Date: Tue, 13 Oct 2026 17:40:11 -0700
Subject: [PATCH 3/6] [MaxTextEngine] Test prepared micro-batches under
 dropless replay and the fused update

Brings prepare_fwd_bwd_input / wait_for_device up to the current engine
(per-micro-batch dropless replay, the final micro-batch's dropless program,
Tunix's fused final update) and makes their tests hold there:

- maxtext_engine_prepared_input_dropless_test.py (new): on the one-layer
  Qwen3.5 MoE rig of the dropless tests, a prepared micro-batch that overflows
  is replayed from its prepared inputs, the step's last micro-batch runs the
  final program, and weights, losses, gradient norms and replay counts are
  bit-identical to the unprepared run, whether it is prepared before or after
  the update; also when the run's first micro-batch is prepared. A control
  that prepares the wrong micro-batch is caught.
- maxtext_engine_prepared_input_test.py: prepares before and after the update
  (the second is where Tunix's stage lands once the update runs inside the
  last train call); the TrainerWorker test drives fwd_bwd(apply_optimizer=True,
  wait_for_update=False) then stage_fwd_bwd, and skips when the installed
  Tunix has no staging; wait_for_device is checked to wait on every queued
  computation, the update's gradient norm included, by recording what it
  blocked on rather than by is_ready() after the fact.
- maxtext_engine_prepared_input_sharded_test.py: re-execs itself with four
  host devices, as maxtext_engine_data_parallel_test.py does; setting
  XLA_FLAGS at import time left one device under pytest and skipped green.
- README: the two methods and the tests.
---
 src/maxtext/training_engine/README.md         |   3 +
 src/maxtext/training_engine/maxtext_engine.py |  12 +-
 ...ext_engine_prepared_input_dropless_test.py | 172 ++++++++++++++
 ...text_engine_prepared_input_sharded_test.py | 124 ++++++----
 .../maxtext_engine_prepared_input_test.py     | 219 +++++++++++-------
 5 files changed, 403 insertions(+), 127 deletions(-)
 create mode 100644 tests/post_training/unit/maxtext_engine_prepared_input_dropless_test.py

diff --git a/src/maxtext/training_engine/README.md b/src/maxtext/training_engine/README.md
index 9793eec4f3..5e08df0117 100644
--- a/src/maxtext/training_engine/README.md
+++ b/src/maxtext/training_engine/README.md
@@ -296,6 +296,8 @@ and does not use the engine.
 | `model_scope(*inputs)` | none; the caller jits | inputs | yields the model and the placed inputs under the sharding rules |
 | `fwd_only(fn, *inputs)` | none; `fn` jits | inputs | `fn(model, ...)` |
 | `eval_step(batch)` | eval | params, rest, batch | loss, aux |
+| `prepare_fwd_bwd_input(micro_batch)` | none; `device_put` | micro-batch | a `PreparedFwdBwdInput`: the micro-batch mapped to loss inputs and placed on `fwd_bwd`'s input shardings, which `fwd_bwd` takes in its place |
+| `wait_for_device()` | none | | the seconds it blocked until the in-flight `fwd_bwd`s and `update` finished; dequeues nothing and writes no metrics |
 
 Every micro-batch runs the same `fwd_bwd` program. The first micro-batch's gradients start the
 accumulator, and `accumulate` adds each later one's to it in place. The accumulator is not passed to
@@ -383,6 +385,7 @@ The engine's tests are in `tests/post_training/unit/` and run on CPU; they requi
 | `maxtext_engine_model_build_test.py` | model construction with and without a mesh set by the caller |
 | `maxtext_engine_checkpoint_test.py` | resuming from a mid-step checkpoint through Orbax reproduces the uninterrupted step |
 | `maxtext_engine_grad_cast_test.py` | `cast_grads_after_all_reduce`: which gradients leave `fwd_bwd` uncast, and that moving their cast changes no number on CPU |
+| `maxtext_engine_prepared_input{,_dropless,_sharded}_test.py` | `prepare_fwd_bwd_input` and `wait_for_device`: a prepared micro-batch trains bit-identically on one device, through Tunix's `TrainerWorker` staging, under the dropless replay and final dropless programs, and on a four-device mesh with Tunix's payload; the wait leaves the in-flight queue and the metrics writes as they were |
 | `maxtext_engine_{constructor,data_parallel,xaot,packing,profiling}_test.py`, `router_replay_engine_test.py` | construction, data parallelism, ahead-of-time compilation, sequence packing, profiling, router replay |
 
 ## Differences from pre_train/train.py
diff --git a/src/maxtext/training_engine/maxtext_engine.py b/src/maxtext/training_engine/maxtext_engine.py
index 6a26a3401e..be10518f87 100644
--- a/src/maxtext/training_engine/maxtext_engine.py
+++ b/src/maxtext/training_engine/maxtext_engine.py
@@ -2440,14 +2440,16 @@ def prepare_fwd_bwd_input(self, payload: abstract_engine.TrainerPayload) -> Prep
     return PreparedFwdBwdInput(batch=batch)
 
   def wait_for_device(self) -> float:
-    """Blocks until the computations dispatched so far have finished; returns the seconds waited.
+    """Blocks until the computations this engine holds in flight have finished; returns the seconds waited.
 
     For a caller that did work after dispatching and wants to know whether it hid behind the
     device (Tunix's `TrainerWorker.stage_fwd_bwd` reports it): ~0 means the device was already
-    idle. A blocking wait, not an `is_ready()` poll: under Pathways a readiness check is an RPC
-    to the proxy, so a poll can read "not ready" for work that has finished. Changes no state:
-    the in-flight queue and its pending metrics are left as they were (see
-    `InflightThrottler.block_until_ready`).
+    idle. It waits on what the in-flight throttler holds: each queued `fwd_bwd`'s loss and the
+    last `update`'s gradient norm, so the update counts. Anything dispatched before those was
+    waited on when it left the queue. A blocking wait, not an `is_ready()` poll: under Pathways
+    a readiness check is an RPC to the proxy, so a poll can read "not ready" for work that has
+    finished. Changes no state: the in-flight queue and its pending metrics are left as they
+    were (see `InflightThrottler.block_until_ready`).
     """
     start = time.monotonic()
     self._throttler.block_until_ready()
diff --git a/tests/post_training/unit/maxtext_engine_prepared_input_dropless_test.py b/tests/post_training/unit/maxtext_engine_prepared_input_dropless_test.py
new file mode 100644
index 0000000000..f9afa5547e
--- /dev/null
+++ b/tests/post_training/unit/maxtext_engine_prepared_input_dropless_test.py
@@ -0,0 +1,172 @@
+# 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
+#
+#    https://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.
+"""prepare_fwd_bwd_input under the dropless replay and the final micro-batch's dropless program.
+
+The MLPerf recipes train with `retry_when_tokens_dropped` (a micro-batch whose MoE ragged buffer overflowed is
+accumulated as zero and rerun dropless, from the inputs it ran on) and `retry_dropless_final_micro_batch` (a step's
+last micro-batch runs the dropless program directly). A micro-batch prepared ahead must take both paths exactly like
+the payload it came from: the overflowing micro-batch is replayed from its prepared inputs, the final one runs the
+final program, and the weights, losses and replay counts are bit-identical.
+
+The rig is `maxtext_engine_final_dropless_test.py`'s: a one-layer Qwen3.5 MoE on CPU whose loss reports the overflow
+flag RoutedMoE would, from which program its graph is, and scales the loss per program so the kernels can be told
+apart.
+"""
+# pylint: disable=protected-access
+
+import os
+import unittest
+from unittest import mock
+
+from flax import nnx
+import jax
+import jax.numpy as jnp
+from jax.sharding import Mesh
+from maxtext.models import models
+from maxtext.training_engine import maxtext_engine
+from maxtext.utils import maxtext_utils
+import numpy as np
+import pytest
+
+from tests.post_training.unit.maxtext_engine_dropless_replay_test import _cfg, _payload, _MARKED
+from tests.post_training.unit.maxtext_engine_final_dropless_test import _program
+
+pytestmark = [pytest.mark.post_training]
+
+# Per-program loss multipliers, so the losses say which kernel ran each micro-batch.
+_SCALES = {"regular": 1.0, "replay": 2.0, "final": 3.0}
+
+# Two steps of three micro-batches. Step 1's first micro-batch overflows, so the one prepared ahead is replayed.
+_STEPS = ((10, _MARKED, 30), (_MARKED, 20, 30))
+
+# Where the next step's first micro-batch is prepared, relative to the step's update.
+_BEFORE_UPDATE = "before_update"
+_AFTER_UPDATE = "after_update"  # where Tunix's stage lands when the update runs inside the last train call
+
+
+class PreparedInputDroplessTest(unittest.TestCase):
+
+  def setUp(self):
+    os.environ["NEW_MODEL_DESIGN"] = "1"
+    os.environ["SKIP_JAX_PRECOMPILE"] = "1"
+
+  def _engine(self):
+    """An engine on the one-layer MoE rig, with the dropless replay and the final dropless program on."""
+    cfg = _cfg(override_logical_axis_rules=True, ragged_buffer_factor=2.5)
+    mesh = Mesh(maxtext_utils.create_device_mesh(cfg), cfg.mesh_axes)
+    real_model = models.Transformer(config=cfg, mesh=mesh, quant=None, model_mode="train", rngs=nnx.Rngs(0))
+    with mock.patch.object(maxtext_engine.model_creation_utils, "from_pretrained", return_value=real_model):
+      engine = maxtext_engine.MaxTextTrainingEngine(cfg, mesh=mesh)
+    # The ring-of-experts prerequisites of retry_when_tokens_dropped cannot run on CPU, so the switches are set
+    # directly, as the dropless replay and final-dropless tests do.
+    engine._dropless_replay = True
+    engine._final_dropless = True
+    base_loss = maxtext_engine.make_router_replay_loss_fn(cfg)
+
+    def loss_fn(model, **batch):
+      program = _program(model)
+      loss, aux = base_loss(model, **batch)
+      aux = {k: v for k, v in aux.items() if k != "has_moe_overflow"}
+      aux["has_moe_overflow"] = jnp.logical_and(batch["inputs"][0, 0] == _MARKED, program == "regular")
+      aux["xent_sum"] = aux["xent_sum"] * _SCALES[program]
+      return loss, aux
+
+    engine.with_gen_model_input_fn(maxtext_engine.router_replay_gen_model_input_fn)
+    engine.with_loss_fn(loss_fn, has_aux=True)
+    engine.compile(None)  # As Tunix's TrainerWorker does at bring-up: the kernels compile on the first fwd_bwd.
+    return engine
+
+  def _train(self, engine, prepare_ahead: bool, order: str = _AFTER_UPDATE, prepare_first: bool = False):
+    """Runs `_STEPS`; returns the final parameters, per-step metrics and the prepared micro-batches' replays.
+
+    With `prepare_ahead`, each step's first micro-batch after the first step is prepared right after the previous
+    step's last `fwd_bwd` (`apply_optimizer=True`, so it ran the final program), before or after the update, and
+    `wait_for_device` follows, as Tunix's `TrainerWorker.stage_fwd_bwd` does.
+    """
+    prepared = None
+    if prepare_first:
+      prepared = engine.prepare_fwd_bwd_input(_payload(_STEPS[0][0]))
+    metrics = []
+    for step, tokens in enumerate(_STEPS):
+      for j, first_token in enumerate(tokens):
+        final = j == len(tokens) - 1
+        if j == 0 and prepared is not None:
+          self.assertIsInstance(prepared, maxtext_engine.PreparedFwdBwdInput)
+          engine.fwd_bwd(prepared, apply_optimizer=False)
+          prepared = None
+        else:
+          engine.fwd_bwd(_payload(first_token), apply_optimizer=final)
+      has_next = step + 1 < len(_STEPS)
+      if prepare_ahead and has_next and order == _BEFORE_UPDATE:
+        prepared = engine.prepare_fwd_bwd_input(_payload(_STEPS[step + 1][0]))
+        self.assertGreaterEqual(engine.wait_for_device(), 0.0)
+      with mock.patch.object(engine, "record_metrics", wraps=engine.record_metrics) as record:
+        engine.update()
+      recorded = {call.args[0]: call.args[1] for call in record.call_args_list}
+      if prepare_ahead and has_next and order == _AFTER_UPDATE:
+        prepared = engine.prepare_fwd_bwd_input(_payload(_STEPS[step + 1][0]))
+        self.assertGreaterEqual(engine.wait_for_device(), 0.0)
+      buffer = engine.get_metrics(clear_cache=True)
+      metrics.append(
+          {
+              "replays": int(recorded["moe_dropless_replays"]),
+              "loss": np.asarray(buffer.weighted_metrics["loss"].compute()).tolist(),
+              "gradient_norm": np.asarray(buffer.scalar_metrics["gradient_norm"]).tolist(),
+          }
+      )
+    self.assertIsNone(prepared)
+    params = [np.asarray(leaf) for leaf in jax.tree.leaves(nnx.state(engine.model, nnx.Param))]
+    return params, metrics
+
+  def _assert_same_training(self, ref, new):
+    ref_params, ref_metrics = ref
+    new_params, new_metrics = new
+    self.assertEqual(ref_metrics, new_metrics)
+    self.assertEqual(len(ref_params), len(new_params))
+    for r, n in zip(ref_params, new_params):
+      np.testing.assert_array_equal(r, n)
+
+  def test_a_prepared_micro_batch_is_replayed_and_trains_bit_identically(self):
+    reference = self._train(self._engine(), prepare_ahead=False)
+    # Not vacuous: one replay per step (the marked micro-batches), and step 1's is the prepared one.
+    self.assertEqual([m["replays"] for m in reference[1]], [1, 1])
+    for order in (_AFTER_UPDATE, _BEFORE_UPDATE):
+      with self.subTest(order=order):
+        engine = self._engine()
+        with mock.patch.object(engine, "_compile_for_batch", wraps=engine._compile_for_batch) as compiles:
+          prepared = self._train(engine, prepare_ahead=True, order=order)
+        self._assert_same_training(reference, prepared)
+        self.assertEqual(compiles.call_count, 1)  # no kernel was rebuilt for the prepared input
+
+  def test_the_first_micro_batch_of_the_run_can_be_prepared(self):
+    # Tunix never stages a run's first micro-batch, but the background compiles that start from it must still
+    # produce the kernels a payload would.
+    reference = self._train(self._engine(), prepare_ahead=False)
+    prepared = self._train(self._engine(), prepare_ahead=True, prepare_first=True)
+    self._assert_same_training(reference, prepared)
+
+  def test_the_comparison_sees_a_different_prepared_micro_batch(self):
+    # Control: preparing the wrong micro-batch (step 1's second instead of its first) is caught.
+    reference = self._train(self._engine(), prepare_ahead=False)
+    engine = self._engine()
+    real_prepare = engine.prepare_fwd_bwd_input
+    with mock.patch.object(
+        engine, "prepare_fwd_bwd_input", side_effect=lambda payload: real_prepare(_payload(_STEPS[1][1]))
+    ):
+      wrong = self._train(engine, prepare_ahead=True)
+    self.assertNotEqual(reference[1], wrong[1])
+
+
+if __name__ == "__main__":
+  unittest.main()
diff --git a/tests/post_training/unit/maxtext_engine_prepared_input_sharded_test.py b/tests/post_training/unit/maxtext_engine_prepared_input_sharded_test.py
index 7a3a9e001e..3ac9365492 100644
--- a/tests/post_training/unit/maxtext_engine_prepared_input_sharded_test.py
+++ b/tests/post_training/unit/maxtext_engine_prepared_input_sharded_test.py
@@ -1,4 +1,4 @@
-# Copyright 2023–2026 Google LLC
+# 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.
@@ -12,28 +12,33 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""prepare_fwd_bwd_input with Tunix's real model input, on a sharded mesh.
+"""prepare_fwd_bwd_input with Tunix's real model input, on a four-device mesh.
 
-`maxtext_engine_prepared_input_test` uses a one-device mesh and a flat toy
-input. Production feeds Tunix's `_algo_model_input`: {"train_example":
-RLTrainerPayload (a flax.struct pytree with static `metadata`), "algo_config":
-a closed-over object, "pad_id", "eos_id"}, on a mesh whose batch axis shards.
-Here the batch is split over 4 host devices, and a prepared micro-batch, with
-`wait_for_device` after it as Tunix's `TrainerWorker.stage_fwd_bwd` does, must
-train bit-identically to the payload it came from.
+`maxtext_engine_prepared_input_test` uses a one-device mesh and a flat toy input. Production feeds Tunix's
+`_algo_model_input`: {"train_example": RLTrainerPayload (a flax.struct pytree with static `metadata`),
+"algo_config": a closed-over object, "pad_id", "eos_id"}, on a mesh whose batch axis shards. Here the batch is split
+over four host devices, and a prepared micro-batch, with `wait_for_device` after it as Tunix's
+`TrainerWorker.stage_fwd_bwd` does, must train bit-identically to the payload it came from.
+
+Everything runs in a subprocess, for the reason `maxtext_engine_data_parallel_test.py` gives: the CPU backend reads
+`--xla_force_host_platform_device_count` only when it initializes, which has already happened by the time pytest
+imports this file, so setting the flag here would leave one device and skip the tests green.
+`test_prepared_input_on_a_four_device_cpu_mesh` re-execs the module with the flag appended and refuses to pass unless
+the child reports tests actually run.
 """
+# pylint: disable=protected-access
 
 import dataclasses
 import functools
 import os
+import re
+import subprocess
+import sys
 import types
+import unittest
 from unittest import mock
 
-# A single device shards nothing, so ask XLA for four host devices before JAX
-# initializes its backend.
-os.environ.setdefault("XLA_FLAGS", "--xla_force_host_platform_device_count=4")
-
-# pylint: disable=wrong-import-position
+from absl import flags
 from absl.testing import absltest
 from flax import nnx
 import jax
@@ -42,17 +47,44 @@
 from maxtext.training_engine import abstract_engine
 from maxtext.training_engine import maxtext_engine
 from maxtext.utils import maxtext_utils
-from tests.utils.test_helpers import get_test_config_path
 import numpy as np
 import optax
 import pytest
 
+from tests.utils.test_helpers import get_test_config_path
+
 # training_engine imports tunix, so these tests need the post-training dependency bundle.
-pytestmark = [pytest.mark.post_training]
+pytestmark = [pytest.mark.post_training, pytest.mark.cpu_only]
+
+_REQUIRED_DEVICES = 4
+_SENTINEL = "MAXTEXT_ENGINE_PREPARED_INPUT_SHARDED_TESTS_PASSED"
+_RAN = re.compile(rf"{_SENTINEL} ran=(\d+)")
 
 B, P, C, L, K = 4, 8, 16, 2, 3
 
 
+@pytest.mark.post_training
+@pytest.mark.cpu_only
+def test_prepared_input_on_a_four_device_cpu_mesh():
+  """The only test pytest collects here; the class below runs inside the child process."""
+  env = os.environ.copy()
+  env["XLA_FLAGS"] = f"{env.get('XLA_FLAGS', '')} --xla_force_host_platform_device_count={_REQUIRED_DEVICES}".strip()
+  env["JAX_PLATFORMS"] = "cpu"
+  # The child imports `tests.utils.test_helpers`, which pytest puts on the path for us and a bare interpreter does
+  # not.
+  repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", ".."))
+  env["PYTHONPATH"] = os.pathsep.join([repo_root, env["PYTHONPATH"]]) if env.get("PYTHONPATH") else repo_root
+
+  result = subprocess.run([sys.executable, __file__], env=env, capture_output=True, text=True, check=False)
+
+  report = f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}"
+  assert result.returncode == 0, report
+  ran = _RAN.search(result.stdout)
+  # An exit status of 0 is also what a run that skipped everything produces.
+  assert ran, f"the child did not report a completed run\n{report}"
+  assert int(ran.group(1)) > 0, f"every test in the child skipped\n{report}"
+
+
 class _Model(nnx.Module):
 
   def __init__(self):
@@ -60,7 +92,7 @@ def __init__(self):
 
 
 def _payload(seed: int):
-  from tunix.experimental.common import datatypes  # pylint: disable=g-import-not-at-top
+  from tunix.experimental.common import datatypes  # pylint: disable=import-outside-toplevel
 
   rng = np.random.default_rng(seed)
   return datatypes.RLTrainerPayload(
@@ -77,16 +109,11 @@ def _payload(seed: int):
 
 def _loss_fn(model, train_example, algo_config, pad_id, eos_id):
   del eos_id
-  w = model.weights.value
+  w = model.weights[...]
   x = train_example.completion_ids.astype(jnp.float32) / 50.0
   m = train_example.completion_mask * (train_example.completion_ids != pad_id)
   routed = jnp.mean(train_example.routed_experts.astype(jnp.float32), (2, 3))
-  pred = (
-      w[0] * x
-      + w[1] * x**2 * algo_config.scale
-      + w[2] * routed[:, P:] / 8.0
-      + train_example.old_per_token_logps
-  )
+  pred = w[0] * x + w[1] * x**2 * algo_config.scale + w[2] * routed[:, P:] / 8.0 + train_example.old_per_token_logps
   per_tok = pred * train_example.advantages[:, None] * m
   return (
       abstract_engine.WeightedMetric(unreduced_sum=jnp.sum(per_tok), denominator=jnp.sum(m)),
@@ -96,10 +123,13 @@ def _loss_fn(model, train_example, algo_config, pad_id, eos_id):
 
 class NestedPayloadPreparedInputTest(absltest.TestCase):
 
+  __test__ = False  # collected only via the subprocess entry point at the top of this file.
+
   def setUp(self):
+    """A four-device mesh and a toy three-weight model, built without a checkpoint."""
     super().setUp()
-    if jax.device_count() < 4:
-      self.skipTest("needs 4 devices; set XLA_FLAGS=--xla_force_host_platform_device_count=4")
+    if jax.device_count() < _REQUIRED_DEVICES:
+      self.skipTest(f"needs {_REQUIRED_DEVICES} devices; run through pytest, which re-execs with them")
     self.config = pyconfig.initialize(
         [None, get_test_config_path()],
         model_name="llama3.1-8b",
@@ -134,13 +164,14 @@ def setUp(self):
     self.algo_config = types.SimpleNamespace(scale=0.75)
 
   def _engine(self):
-    from tunix.experimental.orchestrator import algorithm_adapter  # pylint: disable=g-import-not-at-top
+    """An engine fed through Tunix's `_algo_model_input`, compiled on the first payload."""
+    from tunix.experimental.orchestrator import algorithm_adapter  # pylint: disable=import-outside-toplevel
 
     engine = maxtext_engine.MaxTextTrainingEngine(self.config)
     engine.with_loss_fn(_loss_fn, has_aux=True)
     engine.with_gen_model_input_fn(
         functools.partial(
-            algorithm_adapter._algo_model_input,  # pylint: disable=protected-access
+            algorithm_adapter._algo_model_input,
             algo_config=self.algo_config,
             pad_id=0,
             eos_id=1,
@@ -150,7 +181,7 @@ def _engine(self):
     return engine
 
   def _train(self, engine, payloads, prepare_at):
-    """Two steps of two micro-batches; those in `prepare_at` prepared ahead."""
+    """Two steps of two micro-batches; those in `prepare_at` prepared ahead, after the previous update."""
     prepared = {}
     if 0 in prepare_at:
       prepared[0] = engine.prepare_fwd_bwd_input(payloads[0])
@@ -160,15 +191,15 @@ def _train(self, engine, payloads, prepare_at):
       for j in range(2):
         i = 2 * step + j
         engine.fwd_bwd(prepared.pop(i) if i in prepared else payloads[i])
-        if j == 1 and i + 1 in prepare_at:
-          # Prepared before the update is applied, as stage_fwd_bwd does, which
-          # then waits for the device to report whether that hid.
-          prepared[i + 1] = engine.prepare_fwd_bwd_input(payloads[i + 1])
-          engine.wait_for_device()
       engine.update()
+      if 2 * step + 2 in prepare_at:
+        # Prepared while the update is still on the device, as stage_fwd_bwd does, which then waits for the device
+        # to report whether that hid.
+        prepared[2 * step + 2] = engine.prepare_fwd_bwd_input(payloads[2 * step + 2])
+        engine.wait_for_device()
       buf = engine.get_metrics(clear_cache=True)
       metrics.append(np.asarray(buf.weighted_metrics["loss"].compute()).tolist())
-    return np.asarray(engine.model.weights.value), metrics
+    return np.asarray(engine.model.weights[...]), metrics
 
   def test_bit_identical_on_a_sharded_mesh(self):
     payloads = [_payload(s) for s in (1, 2, 3, 4)]
@@ -186,10 +217,10 @@ def test_bit_identical_on_a_sharded_mesh(self):
   def test_prepared_leaves_are_sharded_like_the_compile(self):
     engine = self._engine()
     prepared = engine.prepare_fwd_bwd_input(_payload(2))
-    dynamic, static = maxtext_engine._split_static_and_dynamic(prepared.batch)  # pylint: disable=protected-access
+    dynamic, static = maxtext_engine._split_static_and_dynamic(prepared.batch)
     self.assertEqual(set(dynamic), {"train_example"})
     self.assertEqual(set(static), {"algo_config", "pad_id", "eos_id"})
-    shardings = engine._batch_data_shardings(dynamic)  # pylint: disable=protected-access
+    shardings = engine._batch_data_shardings(dynamic)
     example, sharding = dynamic["train_example"], shardings["train_example"]
     self.assertEqual(example.metadata, {})  # stripped by _prepare_batch
     split_2d = 0
@@ -205,7 +236,7 @@ def test_prepared_leaves_are_sharded_like_the_compile(self):
       self.assertIsInstance(leaf, jax.Array, field)
       self.assertEqual(leaf.sharding, want, field)
       rows = {s.index[0] for s in leaf.addressable_shards}
-      if leaf.ndim >= 2 and len(rows) == 4:
+      if leaf.ndim >= 2 and len(rows) == _REQUIRED_DEVICES:
         split_2d += 1
     # Not vacuous: the 2-D leaves really are split four ways.
     self.assertGreaterEqual(split_2d, 3)
@@ -221,5 +252,20 @@ def test_one_token_change_is_seen(self):
     self.assertFalse(np.array_equal(ref_w, new_w))
 
 
+_SUITE = (NestedPayloadPreparedInputTest,)
+
+
 if __name__ == "__main__":
-  absltest.main()
+  if jax.device_count() < _REQUIRED_DEVICES:
+    raise SystemExit(
+        f"needs {_REQUIRED_DEVICES} devices, got {jax.device_count()}; run this through pytest, which sets "
+        f"XLA_FLAGS=--xla_force_host_platform_device_count={_REQUIRED_DEVICES}"
+    )
+  flags.FLAGS(sys.argv[:1])  # absltest's create_tempdir reads --test_tmpdir; a bare unittest run parses no flags.
+  _loader = unittest.defaultTestLoader
+  _result = unittest.TextTestRunner(verbosity=2).run(
+      unittest.TestSuite(_loader.loadTestsFromTestCase(cls) for cls in _SUITE)
+  )
+  if not _result.wasSuccessful():
+    sys.exit(1)
+  print(f"{_SENTINEL} ran={_result.testsRun - len(_result.skipped)}")
diff --git a/tests/post_training/unit/maxtext_engine_prepared_input_test.py b/tests/post_training/unit/maxtext_engine_prepared_input_test.py
index 5b2366c7ed..6552caf720 100644
--- a/tests/post_training/unit/maxtext_engine_prepared_input_test.py
+++ b/tests/post_training/unit/maxtext_engine_prepared_input_test.py
@@ -1,4 +1,4 @@
-# Copyright 2023–2026 Google LLC
+# 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.
@@ -19,7 +19,13 @@
 kernel, same inputs, bit-identical weights and metrics. Waiting for the device
 afterwards (what Tunix's `TrainerWorker.stage_fwd_bwd` does, to report whether
 the transfer hid behind the device) must change nothing either.
+
+`maxtext_engine_prepared_input_dropless_test.py` covers the dropless replay and
+the final micro-batch's dropless program, and
+`maxtext_engine_prepared_input_sharded_test.py` a four-device mesh with Tunix's
+real model input.
 """
+# pylint: disable=protected-access
 
 import dataclasses
 import types
@@ -69,7 +75,7 @@ def _payload(seed: int) -> _Payload:
 
 def _loss_fn(model, tokens, mask, scale):
   # Depends on the data and on both weights, so any change to either shows.
-  pred = model.weights.value[0] * tokens + model.weights.value[1] * tokens**2
+  pred = model.weights[...][0] * tokens + model.weights[...][1] * tokens**2
   return (
       abstract_engine.WeightedMetric(
           unreduced_sum=jnp.sum(pred * mask) * scale,
@@ -84,9 +90,15 @@ def _model_input(payload: _Payload) -> dict[str, Any]:
   return {"tokens": payload.token_ids, "mask": payload.token_mask, "scale": 2.0}
 
 
+# Where a step's next first micro-batch is prepared, relative to the step's update.
+_BEFORE_UPDATE = "before_update"  # a separate update call follows the last fwd_bwd
+_AFTER_UPDATE = "after_update"  # the update runs inside the last fwd_bwd call (google/tunix#2719)
+
+
 class PrepareFwdBwdInputTest(absltest.TestCase):
 
   def setUp(self):
+    """A one-device mesh and a toy two-weight model, built without a checkpoint."""
     super().setUp()
     self.config = pyconfig.initialize(
         [None, get_test_config_path()],
@@ -128,92 +140,116 @@ def _engine(self, compiled: bool) -> maxtext_engine.MaxTextTrainingEngine:
       engine.compile(_payload(0))
     return engine
 
-  def _train(self, engine, payloads, prepare_at: set[int], wait: bool = False):
+  def _train(self, engine, payloads, prepare_at: set[int], wait: bool = False, order: str = _BEFORE_UPDATE):
     """Two steps of two micro-batches; micro-batches in `prepare_at` go in prepared.
 
-    A prepared micro-batch is prepared before the previous update is applied, as
-    `TrainerWorker.stage_fwd_bwd` does while that update is still on the device.
-    With `wait`, `wait_for_device` follows each prepare, as it does there too.
+    A step's next first micro-batch is prepared right after the step's last `fwd_bwd`, while that micro-batch
+    and the update are still on the device: before the update with `order=_BEFORE_UPDATE`, after it with
+    `_AFTER_UPDATE` (where Tunix's `TrainerWorker.stage_fwd_bwd` lands once the update runs inside the last
+    train call). With `wait`, `wait_for_device` follows each prepare, as it does there too.
     """
-    prepared = {i: None for i in prepare_at}
-    if 0 in prepared:
-      prepared[0] = engine.prepare_fwd_bwd_input(payloads[0])
+    prepared = {}
+
+    def prepare(i):
+      if i not in prepare_at or i >= len(payloads):
+        return
+      prepared[i] = engine.prepare_fwd_bwd_input(payloads[i])
       if wait:
-        engine.wait_for_device()
+        queued = engine._throttler._inflight_queue.qsize()
+        self.assertGreaterEqual(engine.wait_for_device(), 0.0)
+        # The in-flight queue is left for the next dispatch, or a drain, to empty as before.
+        self.assertEqual(engine._throttler._inflight_queue.qsize(), queued)
+
+    prepare(0)
     metrics = []
     for step in range(2):
       for j in range(2):
         i = 2 * step + j
-        engine.fwd_bwd(prepared[i] if i in prepared else payloads[i])
-        if j == 1 and i + 1 in prepared:
-          prepared[i + 1] = engine.prepare_fwd_bwd_input(payloads[i + 1])
-          if wait:
-            queued = engine._throttler._inflight_queue.qsize()
-            self.assertGreaterEqual(engine.wait_for_device(), 0.0)
-            # The in-flight queue is left for the update to drain as before.
-            self.assertEqual(engine._throttler._inflight_queue.qsize(), queued)
+        engine.fwd_bwd(prepared.pop(i) if i in prepared else payloads[i])
+      if order == _BEFORE_UPDATE:
+        prepare(2 * step + 2)
       engine.update()
+      if order == _AFTER_UPDATE:
+        prepare(2 * step + 2)
       buffer = engine.get_metrics(clear_cache=True)
       metrics.append({k: np.asarray(v).tolist() for k, v in sorted(buffer.scalar_metrics.items())})
       metrics[-1]["loss"] = np.asarray(buffer.weighted_metrics["loss"].compute()).tolist()
-    return np.asarray(engine.model.weights.value), metrics
+    self.assertEqual(prepared, {})
+    return np.asarray(engine.model.weights[...]), metrics
 
   def test_prepared_micro_batches_train_bit_identically(self):
     payloads = [_payload(s) for s in (1, 2, 3, 4)]
     for compiled in (True, False):
-      ref = self._engine(compiled)
-      new = self._engine(compiled)
-      with (
-          mock.patch.object(new, "_compile_for_batch", wraps=new._compile_for_batch) as new_compiles,
-          mock.patch.object(ref, "_compile_for_batch", wraps=ref._compile_for_batch) as ref_compiles,
-      ):
-        ref_w, ref_m = self._train(ref, payloads, prepare_at=set())
-        new_w, new_m = self._train(new, payloads, prepare_at={0, 2})
-      np.testing.assert_array_equal(ref_w, new_w)
-      self.assertEqual(ref_m, new_m)
-      # No kernel was rebuilt for a prepared input.
-      self.assertEqual(ref_compiles.call_count, new_compiles.call_count)
-      self.assertEqual(new.train_step, 2)
+      for order in (_BEFORE_UPDATE, _AFTER_UPDATE):
+        with self.subTest(compiled=compiled, order=order):
+          ref = self._engine(compiled)
+          new = self._engine(compiled)
+          with (
+              mock.patch.object(new, "_compile_for_batch", wraps=new._compile_for_batch) as new_compiles,
+              mock.patch.object(ref, "_compile_for_batch", wraps=ref._compile_for_batch) as ref_compiles,
+          ):
+            ref_w, ref_m = self._train(ref, payloads, prepare_at=set())
+            new_w, new_m = self._train(new, payloads, prepare_at={0, 2}, order=order)
+          np.testing.assert_array_equal(ref_w, new_w)
+          self.assertEqual(ref_m, new_m)
+          # No kernel was rebuilt for a prepared input.
+          self.assertEqual(ref_compiles.call_count, new_compiles.call_count)
+          self.assertEqual(new.train_step, 2)
 
   def test_waiting_for_the_device_changes_nothing(self):
     payloads = [_payload(s) for s in (1, 2, 3, 4)]
     for compiled in (True, False):
-      ref = self._engine(compiled)
-      new = self._engine(compiled)
-      with (
-          mock.patch.object(new, "_compile_for_batch", wraps=new._compile_for_batch) as new_compiles,
-          mock.patch.object(ref, "_compile_for_batch", wraps=ref._compile_for_batch) as ref_compiles,
-          mock.patch.object(new._metrics_logger, "write_metrics", wraps=new._metrics_logger.write_metrics) as new_w,
-          mock.patch.object(ref._metrics_logger, "write_metrics", wraps=ref._metrics_logger.write_metrics) as ref_w,
-      ):
-        ref_weights, ref_m = self._train(ref, payloads, prepare_at=set())
-        new_weights, new_m = self._train(new, payloads, prepare_at={0, 2}, wait=True)
-      np.testing.assert_array_equal(ref_weights, new_weights)
-      self.assertEqual(ref_m, new_m)
-      self.assertEqual(ref_compiles.call_count, new_compiles.call_count)
-      # The same metrics writes, in the same order (the wait writes none).
-      self.assertEqual(
-          [c.args[0].id for c in ref_w.call_args_list],
-          [c.args[0].id for c in new_w.call_args_list],
-      )
-
-  def test_wait_for_device_waits_for_the_update_and_keeps_the_queue(self):
+      for order in (_BEFORE_UPDATE, _AFTER_UPDATE):
+        with self.subTest(compiled=compiled, order=order):
+          ref = self._engine(compiled)
+          new = self._engine(compiled)
+          with (
+              mock.patch.object(new, "_compile_for_batch", wraps=new._compile_for_batch) as new_compiles,
+              mock.patch.object(ref, "_compile_for_batch", wraps=ref._compile_for_batch) as ref_compiles,
+              mock.patch.object(new._metrics_logger, "write_metrics", wraps=new._metrics_logger.write_metrics) as new_w,
+              mock.patch.object(ref._metrics_logger, "write_metrics", wraps=ref._metrics_logger.write_metrics) as ref_w,
+          ):
+            ref_weights, ref_m = self._train(ref, payloads, prepare_at=set())
+            new_weights, new_m = self._train(new, payloads, prepare_at={0, 2}, wait=True, order=order)
+          np.testing.assert_array_equal(ref_weights, new_weights)
+          self.assertEqual(ref_m, new_m)
+          self.assertEqual(ref_compiles.call_count, new_compiles.call_count)
+          # The same metrics writes, in the same order (the wait writes none).
+          self.assertEqual(
+              [c.args[0].id for c in ref_w.call_args_list],
+              [c.args[0].id for c in new_w.call_args_list],
+          )
+
+  def test_wait_for_device_waits_on_every_queued_computation_and_keeps_the_queue(self):
     engine = self._engine(True)
     engine.fwd_bwd(_payload(1))
     engine.fwd_bwd(_payload(2))
     engine.update()
     queue = engine._throttler._inflight_queue
     queued = list(queue.queue)
-    self.assertNotEmpty(queued)
-    waited = engine.wait_for_device()
+    # The update queues its gradient norm with the step's metrics, after the micro-batches' losses.
+    self.assertIsNotNone(queued[-1][1])
+    waited_on = []
+    block = jax.block_until_ready
+
+    def spy(x):
+      waited_on.append(x)
+      return block(x)
+
+    with mock.patch.object(inflight_throttler.jax, "block_until_ready", side_effect=spy):
+      waited = engine.wait_for_device()
     self.assertIsInstance(waited, float)
     self.assertGreaterEqual(waited, 0.0)
+    # Nothing dequeued: the next dispatch or drain finds the same queue.
     self.assertEqual(list(queue.queue), queued)
-    # Everything dispatched before the wait, the update included, is done.
-    for leaves, _ in queued:
-      for leaf in leaves:
-        if isinstance(leaf, jax.Array):
-          self.assertTrue(leaf.is_ready())
+    # Every queued computation was waited on, in order: each micro-batch's loss and the update's gradient norm,
+    # so a caller's wait covers the update.
+    self.assertEqual([id(x) for x in waited_on], [id(leaves) for leaves, _ in queued])
+    token = getattr(engine, "last_update_token", None)
+    if token is not None:
+      # With the readiness tokens of AI-Hypercomputer/maxtext#5571 the update's token is that same norm, so a
+      # caller that waits on it after `wait_for_device` (Tunix's `TrainerWorker.stage_fwd_bwd`) does not wait again.
+      self.assertIs(queued[-1][0][0], token)
 
   def test_the_comparison_sees_a_one_token_difference(self):
     payloads = [_payload(s) for s in (1, 2, 3, 4)]
@@ -226,36 +262,53 @@ def test_the_comparison_sees_a_one_token_difference(self):
     self.assertFalse(np.array_equal(ref_w, new_w))
 
   def test_staged_through_tunix_trainer_worker_trains_bit_identically(self):
-    # The path Tunix's pipelined step boundary takes: TrainerWorker.stage_fwd_bwd
-    # calls prepare_fwd_bwd_input, fwd_bwd_staged trains what it returned.
-    from tunix.experimental.common import datatypes  # pylint: disable=g-import-not-at-top
-    from tunix.experimental.worker import trainer_worker  # pylint: disable=g-import-not-at-top
+    """The path Tunix's pipelined step boundary takes, through a real `TrainerWorker`.
 
-    payloads = [_payload(s) for s in (1, 2, 3, 4)]
+    Each step's last micro-batch runs the update in the same call (`fwd_bwd(apply_optimizer=True)`,
+    google/tunix#2719). With staging, that call returns at the update's dispatch (`wait_for_update=False`) and the
+    next step's first micro-batch is staged right after it: `stage_fwd_bwd` calls `prepare_fwd_bwd_input` and
+    `wait_for_device`, and `fwd_bwd_staged` trains what it returned.
+    """
+    from tunix.experimental.common import datatypes  # pylint: disable=import-outside-toplevel
+    from tunix.experimental.worker import trainer_worker  # pylint: disable=import-outside-toplevel
+
+    if not hasattr(trainer_worker.TrainerWorker, "stage_fwd_bwd"):
+      self.skipTest("this Tunix has no TrainerWorker.stage_fwd_bwd (its pipelined step boundary)")
+
+    payloads = [_payload(s) for s in (1, 2, 3, 4, 5, 6)]
 
     def run(stage: bool):
       engine = self._engine(True)
       worker = trainer_worker.TrainerWorker(trainer_factory=lambda: engine)
       worker.initialize()
-      prepare = mock.patch.object(engine, "prepare_fwd_bwd_input", wraps=engine.prepare_fwd_bwd_input)
-      with prepare as spy:
-        for i, payload in enumerate(payloads):
-          request = datatypes.TrainRequest(request_id=f"r{i}", payload=payload)
-          if stage and i % 2 == 0:
-            resp = worker.stage_fwd_bwd(request)
-            # The engine's wait_for_device answered how long the device ran on.
-            self.assertGreaterEqual(resp.metadata["device_wait_s"], 0.0)
-            worker.fwd_bwd_staged(f"r{i}")
-          else:
-            worker.fwd_bwd(request)
-          if i % 2 == 1:
-            worker.update()
-      return np.asarray(engine.model.weights.value), spy.call_count
-
-    ref_w, ref_prepared = run(stage=False)
-    new_w, new_prepared = run(stage=True)
+      with (
+          mock.patch.object(engine, "prepare_fwd_bwd_input", wraps=engine.prepare_fwd_bwd_input) as prepare,
+          mock.patch.object(engine, "wait_for_device", wraps=engine.wait_for_device) as wait,
+      ):
+        staged_id = None
+        for step in range(3):
+          for j in range(2):
+            i = 2 * step + j
+            kwargs = {"apply_optimizer": True, "wait_for_update": not stage} if j == 1 else {}
+            if staged_id is not None:
+              resp = worker.fwd_bwd_staged(staged_id, **kwargs)
+              staged_id = None
+            else:
+              resp = worker.fwd_bwd(datatypes.TrainRequest(request_id=f"r{i}", payload=payloads[i]), **kwargs)
+            if j == 1:
+              self.assertEqual(resp.metadata["train_step"], step + 1)
+              if stage and i + 1 < len(payloads):
+                staged = worker.stage_fwd_bwd(datatypes.TrainRequest(request_id=f"r{i + 1}", payload=payloads[i + 1]))
+                # The engine's wait_for_device answered how long the device still ran.
+                self.assertGreaterEqual(staged.metadata["device_wait_s"], 0.0)
+                staged_id = f"r{i + 1}"
+      return np.asarray(engine.model.weights[...]), prepare.call_count, wait.call_count
+
+    ref_w, ref_prepared, ref_waits = run(stage=False)
+    new_w, new_prepared, new_waits = run(stage=True)
     np.testing.assert_array_equal(ref_w, new_w)
     self.assertEqual((ref_prepared, new_prepared), (0, 2))
+    self.assertEqual((ref_waits, new_waits), (0, 2))
 
   def test_preparing_places_the_traced_inputs_and_leaves_state_alone(self):
     engine = self._engine(True)
@@ -263,7 +316,7 @@ def test_preparing_places_the_traced_inputs_and_leaves_state_alone(self):
     before = (
         engine.micro_step_count,
         engine.train_step,
-        np.asarray(engine.model.weights.value),
+        np.asarray(engine.model.weights[...]),
     )
     prepared = engine.prepare_fwd_bwd_input(_payload(2))
     self.assertIsInstance(prepared, maxtext_engine.PreparedFwdBwdInput)
@@ -277,7 +330,7 @@ def test_preparing_places_the_traced_inputs_and_leaves_state_alone(self):
     after = (
         engine.micro_step_count,
         engine.train_step,
-        np.asarray(engine.model.weights.value),
+        np.asarray(engine.model.weights[...]),
     )
     self.assertEqual(before[:2], after[:2])
     np.testing.assert_array_equal(before[2], after[2])

From 0b107b4d3bc512824171c81a76b7d673bfb839c0 Mon Sep 17 00:00:00 2001
From: Param Bole <parambole@google.com>
Date: Tue, 13 Oct 2026 19:43:19 -0700
Subject: [PATCH 4/6] [MaxTextEngine] Test the order Tunix uses to stage before
 a replayed update

Tunix's staged step boundary can stage the next step's first micro-batch
between a step's last fwd_bwd and its update: fwd_bwd(last,
apply_optimizer=True), prepare_fwd_bwd_input(next), wait_for_device(),
update(). Under retry_when_tokens_dropped without
retry_dropless_final_micro_batch (the MLPerf recipes) the last micro-batch
can overflow; its flag is still pending at the prepare and its dropless
rerun runs inside update(), after it.

maxtext_engine_stage_before_update_test.py, on the one-layer MoE rig of the
prepared-input dropless test with the final dropless program off: the
staged order trains bit-identically to fwd_bwd(last), update(),
fwd_bwd(next payload) (weights, losses, gradient norms, replay counts) with
the last micro-batch's rerun inside update() at two boundaries and the
prepared micro-batch itself replayed at one; the flag is still pending when
the next micro-batch is prepared and wait_for_device resolves nothing; no
kernel is rebuilt. Also with the final dropless program on. Two controls are
caught: a rerun that runs on the prepared micro-batch instead of its own,
and preparing the wrong micro-batch.
---
 ...maxtext_engine_stage_before_update_test.py | 195 ++++++++++++++++++
 1 file changed, 195 insertions(+)
 create mode 100644 tests/post_training/unit/maxtext_engine_stage_before_update_test.py

diff --git a/tests/post_training/unit/maxtext_engine_stage_before_update_test.py b/tests/post_training/unit/maxtext_engine_stage_before_update_test.py
new file mode 100644
index 0000000000..0106f5a834
--- /dev/null
+++ b/tests/post_training/unit/maxtext_engine_stage_before_update_test.py
@@ -0,0 +1,195 @@
+# 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
+#
+#    https://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 order Tunix uses when it stages a step's first micro-batch before the previous step's update.
+
+Tunix's staged step boundary (`stage_before_update`) calls, at a staged boundary: `fwd_bwd(last, apply_optimizer=True)`,
+then `prepare_fwd_bwd_input(next step's first micro-batch)`, then `wait_for_device()`, then `update()`, and later
+`fwd_bwd(prepared)`. Under `retry_when_tokens_dropped` without `retry_dropless_final_micro_batch` (the MLPerf recipes
+set the first and not the second), the step's last micro-batch can overflow: its flag is still pending when the next
+micro-batch is prepared, and its dropless rerun runs inside `update()`, after the prepare. The prepared micro-batch must
+come out of that untouched, and the whole run must train bit-identically to the base order, `fwd_bwd(last,
+apply_optimizer=True)`, `update()`, `fwd_bwd(next payload)`.
+
+The rig is `maxtext_engine_prepared_input_dropless_test.py`'s (a one-layer Qwen3.5 MoE on CPU whose loss reports the
+overflow flag RoutedMoE would), with the final micro-batch's dropless program off.
+"""
+# pylint: disable=protected-access
+
+import dataclasses
+import os
+import unittest
+from unittest import mock
+
+from flax import nnx
+import jax
+import jax.numpy as jnp
+from jax.sharding import Mesh
+from maxtext.models import models
+from maxtext.training_engine import maxtext_engine
+from maxtext.utils import maxtext_utils
+import numpy as np
+import pytest
+
+from tests.post_training.unit.maxtext_engine_dropless_replay_test import _cfg, _payload, _MARKED
+from tests.post_training.unit.maxtext_engine_final_dropless_test import _program
+
+pytestmark = [pytest.mark.post_training]
+
+# Per-program loss multipliers, so the losses say which kernel ran each micro-batch.
+_SCALES = {"regular": 1.0, "replay": 2.0, "final": 3.0}
+
+# A second overflowing content, so a rerun can be told from the micro-batch prepared while it is pending.
+_MARKED_B = 60
+
+# Three steps of three micro-batches. Every step's last micro-batch overflows. At both staged boundaries (after steps
+# 0 and 1) its rerun runs inside `update()` with the next step's first micro-batch already prepared; after step 2 it
+# runs inside the last `update()`. Step 1's first (prepared) micro-batch overflows too.
+_STEPS = ((10, 20, _MARKED), (_MARKED_B, 20, _MARKED), (40, 20, _MARKED))
+
+
+class StageBeforeUpdateOrderTest(unittest.TestCase):
+
+  def setUp(self):
+    os.environ["NEW_MODEL_DESIGN"] = "1"
+    os.environ["SKIP_JAX_PRECOMPILE"] = "1"
+
+  def _engine(self, final_dropless: bool = False):
+    cfg = _cfg(override_logical_axis_rules=True, ragged_buffer_factor=2.5)
+    mesh = Mesh(maxtext_utils.create_device_mesh(cfg), cfg.mesh_axes)
+    real_model = models.Transformer(config=cfg, mesh=mesh, quant=None, model_mode="train", rngs=nnx.Rngs(0))
+    with mock.patch.object(maxtext_engine.model_creation_utils, "from_pretrained", return_value=real_model):
+      engine = maxtext_engine.MaxTextTrainingEngine(cfg, mesh=mesh)
+    # The ring-of-experts prerequisites of retry_when_tokens_dropped cannot run on CPU, so the switches are set
+    # directly, as the dropless replay and final-dropless tests do.
+    engine._dropless_replay = True
+    engine._final_dropless = final_dropless
+    base_loss = maxtext_engine.make_router_replay_loss_fn(cfg)
+
+    def loss_fn(model, **batch):
+      program = _program(model)
+      loss, aux = base_loss(model, **batch)
+      aux = {k: v for k, v in aux.items() if k != "has_moe_overflow"}
+      marked = jnp.logical_or(batch["inputs"][0, 0] == _MARKED, batch["inputs"][0, 0] == _MARKED_B)
+      aux["has_moe_overflow"] = jnp.logical_and(marked, program == "regular")
+      aux["xent_sum"] = aux["xent_sum"] * _SCALES[program]
+      return loss, aux
+
+    engine.with_gen_model_input_fn(maxtext_engine.router_replay_gen_model_input_fn)
+    engine.with_loss_fn(loss_fn, has_aux=True)
+    engine.compile(None)
+    return engine
+
+  def _train(self, engine, staged: bool, prepare=None):
+    """Runs `_STEPS`; with `staged`, each next step's first micro-batch is prepared before the update.
+
+    Returns the final parameters, per-step metrics, and per staged boundary whether the last micro-batch's flag was
+    still pending when the next one was prepared.
+    """
+    prepare = prepare or engine.prepare_fwd_bwd_input
+    prepared = None
+    pending_at_prepare = []
+    metrics = []
+    for step, tokens in enumerate(_STEPS):
+      for j, first_token in enumerate(tokens):
+        final = j == len(tokens) - 1
+        if j == 0 and prepared is not None:
+          self.assertIsInstance(prepared, maxtext_engine.PreparedFwdBwdInput)
+          engine.fwd_bwd(prepared, apply_optimizer=False)
+          prepared = None
+        else:
+          engine.fwd_bwd(_payload(first_token), apply_optimizer=final)
+      if staged and step + 1 < len(_STEPS):
+        pending_at_prepare.append(engine._pending_overflow_check is not None)
+        prepared = prepare(_payload(_STEPS[step + 1][0]))
+        self.assertGreaterEqual(engine.wait_for_device(), 0.0)
+        # Waiting resolves nothing: the flag is read, and the rerun run, by `update()`.
+        self.assertIsNotNone(engine._pending_overflow_check)
+      with mock.patch.object(engine, "record_metrics", wraps=engine.record_metrics) as record:
+        engine.update()
+      recorded = {call.args[0]: call.args[1] for call in record.call_args_list}
+      buffer = engine.get_metrics(clear_cache=True)
+      metrics.append(
+          {
+              "replays": int(recorded["moe_dropless_replays"]),
+              "loss": np.asarray(buffer.weighted_metrics["loss"].compute()).tolist(),
+              "gradient_norm": np.asarray(buffer.scalar_metrics["gradient_norm"]).tolist(),
+          }
+      )
+    self.assertIsNone(prepared)
+    params = [np.asarray(leaf) for leaf in jax.tree.leaves(nnx.state(engine.model, nnx.Param))]
+    return params, metrics, pending_at_prepare
+
+  def _assert_same_training(self, ref, new):
+    self.assertEqual(ref[1], new[1])
+    self.assertEqual(len(ref[0]), len(new[0]))
+    for r, n in zip(ref[0], new[0]):
+      np.testing.assert_array_equal(r, n)
+
+  def test_staged_before_a_replayed_update_trains_bit_identically(self):
+    reference = self._train(self._engine(), staged=False)
+    # Not vacuous: every step's last micro-batch replays (in update), and step 1's first (the prepared one) replays.
+    self.assertEqual([m["replays"] for m in reference[1]], [1, 2, 1])
+    engine = self._engine()
+    with mock.patch.object(engine, "_compile_for_batch", wraps=engine._compile_for_batch) as compiles:
+      staged = self._train(engine, staged=True)
+    self._assert_same_training(reference, staged)
+    # The last micro-batch's flag was still pending at both prepares, i.e. its rerun ran after the prepare.
+    self.assertEqual(staged[2], [True, True])
+    self.assertEqual(compiles.call_count, 1)  # no kernel was rebuilt for the prepared input
+
+  def test_with_the_final_dropless_program_too(self):
+    # retry_dropless_final_micro_batch (#5673): the last micro-batch runs the final program and cannot overflow, but
+    # its check is still created and read by update(), so the order is the same; only step 1's prepared first
+    # micro-batch replays.
+    reference = self._train(self._engine(final_dropless=True), staged=False)
+    self.assertEqual([m["replays"] for m in reference[1]], [0, 1, 0])
+    staged = self._train(self._engine(final_dropless=True), staged=True)
+    self._assert_same_training(reference, staged)
+    self.assertEqual(staged[2], [True, True])
+
+  def test_the_comparison_sees_a_rerun_of_the_wrong_micro_batch(self):
+    # Control: a rerun inside update() that ran on the prepared micro-batch instead of its own is caught.
+    reference = self._train(self._engine(), staged=False)
+    engine = self._engine()
+    holder = {}
+    real_prepare = engine.prepare_fwd_bwd_input
+    real_resolve = engine._resolve_overflow_check
+
+    def prepare(payload):
+      holder["prepared"] = real_prepare(payload)
+      return holder["prepared"]
+
+    def resolve():
+      check = engine._pending_overflow_check
+      prepared = holder.pop("prepared", None)
+      if check is not None and prepared is not None and check.micro_step > 0:
+        dynamic, _ = maxtext_engine._split_static_and_dynamic(prepared.batch)
+        engine._pending_overflow_check = dataclasses.replace(check, inputs=dynamic)
+      return real_resolve()
+
+    with mock.patch.object(engine, "_resolve_overflow_check", side_effect=resolve):
+      wrong = self._train(engine, staged=True, prepare=prepare)
+    self.assertNotEqual(reference[1], wrong[1])
+
+  def test_the_comparison_sees_a_different_prepared_micro_batch(self):
+    # Control: preparing the wrong micro-batch (the next step's second instead of its first) is caught.
+    reference = self._train(self._engine(), staged=False)
+    engine = self._engine()
+    real_prepare = engine.prepare_fwd_bwd_input
+    wrong = self._train(engine, staged=True, prepare=lambda payload: real_prepare(_payload(20)))
+    self.assertNotEqual(reference[1], wrong[1])
+
+
+if __name__ == "__main__":
+  unittest.main()

From c6b304a30205aabd1868df373ad48f79dd56ccfb Mon Sep 17 00:00:00 2001
From: Param Bole <parambole@google.com>
Date: Tue, 13 Oct 2026 20:12:47 -0700
Subject: [PATCH 5/6] [MaxTextEngine] Drive the staged-before-update order
 through Tunix's TrainerWorker

The same order as the engine-level test, through a real TrainerWorker:
fwd_bwd(last, apply_optimizer=True, defer_update=True), stage_fwd_bwd(next),
apply_deferred_update(), fwd_bwd_staged(next), against the fused calls. It
trains bit-identically (weights, losses, gradient norms) with the last
micro-batch's rerun inside the deferred update, defer_update never reaches
the engine, and the stage reports no device wait (the update call does).
Skipped when the installed Tunix has no apply_deferred_update.
---
 ...maxtext_engine_stage_before_update_test.py | 65 +++++++++++++++++++
 1 file changed, 65 insertions(+)

diff --git a/tests/post_training/unit/maxtext_engine_stage_before_update_test.py b/tests/post_training/unit/maxtext_engine_stage_before_update_test.py
index 0106f5a834..9ea8bfc1e6 100644
--- a/tests/post_training/unit/maxtext_engine_stage_before_update_test.py
+++ b/tests/post_training/unit/maxtext_engine_stage_before_update_test.py
@@ -65,6 +65,7 @@ def setUp(self):
     os.environ["SKIP_JAX_PRECOMPILE"] = "1"
 
   def _engine(self, final_dropless: bool = False):
+    """The one-layer MoE engine with the dropless replay on; its loss raises the overflow flag on marked batches."""
     cfg = _cfg(override_logical_axis_rules=True, ragged_buffer_factor=2.5)
     mesh = Mesh(maxtext_utils.create_device_mesh(cfg), cfg.mesh_axes)
     real_model = models.Transformer(config=cfg, mesh=mesh, quant=None, model_mode="train", rngs=nnx.Rngs(0))
@@ -158,6 +159,70 @@ def test_with_the_final_dropless_program_too(self):
     self._assert_same_training(reference, staged)
     self.assertEqual(staged[2], [True, True])
 
+  def test_through_tunix_trainer_worker_the_deferred_update_trains_bit_identically(self):
+    """The same order, driven by a real Tunix `TrainerWorker`: `fwd_bwd(last, apply_optimizer=True,
+    defer_update=True)` runs the last micro-batch and holds the update, `stage_fwd_bwd` prepares the next step's first
+    micro-batch without waiting, `apply_deferred_update` waits for the device and runs the update (and with it the
+    last micro-batch's rerun), and `fwd_bwd_staged` trains the prepared micro-batch. Against the fused calls
+    `fwd_bwd(last, apply_optimizer=True)` and `fwd_bwd(next payload)`."""
+    from tunix.experimental.common import datatypes  # pylint: disable=import-outside-toplevel
+    from tunix.experimental.worker import trainer_worker  # pylint: disable=import-outside-toplevel
+
+    if not hasattr(trainer_worker.TrainerWorker, "apply_deferred_update"):
+      self.skipTest("this Tunix has no TrainerWorker.apply_deferred_update (staging before the update)")
+
+    def run(staged: bool):
+      engine = self._engine()
+      worker = trainer_worker.TrainerWorker(trainer_factory=lambda: engine)
+      worker.initialize()
+      metrics, device_waits = [], []
+      staged_id = None
+      with mock.patch.object(engine, "fwd_bwd", wraps=engine.fwd_bwd) as fwd_bwd:
+        for step, tokens in enumerate(_STEPS):
+          for j, first_token in enumerate(tokens):
+            final = j == len(tokens) - 1
+            kwargs = {"apply_optimizer": True} if final else {}
+            if final and staged and step + 1 < len(_STEPS):
+              kwargs["defer_update"] = True
+            if j == 0 and staged_id is not None:
+              resp = worker.fwd_bwd_staged(staged_id, **kwargs)
+              staged_id = None
+            else:
+              request = datatypes.TrainRequest(request_id=f"s{step}m{j}", payload=_payload(first_token))
+              resp = worker.fwd_bwd(request, **kwargs)
+          if resp.metadata.get("update_deferred"):
+            nxt = datatypes.TrainRequest(request_id=f"s{step + 1}m0", payload=_payload(_STEPS[step + 1][0]))
+            stage = worker.stage_fwd_bwd(nxt)
+            self.assertTrue(stage.metadata["before_update"])
+            self.assertIsNone(stage.metadata["device_wait_s"])
+            update = worker.apply_deferred_update()
+            device_waits.append(update.metadata["device_wait_s"])
+            self.assertEqual(update.metadata["train_step"], step + 1)
+            staged_id = nxt.request_id
+          else:
+            self.assertEqual(resp.metadata["train_step"], step + 1)
+          buffer = engine.get_metrics(clear_cache=True)
+          metrics.append(
+              {
+                  "loss": np.asarray(buffer.weighted_metrics["loss"].compute()).tolist(),
+                  "gradient_norm": np.asarray(buffer.scalar_metrics["gradient_norm"]).tolist(),
+              }
+          )
+        # `defer_update` is the worker's, never the engine's.
+        self.assertTrue(all("defer_update" not in c.kwargs for c in fwd_bwd.call_args_list))
+      self.assertIsNone(staged_id)
+      self.assertIsNone(worker._deferred_update)
+      params = [np.asarray(leaf) for leaf in jax.tree.leaves(nnx.state(engine.model, nnx.Param))]
+      return params, metrics, device_waits
+
+    reference = run(staged=False)
+    staged = run(staged=True)
+    self.assertEqual(reference[1], staged[1])
+    for r, n in zip(reference[0], staged[0]):
+      np.testing.assert_array_equal(r, n)
+    self.assertEqual(len(staged[2]), len(_STEPS) - 1)
+    self.assertTrue(all(w >= 0.0 for w in staged[2]))
+
   def test_the_comparison_sees_a_rerun_of_the_wrong_micro_batch(self):
     # Control: a rerun inside update() that ran on the prepared micro-batch instead of its own is caught.
     reference = self._train(self._engine(), staged=False)

From c79bcfd162d3c93c08fcdcc7fc8364596a319f5c Mon Sep 17 00:00:00 2001
From: Param Bole <parambole@google.com>
Date: Tue, 13 Oct 2026 23:54:08 -0700
Subject: [PATCH 6/6] [MaxTextEngine] Document that staging leaves a pending
 overflow check alone

Docstrings and README only, no code change.

Tunix stages a step's first micro-batch between the previous step's last
fwd_bwd and its update() when it defers the update. Under
retry_when_tokens_dropped the last micro-batch's overflow flag is still
pending at that point. State the contract this order relies on, which
maxtext_engine_stage_before_update_test.py pins:

- prepare_fwd_bwd_input neither reads nor resolves the pending check. The
  dropless rerun that update() may run reads the inputs that check kept,
  never the prepared ones. The mapping reads only the payload, so the
  prepared input is the same whether the update runs before or after it.
- wait_for_device leaves a pending overflow check, like the in-flight queue
  and its metrics, to the next fwd_bwd or update().
- README: the same note under "One optimizer step", and the test file in the
  tests table.
- Two test docstrings no longer say the MLPerf recipes set
  retry_dropless_final_micro_batch, or disagree about it: the team's trainer
  configuration sets retry_when_tokens_dropped, and the final dropless
  program is an option on top of it.
---
 src/maxtext/training_engine/README.md         |  7 ++++++
 src/maxtext/training_engine/maxtext_engine.py | 23 +++++++++++++------
 ...ext_engine_prepared_input_dropless_test.py | 10 ++++----
 ...maxtext_engine_stage_before_update_test.py | 10 ++++----
 4 files changed, 33 insertions(+), 17 deletions(-)

diff --git a/src/maxtext/training_engine/README.md b/src/maxtext/training_engine/README.md
index 5e08df0117..58b2797dff 100644
--- a/src/maxtext/training_engine/README.md
+++ b/src/maxtext/training_engine/README.md
@@ -299,6 +299,12 @@ and does not use the engine.
 | `prepare_fwd_bwd_input(micro_batch)` | none; `device_put` | micro-batch | a `PreparedFwdBwdInput`: the micro-batch mapped to loss inputs and placed on `fwd_bwd`'s input shardings, which `fwd_bwd` takes in its place |
 | `wait_for_device()` | none | | the seconds it blocked until the in-flight `fwd_bwd`s and `update` finished; dequeues nothing and writes no metrics |
 
+A caller may prepare the next step's first micro-batch between a step's last `fwd_bwd` and its
+`update()`, as Tunix does when it stages before a deferred update. Under `retry_when_tokens_dropped`
+the last micro-batch's overflow flag is then still pending: neither `prepare_fwd_bwd_input` nor
+`wait_for_device` resolves it, `update()` does, and a dropless rerun reads that micro-batch's own kept
+inputs, so the prepared micro-batch is unaffected (`maxtext_engine_stage_before_update_test.py`).
+
 Every micro-batch runs the same `fwd_bwd` program. The first micro-batch's gradients start the
 accumulator, and `accumulate` adds each later one's to it in place. The accumulator is not passed to
 `fwd_bwd`: given it, XLA can fold the add into the backward pass, and on large meshes folding it into
@@ -386,6 +392,7 @@ The engine's tests are in `tests/post_training/unit/` and run on CPU; they requi
 | `maxtext_engine_checkpoint_test.py` | resuming from a mid-step checkpoint through Orbax reproduces the uninterrupted step |
 | `maxtext_engine_grad_cast_test.py` | `cast_grads_after_all_reduce`: which gradients leave `fwd_bwd` uncast, and that moving their cast changes no number on CPU |
 | `maxtext_engine_prepared_input{,_dropless,_sharded}_test.py` | `prepare_fwd_bwd_input` and `wait_for_device`: a prepared micro-batch trains bit-identically on one device, through Tunix's `TrainerWorker` staging, under the dropless replay and final dropless programs, and on a four-device mesh with Tunix's payload; the wait leaves the in-flight queue and the metrics writes as they were |
+| `maxtext_engine_stage_before_update_test.py` | the order Tunix uses to stage before a deferred update (`fwd_bwd(last)`, `prepare_fwd_bwd_input(next)`, `wait_for_device()`, `update()`) under the dropless replay and final dropless programs, directly and through Tunix's `TrainerWorker`: bit-identical to the base order, including a rerun of the last micro-batch inside `update()` |
 | `maxtext_engine_{constructor,data_parallel,xaot,packing,profiling}_test.py`, `router_replay_engine_test.py` | construction, data parallelism, ahead-of-time compilation, sequence packing, profiling, router replay |
 
 ## Differences from pre_train/train.py
diff --git a/src/maxtext/training_engine/maxtext_engine.py b/src/maxtext/training_engine/maxtext_engine.py
index be10518f87..0f1d0ddc1b 100644
--- a/src/maxtext/training_engine/maxtext_engine.py
+++ b/src/maxtext/training_engine/maxtext_engine.py
@@ -2419,12 +2419,19 @@ def prepare_fwd_bwd_input(self, payload: abstract_engine.TrainerPayload) -> Prep
     """Maps a micro-batch to loss inputs and places them on the mesh, ahead of its `fwd_bwd`.
 
     For a caller that receives the micro-batch early (Tunix's `TrainerWorker.stage_fwd_bwd`,
-    while the device still computes the previous step's last micro-batch and update): the
-    input mapping and the host-to-device copy run here, and `fwd_bwd(prepared)` later
-    dispatches without moving data. The traced inputs go to exactly the shardings `fwd_bwd`
-    compiles them onto (`_batch_data_shardings`), with the closed-over half untouched, so the
-    kernel, its inputs and the result are those of `fwd_bwd(payload)`. Holds one micro-batch
-    of inputs in device memory until that `fwd_bwd`.
+    while the device still computes the previous step's last micro-batch, and its update unless
+    the caller deferred it): the input mapping and the host-to-device copy run here, and
+    `fwd_bwd(prepared)` later dispatches without moving data. The traced inputs go to exactly
+    the shardings `fwd_bwd` compiles them onto (`_batch_data_shardings`), with the closed-over
+    half untouched, so the kernel, its inputs and the result are those of `fwd_bwd(payload)`.
+    Holds one micro-batch of inputs in device memory until that `fwd_bwd`.
+
+    Safe between a step's last `fwd_bwd` and its `update()`, the order Tunix uses to stage
+    before a deferred update. Under `retry_when_tokens_dropped` the last micro-batch's
+    overflow flag is still pending there: this neither reads nor resolves it, and the dropless
+    rerun `update()` may then run reads the inputs that micro-batch's pending check kept, never
+    the prepared ones. The mapping reads only the payload (no parameter, gradient or optimizer
+    state), so the prepared input is the same whether the update runs before or after it.
 
     Args:
       payload: Packed micro-batch training input.
@@ -2449,7 +2456,9 @@ def wait_for_device(self) -> float:
     waited on when it left the queue. A blocking wait, not an `is_ready()` poll: under Pathways
     a readiness check is an RPC to the proxy, so a poll can read "not ready" for work that has
     finished. Changes no state: the in-flight queue and its pending metrics are left as they
-    were (see `InflightThrottler.block_until_ready`).
+    were (see `InflightThrottler.block_until_ready`), and so is a pending overflow check: under
+    `retry_when_tokens_dropped` the last micro-batch's flag, and its dropless rerun if it
+    overflowed, are still left to the next `fwd_bwd` or `update()`.
     """
     start = time.monotonic()
     self._throttler.block_until_ready()
diff --git a/tests/post_training/unit/maxtext_engine_prepared_input_dropless_test.py b/tests/post_training/unit/maxtext_engine_prepared_input_dropless_test.py
index f9afa5547e..9db43e0f33 100644
--- a/tests/post_training/unit/maxtext_engine_prepared_input_dropless_test.py
+++ b/tests/post_training/unit/maxtext_engine_prepared_input_dropless_test.py
@@ -13,11 +13,11 @@
 # limitations under the License.
 """prepare_fwd_bwd_input under the dropless replay and the final micro-batch's dropless program.
 
-The MLPerf recipes train with `retry_when_tokens_dropped` (a micro-batch whose MoE ragged buffer overflowed is
-accumulated as zero and rerun dropless, from the inputs it ran on) and `retry_dropless_final_micro_batch` (a step's
-last micro-batch runs the dropless program directly). A micro-batch prepared ahead must take both paths exactly like
-the payload it came from: the overflowing micro-batch is replayed from its prepared inputs, the final one runs the
-final program, and the weights, losses and replay counts are bit-identical.
+The team's MLPerf trainer configuration trains with `retry_when_tokens_dropped` (a micro-batch whose MoE ragged buffer
+overflowed is accumulated as zero and rerun dropless, from the inputs it ran on); `retry_dropless_final_micro_batch` (a
+step's last micro-batch runs the dropless program directly) is an option on top of it. A micro-batch prepared ahead must
+take both paths exactly like the payload it came from: the overflowing micro-batch is replayed from its prepared inputs,
+the final one runs the final program, and the weights, losses and replay counts are bit-identical.
 
 The rig is `maxtext_engine_final_dropless_test.py`'s: a one-layer Qwen3.5 MoE on CPU whose loss reports the overflow
 flag RoutedMoE would, from which program its graph is, and scales the loss per program so the kernels can be told
diff --git a/tests/post_training/unit/maxtext_engine_stage_before_update_test.py b/tests/post_training/unit/maxtext_engine_stage_before_update_test.py
index 9ea8bfc1e6..fa9090c82a 100644
--- a/tests/post_training/unit/maxtext_engine_stage_before_update_test.py
+++ b/tests/post_training/unit/maxtext_engine_stage_before_update_test.py
@@ -15,11 +15,11 @@
 
 Tunix's staged step boundary (`stage_before_update`) calls, at a staged boundary: `fwd_bwd(last, apply_optimizer=True)`,
 then `prepare_fwd_bwd_input(next step's first micro-batch)`, then `wait_for_device()`, then `update()`, and later
-`fwd_bwd(prepared)`. Under `retry_when_tokens_dropped` without `retry_dropless_final_micro_batch` (the MLPerf recipes
-set the first and not the second), the step's last micro-batch can overflow: its flag is still pending when the next
-micro-batch is prepared, and its dropless rerun runs inside `update()`, after the prepare. The prepared micro-batch must
-come out of that untouched, and the whole run must train bit-identically to the base order, `fwd_bwd(last,
-apply_optimizer=True)`, `update()`, `fwd_bwd(next payload)`.
+`fwd_bwd(prepared)`. Under `retry_when_tokens_dropped` without `retry_dropless_final_micro_batch` (the team's MLPerf
+trainer configuration sets the first; the second is opt-in), the step's last micro-batch can overflow: its flag is still
+pending when the next micro-batch is prepared, and its dropless rerun runs inside `update()`, after the prepare. The
+prepared micro-batch must come out of that untouched, and the whole run must train bit-identically to the base order,
+`fwd_bwd(last, apply_optimizer=True)`, `update()`, `fwd_bwd(next payload)`.
 
 The rig is `maxtext_engine_prepared_input_dropless_test.py`'s (a one-layer Qwen3.5 MoE on CPU whose loss reports the
 overflow flag RoutedMoE would), with the final micro-batch's dropless program off.