[DO NOT MERGE] perf(orchestrator): opt-in pipelined step boundary that stages the next micro-batch before the deferred update

[parambole]

Do not merge. This change now ships in #2878 (tunix side) with AI-Hypercomputer/maxtext#5718 (MaxText side). This PR stays open for its review history.

TL;DR

• What: two opt-in flags that shrink the trainer's idle gap at each step boundary. --pipeline_step_boundary packs and sends step k+1's first micro-batch without waiting for step k's bookkeeping. --stage_next_microbatch also uploads it to the trainer while the devices still compute step k's last micro-batch, by default holding step k's update back until the upload is in.
• Measured: boundary device idle 12.46 s → 0.02 s (median) on a 128-chip v7x trainer-only replay rig with 3 micro-batches per step (same-pod A/B). The same stage-before-update order, switched on inside the MLPerf RL trainer bench (Add an RL trainer bench under experimental/rl_trainer_bench AI-Hypercomputer/maxtext#5701) on 128 v7x chips: 107.73 → 100.01 s per step (−7.7 s, +7.7% tokens/s/chip), same pod. Not yet measured end to end.
• Default: both flags off, and off runs exactly today's loop (tests compare every trainer call with a frozen copy of the old loop). Staging needs 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) AI-Hypercomputer/maxtext#5718; without it the flag warns and the boundary runs pipelined but unstaged.
• Risk: small end-to-end gain while rollouts dominate the step (estimated at ~0-41 s per ~26-minute team run); ~5 GB more orchestrator host RAM; the latch meant to turn staging off when uploads are too slow cannot fire under Pathways.

How it works

The default order at a staged boundary (both flags on):

sequenceDiagram
    participant O as Orchestrator
    participant W as TrainerWorker
    participant D as TPU devices
    Note over O: first micro-batch of step k+1 already packed
    O->>W: fwd_bwd(last micro-batch of step k, defer_update=True)
    W->>D: run the last micro-batch of step k
    W-->>O: returns early, update held back
    rect rgba(46, 160, 67, 0.2)
        O->>W: stage_fwd_bwd(first micro-batch of step k+1, about 5 GB)
        Note over W,D: the upload overlaps the last micro-batch on the devices
    end
    O->>W: apply_deferred_update()
    W->>D: wait for the devices, then run the update of step k
    O->>W: get_metrics, save_checkpoint, weight-sync staging
    O->>W: fwd_bwd_staged(request_id), no data attached
    W->>D: first micro-batch of step k+1 starts right away
    Note over O: step k bookkeeping (metrics, logs, on_step_end) runs now

Loading

With the flags off, the devices sit idle from step k's update until step k+1's first micro-batch arrives: through the metrics fetch, save, weight-sync staging, the orchestrator's bookkeeping and packing, and the ~5 GB upload. --pipeline_step_boundary alone takes the bookkeeping and packing off that path; --stage_next_microbatch takes the upload off too. The update is held back because under MaxText's retry_when_tokens_dropped a call that runs the update returns only once the last micro-batch is done, so a stage sent after it would have only the ~1.5 s update left to hide behind.

Results

Trainer-only replay rig on 128 v7x chips: the team's MLPerf trainer cell (fsdp32 x cp4 x ep2, retry_when_tokens_dropped, lr 0), driven by the real StandardRLProgram and TrainerWorker over gRPC with frozen real trajectories. Run parambole-l3m, 3 micro-batches per step; all arms back to back in one run on the same pods, 5 steps each; medians.

Flags	Boundary device idle	Step time
off	12.46 s	64.45 s
off, repeat (noise floor)	11.72 s	61.85 s
--pipeline_step_boundary	11.34 s	62.26 s
+ --stage_next_microbatch --no-stage_before_update	7.43 s	59.23 s
+ --stage_next_microbatch (default order)	0.02 s	50.82 s

On the MLPerf RL trainer bench (AI-Hypercomputer/maxtext#5701, the team's trainer program with its real MaxText/libtpu flags; 128 v7x chips; one pod, same order; the bench driver given the same stage-before-update order behind an off-by-default switch):

Bench arm	Step time, steps 1-4 mean	First call of each step, extra over a regular micro-batch	Regular micro-batch
switch off (stock driver loop)	107.73 s	7.5-7.8 s	10.945 s
switch on	100.01 s	~0 s	10.945 s

−7.7 s per step (pre-registered −7.4, band −8.5 to −5.5), from the step boundary only: the regular micro-batch did not change. Same micro-batches, order and updates in both arms.

Exact in every arm: the same micro-batches in the same order (content digests), the same summed loss denominator at every step, and every micro-batch of step k+1 entering after update k. For scale, a team training run (siv-sk1r, 128-chip trainer) idles 27.6 s per boundary (a different run, not an A/B). On the 2-micro-batch rig (parambole-l3x), where this window is closed by construction, the default order cut idle from 12.40 to 7.56 s, but only about 2.3 s of that survives netting out each boundary's own upload time (upload speed varied between arms).

Caveats: rig numbers only (no rollouts, weight sync or checkpoint save; production has ~9 micro-batches per step). The loss does not repeat step to step at this cell even with the flags off, so numerics rest on the exact checks above and on the CPU tests. The rig ran this branch at 26b06c504 (later commits change only comments, docstrings and a help string) with the MaxText PR's code (MaxText 2d821dfc1).

Changes

• orchestrator/rl_program.py: a step's intake state moves into _StepState. _take_group and _prepack pack step k+1's groups ahead (at most one feed, one step ahead); _send_first_microbatch_ahead sends its first micro-batch before step k's bookkeeping. _stage_planned decides before step k's last call whether to stage; _stage_next_before_update (default) or _maybe_stage_next_microbatch (after the update) stage through _send_stage, which falls back to an unstaged send if the transfer fails.
• worker/trainer_worker.py: fwd_bwd(..., wait_for_update=True, defer_update=False); new stage_fwd_bwd, fwd_bwd_staged and apply_deferred_update. While an update is deferred, the worker refuses every call that needs the step closed; compile and restore_checkpoint drop a stale stage or deferred update.
• orchestrator/distributed_rl_engine.py: stage_train_step (returns StagedTrainStep), apply_deferred_update, trainer_staging_support, and train_step(staged_request_id=..., defer_update=...).
• queue_manager/trajectory_queue_manager.py: BatchOrderedQueueManager also passes a hole the serving cursor has already read past (serving >= next); unchanged for in-order consumers (fuzzed).
• run_deepswe_dist.py, k8s_launcher.sh, mlperf_base.sh: the flags; the recipe exits if staging is on without the pipelined boundary.

Flag	Env var	Default
--[no-]pipeline_step_boundary	PIPELINE_STEP_BOUNDARY	off
--[no-]stage_next_microbatch (needs the flag above)	STAGE_NEXT_MICROBATCH	off
--[no-]stage_before_update (where the stage goes; only acts with staging)	none, CLI only	on

Size: ~1,650 source lines and ~7,400 test lines.

Testing

599 CPU tests pass in the new and extended files (jax 0.11.1). Each new hook, broken on purpose, is caught: 11/11 stage-before-update mutants, plus 9/9 loop, engine, worker and launcher mutants and 11/11 failure-path mutants on the commit before stage-before-update. The existing suites fail on the same 22 pre-existing test ids on the base (atwigg/mlperf at 8d245c561) and on this branch.

export JAX_PLATFORMS=cpu  # with MaxText and google-metrax importable
python -m pytest -q tests/experimental/worker/trainer_worker_{deferred_update,staging}_test.py tests/experimental/orchestrator/distributed_rl_engine_{deferred_update,staging}_test.py tests/experimental/queue_manager/trajectory_queue_manager_test.py tests/experimental/examples/deepswe_dist/mlperf_submission_launcher_test.py
python -m pytest -q -n 6 tests/experimental/orchestrator/rl_program_stage_before_update_test.py
python -m pytest -q -n 6 tests/experimental/orchestrator/rl_program_pipelined_boundary{,_robustness}_test.py

The TPU rig is a harness around the real classes and is not part of this PR; its setup is in the details below.

Risks / rollout

• Rollout: merge with the flags off. Then turn on PIPELINE_STEP_BOUNDARY and watch orchestrator RAM (+5.11 GB for the held micro-batch on an earlier 128-chip rig; skip it where the orchestrator is near its limit, e.g. ~205 of 240.5 GB at 1,024 chips). Then turn on STAGE_NEXT_MICROBATCH once the MaxText pin includes 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) AI-Hypercomputer/maxtext#5718. Rollback is the flag.
• Blind latch: staging should switch itself off when an upload outlasts the last micro-batch (device wait under 50 ms), but under Pathways that wait read 2.7-2.9 s even with the devices already idle, so a slow upload would delay the update and weight sync at every staged boundary.
• Where staging acts: it needs at least 3 micro-batches per step to hide anything. No stage is planned after a step that writes a scored checkpoint (so none once validation starts), or for micro-batches rescored under the live weights first, so never with USE_ROLLOUT_LOGPS=true (the launcher's default; the MLPerf recipes set it to false). Without dropless replay, the default order pays one readiness round trip of idle per boundary (not measured on TPU); --no-stage_before_update may be better there (CLI only for now).
• Before leaving draft: a few wall-clock-ratio tests have failed under heavy machine load (never on a quiet one) and need to hold the feed or become perf tests. [Experimental] Add deadlock-free orchestrator memory backpressure for polling and dispatch #2864 edits the same intake code; whichever lands second ports the other.

Related

• perf(orchestrator): fuse final microbatch fwd_bwd+update RPC and defer get_metrics off weight-sync path #2719 (merged): fused final fwd_bwd + update and deferred get_metrics; this builds on it.
• Reland: pack the next microbatch while the trainer runs the current one (--pipeline_train_microbatches) #2590 (merged): packs micro-batch k+1 while k trains within a step; this carries that across the step boundary.
• feat(orchestrator): separate first-batch dispatch stagger; mlperf_gbs1024_submission holds batch 1 175 s #2845 (merged): first-batch dispatch stagger; tested together with this.
• 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) AI-Hypercomputer/maxtext#5718 (draft): prepare_fwd_bwd_input and wait_for_device, which turn staging on.
• feat(training_engine): expose step-timing readiness tokens for Tunix TRAINER_STEP_TIMING AI-Hypercomputer/maxtext#5571 (open): readiness tokens, so the stage or the deferred update can confirm the update.
• Log the trainer's device-side step timeline (TRAINER_STEP_TIMING) #2713 (open): TRAINER_STEP_TIMING, the instrument behind the team-run numbers.
• Overlapping open PRs: [Experimental] Add deadlock-free orchestrator memory backpressure for polling and dispatch #2864 (same intake code); Split rollout worker pools: pin batch b to pool b % (max_staleness+1), optional per-pool weight sync #2868, orchestrator: hold router-replay routing bit-packed until packed (opt-in, 0.6x per held trajectory) #2866 and test(orchestrator): start-up stagger end to end on the dispatch loop and launcher-to-parser #2849 (one-hunk conflicts, keep both sides).
• Merge order: 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) AI-Hypercomputer/maxtext#5718, then this PR (also safe before it), then a MaxText pin bump before enabling staging.

