feat(training_engine): add prepare_fwd_bwd_input and wait_for_device to overlap the next step's input copy (-7.2% s/step, 397B RL trainer bench, 128 v7x)

src/maxtext/training_engine/README.md

@@ -296,6 +296,14 @@ 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 |
+
+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
@@ -383,6 +391,8 @@ 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_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

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

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,66 @@ 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 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.
+
+    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)
+
+  def wait_for_device(self) -> float:
+    """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. 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`), 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()
+    return time.monotonic() - start
+
   @_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:

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