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)

[parambole]

TL;DR

• 7.2% less time per optimizer step at 128 chips: 105.97 -> 98.31 s (2,783 -> 3,000 tok/s/chip); 21.5% less at 256 chips: 71.47 -> 56.12 s. Trainer only: the RL trainer bench (Add an RL trainer bench under experimental/rl_trainer_bench #5701) with 397B v7x RL trainer: tuned settings, sharded log-prob head, staged step boundary (130.8 -> 98.3 s per optimizer step, trainer bench) google/tunix#2878's trainer settings and log-prob head as measured (the final stack, see Terms), staging off vs on in one TPU job, the calls made by a bench-only driver switch (an earlier form of [DO NOT MERGE] perf(rl_trainer_bench): opt-in STAGE_BEFORE_UPDATE step boundary (-7.2% step time on the 397B bench; stacked on #5701) #5719).
• What this PR adds: two engine methods, used only when called (no new config key). prepare_fwd_bwd_input(payload) turns the next micro-batch into fwd_bwd's inputs and starts copying them to the chips; fwd_bwd(prepared) then skips that work. wait_for_device() waits until the engine's queued work has finished and returns the seconds it waited.
• Why it helps: on the bench the trainer chips sat idle at every step boundary while the next step's first micro-batch (about 5 GB) was copied in; 397B v7x RL trainer: tuned settings, sharded log-prob head, staged step boundary (130.8 -> 98.3 s per optimizer step, trainer bench) google/tunix#2878 now starts that copy while the step's last micro-batch still computes. On the reduced rig (Tunix's own orchestrator, 3 micro-batches per step) boundary idle fell from 11.7-12.5 s to 0.02-0.03 s.
• Validated so far: CPU tests train bit-identically with prepared inputs, including a dropless rerun inside update(). On TPU only order and data were checked (a bench pair and the reduced rig), not loss or gradients. Not run yet: a TPU A/B with real gradients, and any end-to-end RL run.
• To run: 397B v7x RL trainer: tuned settings, sharded log-prob head, staged step boundary (130.8 -> 98.3 s per optimizer step, trainer bench) google/tunix#2878 already pins this PR's head: build its image and launch mlperf_gbs1024_submission.sh start with the image pinned by digest (staging is on by default; STAGE_NEXT_MICROBATCH=false turns it off). After this PR merges, re-pin 397B v7x RL trainer: tuned settings, sharded log-prob head, staged step boundary (130.8 -> 98.3 s per optimizer step, trainer bench) google/tunix#2878 to the merge commit. Draft until the two TPU runs under "Not yet tested" have run and passed.

Terms used below

Term	Meaning
Staging	Preparing the next step's first micro-batch and starting its copy to the chips before this step's update(), so the copy overlaps compute.
Optimizer step	One weight update: 9 micro-batches then one update() at 128 chips (5 at 256 chips, 3 on the reduced rig); 1,152 trajectories (72 prompts x 16 generations) on the bench. Bench numbers are means of steps 1-4; step 0 is warm-up.
tok/s/chip	Micro-batches x rows x 65,536 token slots / chips / seconds per optimizer step (9 x 64 at 128 chips, 5 x 128 at 256). About 9.6% of the slots at 128 chips are padding.
Bench	#5701's RL trainer bench: the trainer half of RL only (no rollout, sandbox or reward), random weights, synthetic sequences of the real shape.
Team cell	The team's trainer settings from before google/tunix#2878, at the bench's layout: ragged_buffer_factor 2.75 with dropless replay, context on device, log-prob chunk 1024. Not the recipe's own earlier default (fully dropless, chunk 512), which was not measured here.
Final stack	google/tunix#2878's trainer settings as measured: 1 libtpu flag and 7 MaxText keys, 2 precision keys (gdn_cp_matmul_precision=high, cast_grads_after_all_reduce=true) and the sharded log-prob head (measured as google/tunix#2873's file). One of the 7, moe_route=device, had no effect on the bench and is left out of the shipped recipe.
Reduced rig	A test harness, not part of this PR: one MaxText trainer on 128 chips driven over gRPC by Tunix's real orchestrator, replaying recorded rollouts; 3 micro-batches per step, learning rate 0.
Arm, same pod	An arm is one side of an A/B (staging off or on). Same pod: the arms ran back to back in one TPU job.
Dropless replay	retry_when_tokens_dropped: a micro-batch that overflows its MoE expert buffers is rerun. Each fwd_bwd reads the previous micro-batch's overflow flag on the host, and update() reads the last one.
Pipelined boundary	Tunix's --pipeline_step_boundary: it sends the next step's first micro-batch without waiting for this step's metrics and logging. On its own it copies nothing to the chips early.
Boundary idle	Trainer chip time with no work between a step's update and the next step's first micro-batch.

What changes

Piece	File	What it does
prepare_fwd_bwd_input(payload) (new)	src/maxtext/training_engine/maxtext_engine.py	Runs the input mapping fwd_bwd runs, then jax.device_puts the traced inputs onto the shardings fwd_bwd is compiled with. Returns a PreparedFwdBwdInput without waiting for the copy.
wait_for_device() (new)	same	Blocks until everything in the in-flight queue has finished. Returns the seconds it waited.
PreparedFwdBwdInput (new)	same	Frozen dataclass with one field, batch.
fwd_bwd (signature and one line)	same	Takes a payload or a PreparedFwdBwdInput, and skips the mapping for the latter. Nothing after the mapping changes.
InflightThrottler.block_until_ready() (new)	src/maxtext/training_engine/inflight_throttler.py	Waits on a snapshot of the queue. Dequeues nothing and writes no metrics.
README	src/maxtext/training_engine/README.md	The two methods, the call-order contract and the new tests.
Tests	4 new files in tests/post_training/unit/	18 tests, plus 3 that run in a 4-device child process.

Size: code +74/-3, README +10, tests +1,090. The methods are not on AbstractTrainingEngine, so callers look them up (Tunix uses getattr). prepare_fwd_bwd_input is safe between a step's last fwd_bwd and its update(), where Tunix stages; contracts and a call-order example are under API contracts below.

How it works

The trainer chips at one step boundary

gantt
    title One step boundary on the bench, 128 chips, seconds
    dateFormat X
    axisFormat %s
    todayMarker off
    section Staging off
    Last micro-batch of step k :b1, 0, 11
    Update :b2, 11, 13
    Idle - micro-batch 1 copied in :crit, b3, 13, 21
    Micro-batch 1 of step k+1 :b4, 21, 32
    section Staging on
    Last micro-batch of step k :a1, 0, 11
    prepare_fwd_bwd_input copies micro-batch 1 :a2, 0, 8
    Update :a3, 11, 13
    Micro-batch 1 of step k+1 :a4, 13, 24

Loading

Schematic in whole seconds, bench at 128 chips. Measured: the 7.6-7.7 s idle fell to about 0 s, and regular micro-batches stayed at 10.75 s. The copy bar is drawn as long as the gap it removes; the bench does not time the copy itself.

The calls at one staged step boundary

flowchart TB
  subgraph S1["Call 1 - last micro-batch of step k"]
    direction LR
    A1["Tunix fwd_bwd<br/>with the update deferred"] --> C1["MaxText fwd_bwd<br/>dispatches the micro-batch"]
  end
  subgraph S2["Call 2 - stage micro-batch 1 of step k+1"]
    direction LR
    A2["Tunix stage_fwd_bwd"] --> C2["NEW prepare_fwd_bwd_input<br/>starts the copy and returns"]
  end
  subgraph S3["Call 3 - deferred update"]
    direction LR
    A3["Tunix apply_deferred_update"] --> C3["NEW wait_for_device<br/>waits for the last micro-batch"] --> D3["MaxText update<br/>reruns it if it overflowed"]
  end
  subgraph S4["Call 4 - step k+1 begins"]
    direction LR
    A4["Tunix metrics, checkpoint<br/>and weight sync, then fwd_bwd_staged"] --> C4["MaxText fwd_bwd<br/>on the prepared input, no copy"]
  end
  S1 --> S2 --> S3 --> S4

Loading

Tunix's call order once staging is on, one row per call; each Tunix call reaches MaxText through its TrainerWorker. NEW marks this PR's methods.

Why the stage goes before update(). With dropless replay (the layout measured here), update() first reads the last micro-batch's overflow flag on the host, so it returns only after that micro-batch has finished. A stage placed after update() can hide only behind the update itself; on the reduced rig (update about 1.6 s) that left 7.43 s idle per boundary. Placed before it, the copy hides behind the whole last micro-batch (13.4 s on the rig, 10.75 s on the bench), and the rig's idle fell to 0.02 s.

Results

RL trainer bench (#5701): staging off vs on, same pod

Mesh fsdp 32 x context 4 x expert 2 at 128 chips (fsdp 64 at 256), ragged_buffer_factor 2.75 with dropless replay, log-prob chunk 1024. Staging ran through an earlier form of #5719's bench-driver switch, which calls prepare_fwd_bwd_input and wait_for_device directly; the two engine methods and the fwd_bwd change were this PR's, and the rest of the engine was the bench image's.

Seconds per optimizer step:

Setup	Off	On	Change
128 chips, final stack (arms ran on, off, on)	105.97	98.31; repeat 98.26	-7.66 s (-7.2%)
256 chips, final stack (off, on)	71.47	56.12	-15.36 s (-21.5%)
128 chips, final stack without the 2 precision keys (off, on)	107.73	100.01	-7.72 s (-7.2%)

Changes are computed from unrounded means (71.474 and 56.115 s at 256 chips), so they can differ from the rounded columns by 0.01 s.

tok/s/chip, same runs:

Setup	Off	On
128 chips, final stack	2,783	3,000
256 chips, final stack	2,292	2,920
128 chips, without the 2 precision keys	2,738	2,949

• At 256 chips the packing has 640 row slots (5 x 128) for the same 560 real rows; at 128 chips it has 576 (9 x 64). So compare the two chip counts in seconds, not in tok/s/chip.
• Where the time went: without staging, each step's 2nd fwd_bwd call took 7.6-7.7 s longer than a regular call at 128 chips (15.1-15.5 s at 256). That call waits for micro-batch 1, including its host-to-device copy. With staging the extra time was gone (-0.04 to -0.05 s at 128, -0.03 to -0.04 s at 256). Regular micro-batches did not change (medians, off vs on: 10.7545 vs 10.754 s at 128 chips, 11.0765 vs 11.075 s at 256).
• At 256 chips the removed wait (15.1-15.5 s) was longer than the roughly 12 s it had to hide behind (the last micro-batch, about 11.1 s, plus the update), yet it hid completely. The prediction written before the run expected partial hiding (-11.7 s) and missed on the good side. Why it hid fully is not established.
• In the final-stack staged arm at 128 chips, prepare_fwd_bwd_input returned in 0.007-0.041 s and wait_for_device then waited 10.49-10.72 s: the last micro-batch was still computing while the copy ran. The repeat arm's wait reached 10.725 s. Without the precision keys: 0.009-0.042 s and 10.69-10.91 s. At 256 chips: 0.013-0.095 s and 10.80-11.05 s.
• Every arm: 0 dropped tokens, 0 dropless replays.
• The bench drains the in-flight queue and reads metrics on the host at every boundary, so its boundary is not the RL trainer's: -7.66 s is a trainer-bench result, not an end-to-end RL gain.

Reduced rig with Tunix's real orchestrator

Tunix's StandardRLProgram, through a subclass that only adds observation hooks, and TrainerWorker, at an earlier revision of the boundary code (26b06c504), with recorded trajectory groups from the first optimizer step of a 397B RL run. No weight sync or checkpoint save. The trainer ran the team cell (ragged_buffer_factor 2.75 with dropless replay; about 13.4 s per micro-batch). Each scored arm ran 5 steps, all arms in one pod.

Arm	Boundary idle, median (s)	Seconds per optimizer step, median
Flags off (two runs)	12.46 / 11.72	64.45 / 61.85
Pipelined boundary, not staged	11.34	62.26
Staged after the update (--no-stage_before_update)	7.43	59.23
Staged before the update (the default once staging is on)	0.02	50.82

• Per boundary the staged-before arm idled 0.017-0.028 s. A stricter estimator, which also counts idle between the last micro-batch and the deferred update, gives a median of 0.026 s (max 0.045 s).
• In the staged-before arm, wait_for_device inside the deferred update still waited about 6.7-9.0 s: the last micro-batch was still computing after the stage returned. Staged after the update it read 2.83-2.94 s, the floor this wait shows on Pathways with nothing left to compute.

Context. google/tunix#2878 tabulates the 128-chip stack under this result, with each stage's numerics check: from the team cell (130.80 s per optimizer step, 2,255 tok/s/chip) to the final stack without staging (106.16 s, mean of 3 pods). 3,000 tok/s/chip is 94% of the pretraining target (3,193 tok/s/chip, from the MaxText pretraining run of this model on 128 v7x chips, rows without padding); on non-padding tokens, about 2,712 (85%).

Correctness

Check	Where	Result
Training with prepared inputs is bit-identical (weights, losses, gradient norms, replay counts), including a rerun of the last micro-batch inside update() after the next one was prepared, and through Tunix's real TrainerWorker; the negative controls are caught	CPU, the 4 new test files	pass: 18 tests and 10 subtests, plus 3 tests in a 4-device child
Same micro-batch inputs in the same order (by a per-micro-batch fingerprint), one update per step, staging lines and call counts as expected	Bench, the 128-chip pair without the precision keys	pass (every exact check). The call sequence, including apply_optimizer, is pinned by the bench switch's CPU tests (#5719), not on TPU
Same micro-batches every step (content digests), equal summed loss denominators, step k+1 only after update k, every staged-before stage at the trainer before the update, no backend compile of 5 s or more after warm-up	Reduced rig, all 6 arms (5 scored plus warm-up; no-compile check: the 5 scored arms)	pass
Real-gradient A/B (bench log-prob gates off)	TPU, 128 chips	not run: written down, but no TPU slices could be obtained
End-to-end RL run (rollout, sandbox, reward)	-	not run

The TPU runs cannot settle numerics. The bench's speed runs log loss and gradient norm as 0 by design: its log-prob gates drop every sequence of a random-weight model. The real-gradient A/B turns the gates off and gives every sequence segment a seeded random advantage. The rig runs at learning rate 0 and its loss moved by more than 2% from step to step even with staging off, so its automatic verdict reads FAIL although 74 of 79 checks passed (2 gradient-norm bands missed, 3 loss bands unscorable); every order, data and decision check passed. Until the real-gradient A/B runs, numerics rest on the CPU tests; google/tunix#2878 lists the caller's own tests.

How to run end to end

Run it now (before merge, for example for gate 2): skip steps 1-2. google/tunix#2878 already pins this PR's head (c79bcfd16) in both MaxText lines of requirements/maxtext_requirements.txt. After merge: steps 1-3.

1. Merge this PR into atwigg/mlperf.
2. Re-pin Tunix. In 397B v7x RL trainer: tuned settings, sharded log-prob head, staged step boundary (130.8 -> 98.3 s per optimizer step, trainer bench) google/tunix#2878, move both MaxText lines (maxtext and maxtext-vllm-adapter) from c79bcfd16 to this PR's squash-merge commit. MaxText squash-merges and deletes the branch, so pin the commit that lands on atwigg/mlperf.
3. Build the image and launch. Build the trellis image from 397B v7x RL trainer: tuned settings, sharded log-prob head, staged step boundary (130.8 -> 98.3 s per optimizer step, trainer bench) google/tunix#2878's tree (bash build_docker.sh --maxtext --raiden --deepswe; its Dockerfile installs MaxText and the vLLM adapter from the pin). Create raiden_wheels/ first: the Dockerfile copies it, and it is not in git (see 397B v7x RL trainer: tuned settings, sharded log-prob head, staged step boundary (130.8 -> 98.3 s per optimizer step, trainer bench) google/tunix#2878's build notes). Push the image, and launch the submission recipe with the image pinned by digest:

export TUNIX_IMAGE="gcr.io/<project>/<user>/trellis@sha256:<digest>"
bash tunix/experimental/examples/recipes/mlperf_gbs1024_submission.sh start

• The recipe's default TUNIX_IMAGE (gcr.io/cloud-tpu-multipod-dev/atwigg/trellis:latest, built on 2026-09-27) has neither this PR nor 397B v7x RL trainer: tuned settings, sharded log-prob head, staged step boundary (130.8 -> 98.3 s per optimizer step, trainer bench) google/tunix#2878's flags. A launch with it fails at argument or config parsing instead of running the wrong code.
• The default launch renders the trainer of the 98.31 s bench arm: the same values on all 75 MaxText keys both carry and the same libtpu flags (17 in total, including 397B v7x RL trainer: tuned settings, sharded log-prob head, staged step boundary (130.8 -> 98.3 s per optimizer step, trainer bench) google/tunix#2878's one). It differs in the batch (GBS 1,024: batch 64, learning rate 2e-6, gradient clip 0.0625; the bench ran GBS 1,152) and in two keys only the bench set: enable_checkpointing=false and the inert moe_route=device.

The switches, with their defaults in mlperf_gbs1024_submission.sh:

Switch	Default	What it does	To turn it off
PIPELINE_STEP_BOUNDARY	1	The pipelined boundary	PIPELINE_STEP_BOUNDARY=false, which turns staging off too
STAGE_NEXT_MICROBATCH	follows PIPELINE_STEP_BOUNDARY	Staging through this PR's two methods	STAGE_NEXT_MICROBATCH=false keeps the pipelined boundary without staging
RAGGED_BUFFER_FACTOR	2.75 with RETRY_WHEN_TOKENS_DROPPED=true	The replay layout where staging was measured	-1 gives the fully dropless layout, not measured with staging. Write -1, not -1.0: the recipe compares the value as a string.

The stage-before-update order is the orchestrator's default (--stage_before_update in run_deepswe_dist.py). The recipe does not pass this flag, so trying the other order (--no-stage_before_update; 7.43 s idle per boundary on the reduced rig) needs an edit to the launcher's orchestrator command.

Every other recipe keeps staging off, and STAGE_NEXT_MICROBATCH without PIPELINE_STEP_BOUNDARY is refused. google/tunix#2878 describes its other switches (TRAINER_TUNED_FLAGS, TUNIX_LOGPS_SHARDED_HEAD).

What the orchestrator log should show:

• [pipeline] STAGED next step's first microbatch on the trainer before the update in <s>; the last microbatch then still computed for <s>, after optimizer steps 1-6 of the recipe's 10-step run (6 of 9 boundaries), or fewer if rollouts are late. Tunix refuses to stage after the steps that write scored checkpoints (from val_start_step 7). A "still computed" value of about 2.7-2.9 s means the chips were already idle (the Pathways floor, measured on the 2026-10-09 build, not on the recipe's 2026-10-13 build), not that time was left.
• Never [pipeline] stage_next_microbatch is set but the trainer cannot stage a microbatch ...: it means the image's MaxText lacks this PR and staging is off.

Not yet tested / gates before un-drafting

This PR stays a draft until these have run on TPU and passed. Each needs TPU time that has not been available.

1. Real-gradient A/B of the staged step boundary at 128 chips. Staging on against two staging-off runs in the same pod, on the bench, through the switch that calls these two methods. Pass if loss, grad_norm, entropy and the forward log-prob metric at steps 0, 1 and 2 stay inside the band between the two staging-off runs.
2. One end-to-end smoke run from an image built at the pin. Checkpointing on, real rollouts, and if possible a forced overflow at 128 chips, so that a dropless rerun follows a prepare on TPU.

Also pending on the Tunix side, outside this PR's code: google/tunix#2878's same-pod real-gradient A/B of its first commit's settings (the libtpu flag and the six MaxText keys it ships; moe_route is now left out) against the team cell, with pass bounds written down before the run.

Not yet tested	Status
Tunix's orchestrator path at the tuned trainer settings	Ran on TPU only in the reduced rig (the team cell, 13.4 s per micro-batch, no weight sync or checkpoint). With the tuned settings (10.75 s per micro-batch) the stage's gRPC transfer of about 5 GB has not been timed.
Fully dropless layout (ragged_buffer_factor -1, no replay)	Not measured with staging. There update() does not wait for the last micro-batch, so the deferred update's wait_for_device may hold it back by one readiness round trip (about 2.7-2.9 s on Pathways, inferred). The submission recipe defaults to the replay layout.
This head on its own base, on TPU	The TPU runs used this PR's engine changes (equal to this head's except docstrings), overlaid on the bench image's MaxText. That MaxText differs from this PR's base in 395 files, including maxtext_engine.py itself and 2 other training_engine/ files.
Pathways build	Measured on the 2026-10-09 build; google/tunix#2878's recipe uses the 2026-10-13 build.

Commits in this PR

Six linear commits on 60425bfbd (atwigg/mlperf):

#	Commit	What it does
1	[MaxTextEngine] Prepare a micro-batch's inputs on the mesh ahead of its fwd_bwd	Adds PreparedFwdBwdInput, prepare_fwd_bwd_input and the fwd_bwd change, with the one-device tests.
2	[MaxTextEngine] Report how long a caller waited for the device	Adds wait_for_device and InflightThrottler.block_until_ready, their tests and the 4-device test.
3	[MaxTextEngine] Test prepared micro-batches under dropless replay and the fused update	Brings both methods up to the current engine (dropless replay, final dropless program, Tunix's fused update); adds the dropless test, reworks the TrainerWorker staging test, makes the 4-device test re-exec itself; README.
4	[MaxTextEngine] Test the order Tunix uses to stage before a replayed update	Adds maxtext_engine_stage_before_update_test.py: the staged-before-update order with reruns inside update(), and two negative controls.
5	[MaxTextEngine] Drive the staged-before-update order through Tunix's TrainerWorker	The same order through a real TrainerWorker (defer_update, apply_deferred_update).
6	[MaxTextEngine] Document that staging leaves a pending overflow check alone	Docstrings and README only.

Superseded PRs (do not merge)

Each is titled [DO NOT MERGE] and opens with a pointer note.

PR	What it was	Status
#5719	The STAGE_BEFORE_UPDATE switch in #5701's bench driver, which calls these two methods. An earlier form of it ran the three staging A/Bs in Results; the earlier stages of the stack (130.80 s to 106.16 s) ran on #5701's driver without it.	Bench only
#5723	Integration build: atwigg/mlperf plus this PR	Superseded by this PR
google/tunix#2877	The Tunix caller (the staged step boundary), now in google/tunix#2878	Kept for its review history and the reduced-rig write-up
google/tunix#2873	The sharded log-prob head, now in google/tunix#2878	Kept for its review history and as the exact file measured on TPU (a29991d3f)
google/tunix#2876	Lets the old log-prob head compile at chunks above 1024 under context parallelism	Reference only
google/tunix#2887	Integration build of every Tunix change of the 128-chip run	Superseded by google/tunix#2878

Known limitations and follow-ups

None of these blocks review.

Item	Status and follow-up
wait_for_device on Pathways	Read 2.69-2.94 s with nothing left to compute (two rig runs, 2026-10-09 Pathways build), so its docstring's "~0 means the device was already idle" does not hold there. Tunix turns staging off for the run when this wait reads under 0.05 s; on Pathways that never triggers, so a stage that outlasted the last micro-batch would delay the update without a warning. Follow-up: correct the docstring here; calibrate the switch in Tunix.
Lazily jitted fwd_bwd (compile(None))	The first prepared input re-traces and compiles fwd_bwd once, at the first staged boundary. AOT kernels (compile(dummy_data)) are unaffected, and the submission recipe compiles them at warm-up. The tests count _compile_for_batch calls and cannot see the re-trace. Follow-up: compile the lowered program on that path, as the background compiles already do, or warn.
Test gaps	No test checks that prepare_fwd_bwd_input does not block, and none pins the seconds wait_for_device returns: four blocking variants of the prepare, and three wrong-value variants of the wait, pass the tests they were run against. The only multi-device test uses an all-fsdp mesh (production is fsdp 32 x context 4 x expert 2). Follow-up: add those tests.
Device memory	One extra micro-batch of inputs from the prepare until its fwd_bwd (under replay, until its overflow check resolves): about 20 MB per device on the bench (5.15 GB over 256 devices; derived). The in-flight throttler does not count it and no compile can see it. It fit on the bench.
What the wait covers	The in-flight queue only (queued losses and the update's gradient norm), not the accumulate and gradient-cast programs or a prepared input's own copy.
Multi-controller JAX	Not covered: prepared inputs are placed with jax.device_put on every traced leaf, without the zero-width and non-fully-addressable skips of _place_read_only_inputs.
A MaxText without this PR	Tunix warns once and runs the boundary unstaged, so a stale image loses the win without failing. Follow-up in Tunix: fail at bring-up when STAGE_NEXT_MICROBATCH is set and the trainer cannot stage.
Tunix-side follow-ups	Listed in google/tunix#2878: how the tuned keys combine with a caller's MaxText flags, the drop-audit default, and two test gaps.

API contracts, dropless replay and the fused update

prepare_fwd_bwd_input(payload) -> PreparedFwdBwdInput

• Runs _prepare_batch(payload), the mapping fwd_bwd runs. With a mesh and a dict batch it jax.device_puts the traced half onto _batch_data_shardings (the per-leaf in_shardings of fwd_bwd's compile) and leaves the closed-over half (Tunix's algo_config, pad_id, eos_id) as is. Otherwise it returns the mapped batch, and fwd_bwd places it as before.
• Returns without waiting for the copy. Changes no step state (micro-step count, train step, weights, in-flight queue, pending overflow check); the only side effect is a warn-once flag for a batch that does not divide its mesh axis. Takes no lock.
• Safe between a step's last fwd_bwd and its update(): under dropless replay it neither reads nor resolves the last micro-batch's pending overflow flag, and the mapping reads only the payload, so the prepared input is the same before or after the update.

wait_for_device() -> float

• Blocks until every entry of the in-flight queue is ready (each queued fwd_bwd loss and the last update's gradient norm; at most max_inflight_computations, 2 by default) and returns the seconds it waited.
• Dequeues nothing and writes no metrics, so the next wait_for_next / wait_for_all and their metrics are unchanged. A pending overflow check, and its rerun, stay with the next fwd_bwd or update().
• Raises what jax.block_until_ready raises, so a failed micro-batch or update surfaces here; Tunix reports it as the update's failure.
• Blocks instead of polling is_ready(): under Pathways a readiness check is a proxy RPC and can read "not ready" for finished work.
• Composes with feat(training_engine): expose step-timing readiness tokens for Tunix TRAINER_STEP_TIMING #5571 (open): its last_update_token is the gradient norm waited on here.

Path	With a prepared input	Pinned by
Per-micro-batch dropless replay (retry_when_tokens_dropped, #5570)	A prepared micro-batch that overflows is rerun from its own prepared inputs. The last micro-batch's rerun inside update() reads that micro-batch's kept inputs, even when the next one was already prepared.	maxtext_engine_prepared_input_dropless_test.py; maxtext_engine_stage_before_update_test.py (replays per step [1, 2, 1])
The final micro-batch's dropless program (retry_dropless_final_micro_batch, #5673)	Chosen by apply_optimizer in fwd_bwd, exactly as for a payload.	Both files above, with the final program on (replays [0, 1, 0] in the order test)
Tunix's fused final update (google/tunix#2719)	In Tunix's staged order the worker defers the update until after the stage, so the engine still gets one update() per step. With --no-stage_before_update the stage lands after the fused update; that is bit-identical too.	maxtext_engine_prepared_input_test.py; maxtext_engine_stage_before_update_test.py (both also through TrainerWorker)

The supported call order, for any caller:

engine.fwd_bwd(last_micro_batch, apply_optimizer=True)
prepared = engine.prepare_fwd_bwd_input(next_step_first_micro_batch)  # returns while the copy runs
waited_s = engine.wait_for_device()  # optional: how long the chips were still busy
engine.update()
# ... metrics, checkpoint, weight sync ...
engine.fwd_bwd(prepared, apply_optimizer=False)  # dispatches without mapping or copying

Repro: CPU tests and the bench A/B

CPU tests, from the repo root with Tunix importable:

export JAX_PLATFORMS=cpu
python -m pytest -q tests/post_training/unit/maxtext_engine_prepared_input_test.py \
  tests/post_training/unit/maxtext_engine_prepared_input_dropless_test.py \
  tests/post_training/unit/maxtext_engine_prepared_input_sharded_test.py \
  tests/post_training/unit/maxtext_engine_stage_before_update_test.py

The two TrainerWorker tests skip unless the installed Tunix has stage_fwd_bwd and apply_deferred_update (google/tunix#2878).

Bench A/B at 128 and 256 chips

This route was checked against the measured arm offline, not re-run on TPU: the measured arms were rendered from the team recipe's launch render, with their flags set explicitly.

1. Bench: Add an RL trainer bench under experimental/rl_trainer_bench #5701 (src/maxtext/experimental/rl_trainer_bench) at its commit "rl_trainer_bench: rename e2e identifiers to bench/launch names", with [DO NOT MERGE] perf(rl_trainer_bench): opt-in STAGE_BEFORE_UPDATE step boundary (-7.2% step time on the 397B bench; stacked on #5701) #5719's commit "rl_trainer_bench: STAGE_BEFORE_UPDATE switch for the step boundary" on top. The TPU runs used an earlier form of the same switch.

2. Image: a client image whose MaxText has this PR's training-engine changes and whose Tunix does not fuse the update into the last train call (with perf(orchestrator): fuse final microbatch fwd_bwd+update RPC and defer get_metrics off weight-sync path google/tunix#2719's fused update the switch raises at the first boundary). The measured runs used, on the Pathways build of 2026-10-09:

   • client gcr.io/tpu-vm-gke-testing/sivaibhav/trellis-experimental@sha256:b4647c188b25638e644c7a6fc53fb8e0751741525a277a1bcf07dca92ba9a6ac
   • Pathways server us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/gke/datenglin/unsanitized_server@sha256:662f33b0ca6cd236a12534e4c78a8b7d715868669366738a015401925a39cb86
   • Pathways proxy us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/gke/datenglin/unsanitized_proxy_server@sha256:e714351cdde50d306dd62f127cffabcdaf98dff80ad22e8b4c66c5a1d1df6ddb

   At pod start they overlaid files carrying this PR's training-engine changes on the image's MaxText, and installed the head file measured on TPU (google/tunix a29991d3f).

3. Launch render (Add an RL trainer bench under experimental/rl_trainer_bench #5701's README, step 1): DRY_RUN=true bash tunix/experimental/examples/recipes/mlperf_gbs1024_submission.sh start from 397B v7x RL trainer: tuned settings, sharded log-prob head, staged step boundary (130.8 -> 98.3 s per optimizer step, trainer bench) google/tunix#2878, with BATCH_SIZE=72 (1,152 trajectories per optimizer step, as measured).

4. Bench render (Add an RL trainer bench under experimental/rl_trainer_bench #5701's README, step 2): render_bench.sh with STEPS=5 MICRO_STEPS=9 and EXTRA_MAXTEXT_FLAGS="enable_checkpointing=false", as the measured arms set (they also set moe_route=device, which did nothing on that MaxText). Render twice, identical except DRIVER_ENV="STAGE_BEFORE_UPDATE=on" vs off.

5. Run both arms back to back in one pod.

6. Read the client log: completed step: <n>, seconds: <s>, ... fwd_bwd dispatch s: [...] for steps 1-4, and STAGE_BEFORE_UPDATE step=<k> staged=... stage_s=... device_wait_s=.... A healthy staged arm shows each step's 2nd dispatch entry near the regular call time, and device_wait_s far above stage_s.

7. 256 chips: render with TRAINER_SLICE=4x8x8 (fsdp 64 x context 4 x expert 2) and MICRO_STEPS=5 (micro-batches of 128 rows), and add FINAL_PAD_FRACTION=0.625 to DRIVER_ENV, as the measured arms did, so a step keeps the 560 real rows of the 128-chip runs.

Evidence: the tests in this PR and the recorded CPU runs

File	What it pins
maxtext_engine_prepared_input_test.py (one device, toy model, Tunix-shaped input; 9 tests)	Prepared micro-batches train bit-identically to the payloads (weights, every metric), compiled and eager, prepared before or after the update, with no extra _compile_for_batch call (the same count as the payload run); also with wait_for_device after each prepare, queue and metric writes unchanged. A spy on jax.block_until_ready sees every queued entry in order. Traced leaves carry exactly _batch_data_shardings; step state is untouched. Through Tunix's real TrainerWorker (skipped without staging). InflightThrottler.block_until_ready neither dequeues nor writes (3 tests). Positive control: a one-token change is seen.
maxtext_engine_prepared_input_dropless_test.py (one-layer Qwen3.5 MoE rig; replay and the final program on; 3 tests)	An overflowing prepared micro-batch is replayed from its prepared inputs, bit-identically (weights, losses, gradient norms, replay counts), prepared before or after the update. A prepared first micro-batch starts the payload's background compiles. Negative control: preparing the wrong micro-batch is caught.
maxtext_engine_prepared_input_sharded_test.py (1 test, plus 3 in its child)	Re-runs itself with 4 host devices and fails unless the child ran its tests. The child checks bit-identical training with Tunix's nested payload split four ways, leaf shardings equal to the compile's, and a one-token change.
maxtext_engine_stage_before_update_test.py (dropless rig on the compile(None) path; 3 steps of 3 micro-batches, each step's last one overflows; 5 tests)	fwd_bwd(last), prepare(next), wait_for_device(), update(), fwd_bwd(prepared) equals the base order: replays per step [1, 2, 1], the flag still pending after the wait, one _compile_for_batch call (the JAX re-trace on this path is not counted). The same with the final dropless program on ([0, 1, 0]), and through TrainerWorker with the deferred update (defer_update never reaches the engine). Negative controls: a rerun of the prepared micro-batch instead of its own, and preparing step k+1's 2nd micro-batch instead of its 1st, are both caught.

Recorded runs (CPU, jax 0.11.1):

• These 4 files at this head, with a Tunix whose TrainerWorker code equals 397B v7x RL trainer: tuned settings, sharded log-prob head, staged step boundary (130.8 -> 98.3 s per optimizer step, trainer bench) google/tunix#2878's (only docstrings differ): 18 passed, 10 subtests.
• The same 4 files with a Tunix that has no staging (atwigg/mlperf): 16 passed, 2 skipped (the two TrainerWorker tests).
• Engine suites, at a commit whose engine code equals this head (only the README and two test docstrings changed since): maxtext_engine_test.py 111 passed on base and branch, with the same 1 failure; the other suites (maxtext_engine_*_test.py, router_replay_engine_test.py) 130 passed and 13 failed of 143 on base, 148 passed and 13 failed of 161 on this branch. 12 failures are identical pre-existing ones; the 13th is a libtpu lockfile collision between parallel test workers, which hit a different test on each side.

Evidence: TPU run records (predictions, arm order, rig checks)

Bench (predictions as predicted value [allowed band])

• 128 chips, final stack: arms ran staged, off, staged (repeat) in one pod on 2026-10-14; the written plan listed the off arm first. Written before the run: off 106.15 s [105.5, 107.0], on 98.48 s [97.5, 99.8], change -7.68 s [-8.2, -7.1], repeat within 0.3 s. Measured: 105.97, 98.31, -7.66, and a repeat of 98.26 (0.05 s apart). All inside.
• 128 chips without the precision keys: arms ran off, then on: 107.73 -> 100.01 s. The step tail (last micro-batch, update, drain) stayed at 12.556 vs 12.548 s (medians of steps 1-4).
• 256 chips: predicted -11.7 s [-12.4, -7.6] with partial hiding (a 15.57 s copy against a ~12 s window); measured -15.36 s with complete hiding, outside the band on the good side (see Results). The off arm (71.47 s) also came in under its band [71.55, 72.38].

Reduced rig

• Tunix's boundary code was at 26b06c504; since then it changed only in comments, docstrings, one flag's help text and the keyword order in one constructor, plus an unrelated base merge.
• Input: 16 recorded trajectory groups, each used twice per step, which pack into 3 micro-batches of 64 rows of 65,536 tokens (224, 226 and 62 trajectories). The trainer kept padding out of MoE routing and ran the per-micro-batch drop audit.
• Arm order in the pod: warm-up (staged), flags off, staged before the update, flags off again, pipelined but not staged, staged after the update.
• Written before the run, as predicted value [allowed band]: idle 0.016 s [0, 3] staged before the update and 12.0 s [9, 15] with the flags off; step 52.6 s [48.1, 57.6] and 64.6 s [59.6, 70.6]. All four landed inside their bands, as did all 20 recorded predictions (separate from the scorer's 79 checks).
• Checks: those in the Correctness table (no-compile check: the 5 scored arms, since the warm-up arm absorbs first-time compiles; summed loss denominator 2,040,178 in every arm), plus logged loss and gradient norm equal to the trainer's own and 0 dropless replays in every arm.
• Decision checks written before the run, all passed: staged-before idle <= 3.0 s (0.026 s on the pre-registered reading, which takes the larger of the two idle estimators per boundary; 0.020 s raw), and a saving >= max(7.0 s, 3 x the 0.74 s difference between the two flags-off runs) (12.44 s). The staged-before arm also had to stage at every boundary, always before the update, with a device wait of at least 4 s each time; it did. The step-time saving, 13.63 s, landed inside its recorded prediction band [7, 16].
• The in-pod verdict line (FAIL, 74 of 79; see Correctness) missed only the gradient-norm bands of the pipelined and staged-after arms; the 3 loss bands were unscorable because the flags-off loss did not repeat within 2%.
• An earlier 2-micro-batch run of the same rig passed the same order checks but could not show the win (a step's 2nd fwd_bwd returns only when its micro-batch is done), so its step-time-saving prediction missed (4.78 s against 9-20 s). Its staged-after control expected Tunix to switch staging off; it did not, because with nothing left to compute wait_for_device read 2.69-2.87 s at all 8 staged boundaries.

Corrections to commit messages (pushed history is not rewritten)

[MaxTextEngine] Prepare a micro-batch's inputs on the mesh ahead of its fwd_bwd

• "no recompile" holds for AOT kernels only; on the lazily jitted path the first prepared input re-traces fwd_bwd once (see Known limitations).
• "until it is trained": under dropless replay the inputs stay on the devices until that micro-batch's overflow check resolves.
• "last micro-batch and update": in the order Tunix uses once staging is on, the update is deferred until after the prepare.

[MaxTextEngine] Report how long a caller waited for the device

• "every computation dispatched so far": it waits on the in-flight queue only (see "What the wait covers" under Known limitations).
• "a wait of ~0 means the device was already idle ... the RL loop stops staging": on Pathways an idle-device wait read 2.69-2.94 s, so Tunix never switches staging off there. In the order Tunix uses once staging is on, the wait runs inside the deferred update, after the stage.
• The last sentence calls the 4-device test "ported". It is new in this PR: maxtext_engine_prepared_input_sharded_test.py, with Tunix's real model input.

[MaxTextEngine] Test prepared micro-batches under dropless replay and the fused update

• "the second is where Tunix's stage lands once the update runs inside the last train call": in the order 397B v7x RL trainer: tuned settings, sharded log-prob head, staged step boundary (130.8 -> 98.3 s per optimizer step, trainer bench) google/tunix#2878 uses once staging is on, the worker defers the update, so the stage lands before it. Preparing after the update is the --no-stage_before_update order.

[MaxTextEngine] Test the order Tunix uses to stage before a replayed update

• "no kernel is rebuilt": the test counts _compile_for_batch calls on the compile(None) path, where the first prepared input re-traces fwd_bwd once (one more backend compile than the unstaged order).
• "(the MLPerf recipes)": the MLPerf recipes do not default to this layout; see the correction under the last commit below.

[MaxTextEngine] Document that staging leaves a pending overflow check alone

• "the team's trainer configuration" means the team cell measured on the RL trainer bench: ragged_buffer_factor 2.75 with retry_when_tokens_dropped. In the MLPerf recipes this layout is opt-in (RETRY_WHEN_TOKENS_DROPPED defaults to false); only 397B v7x RL trainer: tuned settings, sharded log-prob head, staged step boundary (130.8 -> 98.3 s per optimizer step, trainer bench) google/tunix#2878's submission recipe uses it by default.

[gemini-code-assist]

Code Review

This pull request introduces prepare_fwd_bwd_input and wait_for_device to the MaxTextTrainingEngine to support staging micro-batches and pipelining step boundaries, particularly for Tunix. It also adds block_until_ready to the InflightThrottler to wait for inflight computations without dequeuing them or writing metrics, accompanied by comprehensive unit tests. The reviewer suggests defensively mapping the device placement over dynamic batch leaves in prepare_fwd_bwd_input to robustly handle edge cases like zero-width or non-fully-addressable arrays in multi-controller JAX environments.

[gemini-code-assist]

To ensure robustness in multi-controller/multi-host JAX environments and handle edge cases like zero-width or non-fully-addressable arrays (similar to the defensive logic implemented in _place_read_only_inputs), we should defensively map the placement over the dynamic batch leaves. This avoids potential runtime errors or redundant device_put operations on already-sharded or non-addressable arrays.

    batch = self._prepare_batch(payload)
    if self._mesh is not None and isinstance(batch, dict):
      dynamic_batch, static_batch = _split_static_and_dynamic(batch)
      shardings = self._batch_data_shardings(dynamic_batch)

      def place(leaf: Any, sharding: Any) -> Any:
        if not isinstance(leaf, (jax.Array, np.ndarray)):
          return leaf
        if leaf.ndim == 0 or 0 in leaf.shape:
          return leaf
        if isinstance(leaf, jax.Array) and not leaf.is_fully_addressable:
          return leaf
        if sharding is None or getattr(leaf, 'sharding', None) == sharding:
          return leaf
        return jax.device_put(leaf, sharding)

      dynamic_batch = jax.tree.map(place, dynamic_batch, shardings)
      batch = {**dynamic_batch, **static_batch}
    return PreparedFwdBwdInput(batch=batch)