Full evidence and notes

The previous description, verbatim: the full rig setup, both TPU runs with every caveat, the CPU test table and the reproduction steps.

Summary

Opt-in changes that take orchestrator and transfer work off the trainer's step boundary. Everything is off by
default, and off runs exactly today's loop.

• --pipeline_step_boundary (PIPELINE_STEP_BOUNDARY=true):
  • packs step k+1's first micro-batch while step k's last micro-batches train;
  • sends it as soon as step k's checkpoint, weight sync (only its staging, with async weight sync), commit and
    deferred metrics fetch are done;
  • runs step k's step metrics, step log, on_step_end and host-memory release once that call has returned, while the
    device computes it.
• --stage_next_microbatch (STAGE_NEXT_MICROBATCH=true; needs the first flag): also ships that micro-batch's data to
  the trainer host while the device still computes step k's last micro-batch, so its train call carries no data.
  • By default (--stage_before_update) step k's last train call holds the update back (defer_update=True). The
    stage goes in between, and apply_deferred_update then runs the update.
  • --no-stage_before_update stages after the update's dispatch instead.
  • Staging needs MaxText's wait_for_device, and uses prepare_fwd_bwd_input when present (both from
    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) AI-Hypercomputer/maxtext#5718). With a MaxText that lacks wait_for_device, the flag logs a warning and the boundary
    runs pipelined but unstaged.

The trainer receives the same micro-batches, in the same order, and the same updates. Staging adds one data-only stage
call, and staging before the update moves the update into its own call.

Why

At every step boundary the trainer's devices sit idle until step k+1's first micro-batch is on them. Measured on the
team's 128-chip trainer (run siv-sk1r, its TRAINER_STEP_TIMING lines, 6 normal boundaries):

Leg after step k's update lands	Median (range)
device idle until step k+1's first micro-batch starts	27.6 s (26.2-31.6)
orchestrator bookkeeping, host-memory release and packing ([pipeline] COMMIT to the next train call)	5.75 s (5.43-6.54)
payload gRPC (64 rows with router-replay data), then trainer-side input work	~11.2 s, then ~3.9 s

--pipeline_step_boundary takes the bookkeeping leg off the boundary. Staging targets the transfer and input legs.

Why stage before the update: under MaxText's retry_when_tokens_dropped (opt-in in the recipes as
RETRY_WHEN_TOKENS_DROPPED=true; the team's 128-chip trainer runs and the TPU cell below use it with
ragged_buffer_factor=2.75), update() reads the last micro-batch's overflow flag on the host before it dispatches the
update. So a train call that runs the update returns only once the last micro-batch has finished. A stage sent after
that call finds the device with only the ~1.5 s update left, so it hides almost nothing, and the latch switches staging
off at the first boundary (shown on CPU, both ways). With the update deferred, the last train call returns as soon as
the micro-batch before it is done. The stage then overlaps the whole last micro-batch, and the update goes out when
update() would have dispatched it anyway (from the source; not measured separately). That needs at least 3
micro-batches per step (see caveat 1 under Evidence). The default order is chosen for replay; without replay,
--no-stage_before_update may be the better one (see Risks).

How it works

The order on the trainer at a boundary where a stage is planned:

	trainer calls, in order
flags off (today)	fwd_bwd(last) + update (one call, #2719) -> get_metrics -> save -> weight-sync staging -> fwd_bwd(k+1 first, payload)
--pipeline_step_boundary	the same calls; step k's bookkeeping runs after fwd_bwd(k+1 first) returns, while the device computes it
--stage_next_microbatch (default: before the update)	fwd_bwd(last, defer_update=True) -> stage_fwd_bwd(k+1 first) -> apply_deferred_update -> get_metrics -> save -> weight-sync staging -> fwd_bwd_staged(k+1 first)
--stage_next_microbatch --no-stage_before_update	fwd_bwd(last) + update (returns at the update's dispatch) -> stage_fwd_bwd(k+1 first) -> get_metrics -> save -> ...

With a trainer that confirms its update (last_update_token, AI-Hypercomputer/maxtext#5571), get_metrics stays
deferred until after the save, as #2719 does for a confirmed update.

The rules every mode keeps:

• get_metrics follows step k's update and precedes step k+1's first fwd_bwd;
• the MLPerf checkpoint stamp is taken after the update has landed;
• one update per step, never two.

The decision to stage is made before step k's last call, because that call has to know whether to defer. No stage is
planned:

• after a step that writes a scored checkpoint, so none once validation starts (every step from val_start_step on
  writes one);
• at the last step, or when the next step's first micro-batch is already known to be its last;
• for micro-batches rescored under the live weights first. So staging never happens with USE_ROLLOUT_LOGPS=true, the
  launcher's default; the MLPerf recipes set it to false.

A planned stage is skipped when the next micro-batch is not packed by the time the last call returns. The update is
then still deferred, and pays the device wait before it.

While an update is deferred, the worker refuses fwd_bwd, update, get_metrics, save_checkpoint,
prepare_weight_sync, per_token_logps, eval_step and run_eval. A stage that fails on the way to the trainer falls
back: the micro-batch is sent whole, staging switches off for the run, and the deferred update still runs. A stage the
trainer ran and that raised ends the run, as a failed train call does. In the default order that happens before step
k's update and checkpoint, which the base loop would have saved.

apply_deferred_update waits for the device first and reports how long (device_wait_s); while an update is deferred,
stage_fwd_bwd reports no wait. A wait under 50 ms means the upload outlasted the last micro-batch, and staging latches
off for the run. The wait does not cover the host-to-device copy of the staged input (caveat 5).

Changes

• tunix/experimental/orchestrator/rl_program.py
  • A step's intake state moves into _StepState.
  • Step k+1's groups are taken and packed ahead (_take_group, _prepack): at most one feed's output, one step ahead.
  • The send-ahead (_send_first_microbatch_ahead, logged [pipeline] SENT_AHEAD step=k).
  • Staging: _stage_planned, _trainer_can_stage_before_update, _stage_next_before_update, the
    after-the-update hook _maybe_stage_next_microbatch, and the shared send-and-fall-back path _send_stage.
• tunix/experimental/worker/trainer_worker.py
  • fwd_bwd(..., *, wait_for_update=True, defer_update=False).
  • stage_fwd_bwd holds one micro-batch. It prepares it with prepare_fwd_bwd_input when the trainer has that. While
    an update is deferred it waits for nothing; otherwise it waits for the device and confirms the update.
  • fwd_bwd_staged and apply_deferred_update.
  • info() advertises stage_fwd_bwd, reports_device_wait and stage_before_update.
  • compile and restore_checkpoint drop a stale stage or deferred update left behind by a dead orchestrator.
• tunix/experimental/orchestrator/distributed_rl_engine.py: stage_train_step, trainer_staging_support,
  apply_deferred_update, and train_step(staged_request_id=..., defer_update=...). These go through perf(orchestrator): fuse final microbatch fwd_bwd+update RPC and defer get_metrics off weight-sync path #2719's response
  handling, so a worker that ran the update after all counts as a fused call (one update, never two).
• tunix/experimental/queue_manager/trajectory_queue_manager.py: BatchOrderedQueueManager passes a hole the serving
  cursor has already read past (serving >= next). Without read-ahead this gives traces identical to today's rule
  (fuzzed).
• run_deepswe_dist.py, k8s_launcher.sh, mlperf_base.sh: the flags below. The recipe refuses staging without the
  pipelined boundary.

Size: ~1,650 source lines and ~7,400 test lines.

Flags

flag	env var (launcher, recipe)	default	effect
--[no-]pipeline_step_boundary	PIPELINE_STEP_BOUNDARY	off	send-ahead plus overlapped bookkeeping
--[no-]stage_next_microbatch	STAGE_NEXT_MICROBATCH	off	also stage; needs the flag above (otherwise the recipe exits and StandardRLProgram raises)
--[no-]stage_before_update	none (CLI only)	on	where the stage goes; only acts with --stage_next_microbatch and a trainer that advertises stage_before_update, otherwise the stage goes after the update's dispatch, with a warning

Evidence

TPU, 128 v7x chips (our run parambole-l3x, 2026-10-14)

A trainer-only replay rig at the team's MLPerf trainer cell: fsdp32 x cp4 x ep2, ragged_buffer_factor=2.75,
retry_when_tokens_dropped=true, lr 0. It ran this PR's Tunix code (byte-identical to this branch at 26b06c504; the
later commits change only comments, docstrings and one help string, and merge two upstream files this PR does not
touch) with the MaxText PR's code on the image's MaxText. Five steps per scored arm, all arms in one trainer process. All numbers are MEASURED except the
derived column.

arm	flags	boundary device idle, median (s)	step wall, median (s)	idle net of that boundary's own upload, median (s), DERIVED	stages
A0	off	12.40	40.70	5.25	none
A0p	off (repeat)	11.59	39.89	5.19	none
A1	--pipeline_step_boundary	11.34	39.64	3.86	none
A2	+ --stage_next_microbatch --no-stage_before_update	8.59	36.88	2.31	4/4 after the update; device wait 2.74-2.87 s; not latched
A3	+ --stage_next_microbatch (this PR's default)	7.56	35.92	2.91	4/4 before the update; device wait 2.69-2.84 s; not latched

"Boundary device idle" is measured from step k's last micro-batch finishing on the device to step k+1's first
micro-batch starting, minus the update's own device time:
G_full = (k+1's first micro-batch starts) - (update landed) + max(0, update dispatched - last micro-batch finished).
A3's four boundaries (6.86-8.17 s) are all below every flags-off boundary (9.40-12.76 s).

Exact checks:

• in all 6 arms, including the warm-up arm, every step trained the same micro-batches as the flags-off arm (per-step
  content digests), with equal summed loss denominators (1,020,089);
• in all 6 arms, every micro-batch of step i entered the trainer after update i-1 (by the train_step it reported), and
  the program's logged loss and grad_norm equal the trainer's;
• in the 5 scored arms, no backend compile of 5 s or more after the arm's first fwd_bwd (the warm-up arm took the
  first-time compiles);
• every arm completed its planned updates (3 in the warm-up arm, 5 in the others) with status ok.

Read these numbers with the caveats:

1. This rig cannot show the overlap the change is for, except at one boundary. It has 2 micro-batches per step;
   production has about 9. MaxText flushes the previous step's metrics inside a step's second fwd_bwd, and that flush
   waits for the device. With 2 micro-batches per step the second call is the last one, so it returns only once its
   micro-batch has finished. At every scored boundary of every arm the micro-batch was already done when the call
   returned (the readiness watcher, started at the return, fired 0.005-0.020 s later), so no scored stage overlapped
   compute. The one open window was the warm-up arm's first boundary, before any metrics were queued: there the last
   call returned 13.0 s early, the stage hid behind 5.56 s of remaining compute, and the boundary idle was 0.98 s (n=1).
   On the scored boundaries, what A3 gains comes from doing the upload and the input placement before the update
   instead of after the step's bookkeeping.

2. Upload speed varied. The same 5.15 GB payload took 4.0-7.6 s to reach the trainer, and A3's four stages happened
   to be fast (4.0-5.1 s) while A0's uploads were slow (6.9-7.5 s). Net of each boundary's own upload time (the derived
   column), stage-before-update saves about 2.3 s per boundary against flags off on this rig, not 4.8 s. Staging after
   the update does a little better here (2.3 vs 2.9 s net), INFERRED to be because its upload starts 1.5 s before the
   update lands and so overlaps it. Without compute to overlap, there is nothing for the before-update order to win.
   Placing a prepared input on the devices took a further 3.7-3.9 s after prepare_fwd_bwd_input returned, about what
   an unprepared payload costs inside its train call (3.7-4.0 s here).

3. Pre-registered before either attempt of this run, for this rig: a closed window at every boundary but the
   warm-up arm's first; A3's pre-update idle 8-11 s; A3's boundary idle about A2's + 1.6 s; the in-pod gap reading a
   false win of 0.4-1 s for A3; the A2 latch not tripping; and A3's step wall not below A1's. Measured: closed window
   (held), 7.54 s, A2 - 1.03 s raw (+0.60 s net of upload), 0.016 s, not tripped (held), and 3.7 s below A1. So two
   held and four missed, all in A3's favour, mostly through the upload speed in caveat 2.

4. The rig's in-pod verdict reads FAIL (63 of 70 gates; 9 of 11 predictions), and it cannot certify numerics
   either way. No exact check failed. The 7 misses:

   • 3: per-step loss within ±2% of the flags-off median, for A1, A2 and A3. Under the reading rule written before the
     run these are unscorable, because the flags-off arms fail the same band themselves (only 1 of 5 steps each is
     within ±2%, at identical inputs and lr 0). The trainer's loss is not repeatable step to step at this cell: the two
     flags-off arms' per-step losses range from -0.0415 to -0.0621 (cause not identified), and the scored staged
     arms' (A2, A3; -0.0471 to -0.0885) fall partly outside that range;
   • 1: A1's grad_norm left its ±2.5% band at 3 of 5 steps (+2.6%, -2.6%, and one step at +15%); A0, A0p, A2 and A3
     stayed inside it at every step;
   • 1: the pre-registered A2 latch did not trip (caveat 5);
   • 2: timing bands, namely the pipelined saving (1.06 s against a 1.5-8 s band) and an ordering gate that assumed A2
     latches.

   The two missed predictions are the step-wall savings: staged 4.78 s against 9-20 s, pipelined 1.06 s against
   1.5-8 s. Some hits are not evidence either: the three in-pod gap gates for A3 use a gap measured from the update's
   landing, which the reading rule calls a false win on this rig, and the "every stage hid" wait gates pass only
   because of the wait floor in caveat 5. The numerics claim rests on the exact checks above and on the CPU tests
   below.

5. The 50 ms latch is blind on this stack. The device wait the latch reads (in apply_deferred_update for A3, in
   stage_fwd_bwd for A2) read 2.7-2.9 s at all 8 scored stages, although the last micro-batch had finished before
   the stage began. It also returned 0.97-1.01 s before the staged input's copy was ready. So on this stack a stage
   that does not hide cannot switch itself off, and a copy that overruns is not seen.

6. Production gain is INFERRED, not measured. With ~9 micro-batches per step the last call returns about one
   micro-batch (13.4 s at this cell) before the device finishes. That is room for production's ~11.2 s upload (8.45 s
   median outside the weight-sync transfer window, 11.58 s inside it). If the ~2.8 s wait floor measured here also
   applies while compute remains, the update slips by up to ~1 s; the placement can finish after the update, before
   the staged call. Estimates made before this run: 16.8-20.7 s of boundary idle hidden per trainer-bound boundary.
   End to end at today's recipes the gain is small, because rollouts usually dominate the step: about 16-18 s per
   ~26-minute team run in one analysis of three team runs (about 0 in a run whose scored step waited on rollouts), and
   17-41 s per run in another; the pipelined boundary alone, about 12 s. None of these include the ~2.8 s device-wait
   floor measured here. The 3-micro-batch run below settles the rig question; an end-to-end run is still needed before quoting
   a whole-run speedup.

TPU, 128 v7x chips, 3 micro-batches per step (our run parambole-l3m, 2026-10-14): the decisive A/B

The rig above has 2 micro-batches per step, which closes the staging window. This run serves the same frozen real
trajectory groups twice per step, so every step has exactly 3 micro-batches (checked at every step of every arm). The
decision rules, bands and scorer were written and stored before the run.

arm	flags	boundary idle, median (s)	step time, median (s)
A0	off	12.46	64.45
A0'	off (repeat)	11.72	61.85
A1	--pipeline_step_boundary	11.34	62.26
A2	+ staging after the update	7.43	59.23
A3	+ staging before the update (default)	0.02	50.82

• Boundary idle is the device idle between step k's update landing and step k+1's first micro-batch starting, taken as
  the larger of two independent estimators. Every one of A3's 4 boundaries is at most 0.045 s.
• Saving A0 - A3: 12.44 s per boundary, against a 0.74 s difference between the two flags-off arms.
• Step time: -11.0 to -13.6 s per step against the two flags-off arms (predicted before the run: 64.6 -> 52.6 s).
• Correctness, exact in every arm: the same micro-batches in the same order (digests), the same summed loss
  denominator at every step (2,040,178), and every micro-batch of step k+1 entering after update k. In A3 the stage
  reached the trainer before the update at every boundary.
• The loss/grad-norm noise gates miss in the two flags-off arms against each other as well (learning rate 0; the
  trainer at this cell does not repeat step to step), so numerics rest on the exact checks above plus the CPU tests.
• The code under test was the overlay this PR is built from (tunix 26b06c504, MaxText 2d821dfc1); later commits on
  this branch change comments, docstrings and a help string only.
• This is a rig number. End to end, the gain depends on how often the trainer, not rollouts, is on the critical path
  (see the estimates above).

CPU tests (jax 0.11.1, on a shared Linux machine; logs available)

Run	Result	Tag
Worker and engine deferred-update tests (31, 8), worker and engine staging tests (30, 17), extended queue-manager test (31)	117 passed	MEASURED
Launcher test (renders the recipe with both flags)	19 passed	MEASURED
rl_program_stage_before_update_test.py	165 passed, 150 subtests (6 workers, 7 min 08 s)	MEASURED
rl_program_pipelined_boundary_test.py + ..._robustness_test.py (the after-the-update hook, pinned with stage_before_update=False)	298 passed, 259 subtests (6 workers, 7 min 59 s). An earlier run of the same code (before the last, text-only commits), at a 1-minute machine load sampled at 65-270 during the run, failed one timing check: a boundary was sent ahead but not staged, because its micro-batch was packed late. The strict equivalence checks in that test passed, and it passed 3 of 3 serial reruns	MEASURED
Existing tests/experimental/{orchestrator,worker,queue_manager} + launcher test, base atwigg/mlperf (8d245c561) vs this branch (the three long files above run separately)	base 802 cases, head 891 cases: the same 22 failing ids, all failing on the base too: 11 need sqlalchemy, 1 needs tpu_sync (collection error), 5 PaddedBatchAssemblerTest cases are stale on the base (they treat _pack_chunk's tuple as a payload), and 5 AbstractWorkerTest cases fail in this environment. The AbstractWorkerTest ones construct TrainerWorker, so they give no signal about this PR's worker changes. Two RLProgramTest cases that write to hard-coded /tmp paths were deselected at the head; both pass on the base and passed in an earlier head run. The comparison reads DIFFERENT when one failure is removed (control)	MEASURED
Each new hook broken on purpose, a fresh copy per mutant	stage-before-update: 11 of 11 caught (2-37 failing tests each), on 26b06c504, the code of this branch. Loop, engine, worker and launcher: 9 of 9; failure paths: 11 of 11; both on the code before the stage-before-update commit, which then reworked some of their targets (the fall-back path, now _send_stage and shared by both orders, and the latch). The unmutated controls pass, and a mutation whose target is missing reports NOT-APPLIED	MEASURED
CPU rig: real StandardRLProgram + TrainerWorker + gRPC against a timed trainer stub that blocks like MaxText's replay	median boundary idle: off 3.33 s, pipelined 2.02 s, after the update 2.02 s (it latched at boundary 1, wait 2 µs), before the update 0.12 s; every stage before its update; loss and grad_norm bit-identical across arms	MEASURED
Same rig, real MaxText engine on 8 CPU devices, replay on	off, pipelined, both staged orders and off again: bit-identical grad_norm and loss at all 5 steps while the weights move (42 of 43 gates; the miss is a control that cannot fire on this engine). On CPU the before-update arm staged once and then latched, as pre-registered (XLA:CPU returns the last call only when its micro-batch is done), so it covers one staged boundary	MEASURED
Lint	pyink 25.12.0 clean on every changed line. pylint (the repo's pre-commit hook, --exit-zero): a few long docstring or pragma lines and not-callable warnings on getattr + callable() probes, a pattern the base already uses	MEASURED

How flags-off is shown to be today's loop: every scenario also runs through a byte-identical frozen copy of the loop as
it was before this change, pinned by the sha256 of its source, which equals atwigg/mlperf's _train_loop. Flags off
must match it in every trainer call, program event and metric. Both flag modes must match it in every trainer call:

• the micro-batch content hash and the weights version it runs on;
• the update and each get_metrics fetch;
• the checkpoint metadata and the weight-sync calls.

This is checked over 14 scenarios x 3 modes plus randomized draws, against a trainer that can confirm its update and
one that cannot. The stage-before-update file does the same against seven trainer doubles: two that do not block
(with and without an update token), two that block like MaxText's replay but never overflow, and three that also
overflow and really rerun (the last micro-batch, with and without a token, or the first). A separate test has every
micro-batch overflow. A replayed last micro-batch leaves the staged micro-batch trained once and unchanged.

How to reproduce

CPU, from the repo root, with MaxText and google-metrax importable:

export JAX_PLATFORMS=cpu
# 136 tests: the 117 in the first table row plus the launcher's 19
python -m pytest -q tests/experimental/worker/trainer_worker_deferred_update_test.py \
  tests/experimental/orchestrator/distributed_rl_engine_deferred_update_test.py \
  tests/experimental/worker/trainer_worker_staging_test.py \
  tests/experimental/orchestrator/distributed_rl_engine_staging_test.py \
  tests/experimental/queue_manager/trajectory_queue_manager_test.py \
  tests/experimental/examples/deepswe_dist/mlperf_submission_launcher_test.py
python -m pytest -q -n 6 tests/experimental/orchestrator/rl_program_stage_before_update_test.py
python -m pytest -q -n 6 tests/experimental/orchestrator/rl_program_pipelined_boundary_test.py \
  tests/experimental/orchestrator/rl_program_pipelined_boundary_robustness_test.py

TPU: the rig is a harness around the real classes. It is not part of this PR, and I can share it.

• One trainer process runs MaxText on 128 v7x chips (256 devices) at the cell above.
• Tunix's StandardRLProgram and TrainerWorker drive it over gRPC.
• Input: 256 frozen trajectories from a team run's first step, i.e. 2 micro-batches of 64 packed 65,536-token rows per
  step, 5.15 GB of payload each. No rollout, weight sync or checkpoint save.
• Arms, back to back: W (warm-up, staged, 3 steps), then A0 (flags off), A3, A2, A1 and A0p (flags off again), 5 steps
  each.
• Images: client trellis-experimental@sha256:b4647c188b25638e644c7a6fc53fb8e0751741525a277a1bcf07dca92ba9a6ac with
  this PR's files overlaid; Pathways server and proxy from the Oct-9 build (662f33b0…, e714351c…).
• The readout comes from trainer-side timestamps. A micro-batch's device start is its ready time minus the calibrated
  13.38 s micro-batch time.

To try it in a recipe: PIPELINE_STEP_BOUNDARY=true, then optionally STAGE_NEXT_MICROBATCH=true. Staging needs a
MaxText pin that includes the MaxText PR.

Risks and rollout

• Off by default, and off is today's loop (the frozen-copy tests). Rollback is the flag, or a revert.
• Suggested rollout:
  1. Merge with the flags off.
  2. Turn on PIPELINE_STEP_BOUNDARY and measure orchestrator RAM. It holds one packed micro-batch ahead: +5.11 GB on
     an earlier 128-chip rig, and up to ~5.2 GB at 64 rows (inferred). Do not enable it where the orchestrator is near
     its limit (at 1,024 chips it already peaks at ~205 GB of 240.5 GB).
  3. Turn on STAGE_NEXT_MICROBATCH after the MaxText pin bump (the 3-micro-batch A/B above is done: boundary idle
     12.46 s -> 0.02 s).
• With the flag on, step k's Train step line and on_step_end come after [pipeline] SENT_AHEAD step=k, and
  on_step_begin(k+1) fires after step k+1's first call returns. Under ARRIVAL order the dispatch window opens one train
  call later. PROMPT_BATCH, the recipe default, is unaffected.
• Staging before the update without replay: the deferred call returns at dispatch, and apply_deferred_update waits
  for the device before it dispatches the update. That costs at least one readiness round trip of device idle per
  staged boundary, which the after-the-update order does not pay (not measured on TPU). For runs without replay,
  --no-stage_before_update may be the better order; the launcher has no env var for it yet.
• If the upload and placement outlast the last micro-batch, that boundary's update and weight sync slip by the
  difference. The latch is meant to switch staging off after the first such boundary, but under Pathways it cannot see
  it (caveat 5), so the slip would recur at every staged boundary.
• A stage the trainer ran and that raised ends the run before step k's update and checkpoint (see How it works).
• Trainer memory, staging only: with prepare_fwd_bwd_input, about 20 MB per device of prepared inputs from the stage
  until that micro-batch trains (5.15 GB over 256 devices); without it, the 5.15 GB payload on the trainer host.
• fwd_bwd_staged called while an update is deferred would drop the staged micro-batch before fwd_bwd refuses. The
  program never makes that call (it applies the update first), and no test covers it.
• Known open item (low): if the reference log-probs call fails on the unstaged send-ahead, the run keeps step k's
  checkpoint but loses step k's step line, StepMetrics and on_step_end, which the base loop keeps. The run is ending
  in both cases.
• Some tests depend on wall-clock ratios and have failed under heavy load on a shared machine, never on a quiet one:
  BoundaryTimingTest, StagingLatchTest, the staged FailedUpdateTest cases, and the per-boundary stage checks in
  PipelinedBoundaryEquivalenceTest and StageBeforeUpdateEquivalenceTest. One of those checks failed once in the
  earlier run above and passed 3 of 3 serial reruns. Before this leaves draft, those checks should count only boundaries
  whose next micro-batch was packed in time (or the tests should hold the feed), or be marked as perf tests.
• Overlaps with open PRs: [Experimental] Add deadlock-free orchestrator memory backpressure for polling and dispatch #2864 (memory backpressure) edits the same intake code (4 hunks in rl_program.py), so
  whichever lands second ports the other's bookkeeping. Split rollout worker pools: pin batch b to pool b % (max_staleness+1), optional per-pool weight sync #2868, orchestrator: hold router-replay routing bit-packed until packed (opt-in, 0.6x per held trajectory) #2866 (__init__ keyword lines) and test(orchestrator): start-up stagger end to end on the dispatch loop and launcher-to-parser #2849 (launcher test
  methods) conflict trivially, 1 hunk each, keep both sides. feat(orchestrator): support MIN_ROLLOUT_REPLICAS and full-pipeline trajectory queueing metrics #2847, recipes(397b v7x): knobs for remat, GDN CP mode and expert-weight prefetch #2860, [DO NOT MERGE] Opt-in sharded GRPO log-prob head: gather the vocab table once, scan local tokens (-11.8% step time, v7x 397B bench) #2873 and fix(raiden): pin FFI control-plane listeners to primary host IP #2874 merge cleanly with this
  branch.

Related

• perf(orchestrator): fuse final microbatch fwd_bwd+update RPC and defer get_metrics off weight-sync path #2719 (merged): fused final fwd_bwd + update and the deferred get_metrics; this builds on it.
• Reland: pack the next microbatch while the trainer runs the current one (--pipeline_train_microbatches) #2590 (merged): packs micro-batch k+1 while k trains, within a step; this carries that across the step boundary.
• feat(orchestrator): separate first-batch dispatch stagger; mlperf_gbs1024_submission holds batch 1 175 s #2845 (merged): first-batch dispatch stagger; tested together here.
• 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) AI-Hypercomputer/maxtext#5718: prepare_fwd_bwd_input and wait_for_device, which turn staging on, plus engine
  tests of the stage-before-update order under dropless replay.
• feat(training_engine): expose step-timing readiness tokens for Tunix TRAINER_STEP_TIMING AI-Hypercomputer/maxtext#5571 (open): readiness tokens. With them, the stage or the deferred update confirms the
  update, and get_metrics stays deferred.
• Log the trainer's device-side step timeline (TRAINER_STEP_TIMING) #2713 (open): TRAINER_STEP_TIMING, the instrument behind the numbers in "Why".

Merge order

1. 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) AI-Hypercomputer/maxtext#5718 (additive).
2. This PR (flags off). It is also safe before 1: staging then stays dormant with a warning.
3. Bump requirements/maxtext_requirements.txt to a MaxText commit that has 1 (and #5571 for update confirmation)
   before enabling STAGE_NEXT_MICROBATCH.

[ultrons]

Heads-up from an e2e run of this PR together with #2713 (TRAINER_STEP_TIMING), on the 1013b image lineage (Tunix 21b62f2 + #2873 + #2878 + this PR's head 42ada6c + #2713), 128-chip trainer, 128 engines, PIPELINE_STEP_BOUNDARY=true STAGE_NEXT_MICROBATCH=true.

The run crashed at the first step boundary (end of step 0):

[pipeline] STAGED next step's first microbatch on the trainer before the update in 7.58s; the last microbatch then still computed for 7.00s
Synchronizing weights (target policy_version=1)...
[pipeline] SENT_AHEAD step=0 first microbatch of step 1 is on its way to the trainer; step 0's metrics follow once it is there.
Executing train_step on actor worker (accumulate_gradients=True, apply_optimizer=False)...
...
Exception in StandardRLProgram execution: cannot access local variable 'request' where it is not associated with a value
  File ".../orchestrator/rl_program.py", line 2212, in _train_loop -> _await_pending_train -> _timed_train_step
  File ".../orchestrator/distributed_rl_engine.py", line 1065, in train_step
    request.request_id,

The failing line is the "train_step fwd_bwd RPC returned ..." log statement from #2713, which reads request.request_id; on this PR's staged path (fwd_bwd_staged, staged_request_id set) no TrainRequest is built, so the name is unbound. So it is an interaction between the two PRs rather than a defect in either alone: #2873 + #2878 + this PR without #2713 would not hit it, and #2713 without this PR would not either. The weight-sync round that was in flight then ended in unknown_transfer_state as a consequence of the exception, not as its cause.

A one-line fix on the #2713 side (log staged_request_id if staged_request_id is not None else request.request_id) resolves it; worth keeping in mind when the two land. Happy to share the full orchestrator log.

[ultrons]

Follow-up with the exact fix we are running with, for review (interaction between this PR's staged path and the #2713 log line in DistributedRLEngine.train_step):

@@ -1062,7 +1062,7 @@ class DistributedRLEngine(rl_engine_interface.AbstractRLEngine):
           " (request_id=%s, apply_optimizer=%s)",
           role_name,
           time.monotonic() - t_rpc,
-          request.request_id,
+          staged_request_id if staged_request_id is not None else request.request_id,
           apply_optimizer,
       )
       if is_actor:

Commit on our branch: sivaibhav/mlperf-1013b-allprs @ 3d301c4 (ultrons/tunix), on top of 21b62f2 + #2873 + #2878 + this PR's head 42ada6c + #2713. With it the same e2e configuration (128-chip trainer, 128 engines, PIPELINE_STEP_BOUNDARY=true STAGE_NEXT_MICROBATCH=true, sharded head on) is relaunched and in start-up as of this comment; we will report once it has passed the first step boundary, and its step times. The change belongs with whichever of #2713 and this PR lands second: the log statement is #2713's, the staged_request_id branch is this PR's.