Repository navigation
Conversation
…opt-in) Re-applies the pipelined step boundary (parambole/u10-L3-pipelined-boundary at 0db0ac237) on #2719 merged with atwigg/mlperf. Two flags, both off by default, and off is the base loop exactly. --pipeline_step_boundary: while step k's last micro-batches train, take and pack step k+1's groups. Send step k+1's first micro-batch once step k's checkpoint save, weight-sync staging, commit and (deferred) metrics fetch are done, and run step k's step metrics, logging, on_step_end and host-memory release while that call is in flight. The trainer gets the same calls, with the same inputs, in the same order. --stage_next_microbatch: also ship that micro-batch to the trainer right after step k's last train call, while the device still computes, so the call that later runs it carries no data. #2719 runs the update inside that last call and, once the trainer exposes last_update_token, returns only after it landed, which would leave the transfer nothing to hide behind. So when a stage is planned the call is sent with wait_for_update=False (it returns at the update's dispatch) and the stage confirms the update instead (update_ready), which keeps #2719's deferred get_metrics at a staged boundary. A planned stage that finds nothing packed in time falls back to get_metrics ahead of the save. Trainer side: TrainerWorker.fwd_bwd(wait_for_update=...); stage_fwd_bwd holds one micro-batch and reports device_wait_s and update_ready; fwd_bwd_staged runs it through fwd_bwd, so a staged micro-batch gets the fused update too. DistributedRLEngine: stage_train_step, trainer_staging_support, and train_step(staged_request_id=...) sharing #2719's fused-response handling. BatchOrderedQueueManager passes a hole the serving cursor already read past. Recipe, launcher and driver plumbing for PIPELINE_STEP_BOUNDARY and STAGE_NEXT_MICROBATCH.
…start-up stagger Runs rollout_dispatch_stage with #2845's first-batch stagger in front of the train loop (max_staleness 1, so only the stagger holds batch 1): batch 1 is still held for the first gap in every mode, the prepack waits on the held batch without holding up step 0's boundary or the dispatcher, the next boundary is pipelined (and staged), and the trainer calls are the base loop's. Control: without the stagger batch 1 goes out at once.
On a box at load ~170 the first version failed in every mode, the base loop included: it assumed step 0 ends well within the 1 s gap, and with a background sync the dispatch window then rightly waits for step 0's round. Use a blocking sync so only the stagger holds batch 1, measure from the stagger's own reference point, keep the lower bound exact and allow for load above it, check the boundary-0 claim whenever its precondition holds, and stamp each delivered group with a fixed policy version so admission and metrics do not depend on timing.
Three fixes from the review of the pipelined step boundary. All are on the
staged path, which stays dormant until MaxText has wait_for_device.
* rl_program: a stage that fails on the way to the trainer, or in the
reference log-probs it needs, no longer ends the run before the step's
get_metrics, checkpoint and on_step_end. The micro-batch goes back to the
front of the next step's ready list and is sent whole in that step's
first train call, as without staging; staging is off for the rest of the
run. A stage the trainer ran and that raised there still ends the run
with its own error: the worker is then in ERROR, so no checkpoint could
follow anyway.
* trainer_worker: stage_fwd_bwd reports a failure in its device wait or in
its update wait as the optimizer update's ("optimizer update() failed
... seen by stage_fwd_bwd"), as the fused fwd_bwd does. MaxText's
wait_for_device waits on the update's gradient norm, so the device wait
is where a failed update shows up first.
* trainer_worker: compile and restore_checkpoint drop a micro-batch staged
before them, so an orchestrator that died without stop() no longer makes
every successor's first stage fail.
Tests. New rl_program_pipelined_boundary_robustness_test, adapted from the
review: the take-ahead's memory bound, also in a production-like streaming
regime, with a runaway it catches and the equivalence suite does not; a
failed update in every mode against a token-wait and a device-wait
trainer; a stage that fails in transport, and a reference call that fails
in transport or in the reference worker, fall back and train exactly what
the base loop trains (with a negative control); the hole-rule fuzz; a
stale stage after a crash. trainer_worker_staging_test: the attribution
of both waits, and compile/restore dropping a held stage.
With --stage_next_microbatch, the next step's first micro-batch went to the trainer right after the step's last train call, which runs the update and, when a stage follows, returns at the update's dispatch. Under MaxText's retry_when_tokens_dropped (dropless replay, maxtext#5570, and the final dropless micro-batch of #5673, in every MLPerf recipe today) update() reads the last micro-batch's overflow flag on the host before it dispatches the update, so that call returns only once the last micro-batch has finished. The transfer then found the device idle (only the update left), the device wait read ~0, and the latch switched staging off at the first boundary. Stage between the last micro-batch and its update instead, when the trainer can defer the update: - TrainerWorker.fwd_bwd(apply_optimizer=True, defer_update=True) runs the step's last micro-batch (the trainer still sees apply_optimizer, so the final dropless program is still chosen) and holds the update. apply_deferred_update() waits for the device (device_wait_s: how long the last micro-batch still ran once the stage was in place), runs the update, and confirms it when the trainer exposes last_update_token. A stage while an update is deferred waits for nothing. While it is deferred the worker refuses whatever needs the step closed (another micro-batch, update, get_metrics, a checkpoint, weight sync, actor log-probs, eval); compile and restore_checkpoint drop it, as they drop a stale stage. info() advertises stage_before_update. - DistributedRLEngine.train_step(defer_update=True) and apply_deferred_update(); trainer_staging_support() reports stage_before_update. A worker that runs the update anyway is handled as a fused call (one update, never two). - StandardRLProgram(stage_before_update=True), --[no-]stage_before_update: at a planned stage, when the trainer supports it, the last call is sent with its update deferred, then _stage_next_before_update stages the next micro-batch under the same refusals as before and runs the update; the latch reads the deferred update's device wait. Otherwise (an older trainer image, or --no-stage_before_update) the stage goes after the last call, as before. The shared send/fallback path moves into _send_stage. A replay of the last micro-batch inside update() needs nothing from the staged micro-batch: its inputs are mapped from its own payload alone, and the rerun reads the inputs it kept. Tests: - rl_program_stage_before_update_test.py: trainer doubles with MaxText's host blocking and real replays (the last, the first, or every micro-batch overflows and is rerun when its flag is read). Every harness scenario against seven doubles trains exactly what the base loop trains (strict check, reruns included) with every stage before its update; replays leave the staged micro-batch trained once and unchanged; under replay the stage hides (tr-s3 ratios: device waits 5.4-5.5 units, boundary gap 6.7 vs 15.3 for the old hook, which latches after one stage); feed-bound still latches; randomized draws; hook selection, checkpoint stamps, failure fallbacks. Controls: a rerun that reads the staged micro-batch, an update run before the stage, a stamp at the last call's return, an update measured before the device wait, and an update never applied are each caught. - trainer_worker_deferred_update_test.py, distributed_rl_engine_deferred_update_test.py. - The pipelined-boundary suite pins the old hook (stage_before_update=False); the staging-support dicts gain stage_before_update.
…DED_HEAD) compute_chunked_logps scans the sequence under GSPMD. On the MLPerf Qwen3.5-397B trainer cell (fsdp32 x cp4 x ep2 on 256 devices, 64 x 65,536-token packed micro-batches, compute_logps_chunk_size 1024) the compiled program all-gathers the 2 GB vocab table in every scan step (forward and rematerialized backward) and all-reduces a 2 GB table gradient in every backward step: 129 table all-gathers and 65 table all-reduces per micro-batch. Because the scan runs over a context-sharded sequence, every context rank also computes every chunk, so the vocab matmuls are 4x redundant. With TUNIX_LOGPS_SHARDED_HEAD=local (or 1), packed rows go through a shard_map'ed head that runs the same per-chunk math: - each head parameter is all-gathered once per pass; - each device scans only the tokens it holds (dynamic slices, no reshape of the sequence); - the table gradient is accumulated per device in float32 and reduce-scattered once; - the one-token shifts (targets left, outputs right) are explicit collective-permutes inside the shard_map, so they do not depend on how XLA partitions a shift of a sequence-sharded array. The head falls back to compute_chunked_logps, with a log line, when there is no mesh in context, when some mesh axis shards neither the batch nor the sequence, or when a head parameter does not divide over the mesh. Unset (the default), nothing changes.
Comments only, no code change: the ASTs of both modules are identical before and after. - rl_program.py, at the step's final micro-batch: why _stage_next_before_update, which always applies the update, runs before get_metrics, the checkpoint and the weight sync (the worker refuses them while an update is deferred); what a result without update_deferred means (the trainer ran the update inside the call, so the stage goes after it, as with --no-stage_before_update); and why stage_after is cleared once the deferred path ran (staging after the update's dispatch is the order this path replaces, and under dropless replay it would find the device idle and latch staging off). - trainer_worker.py, apply_deferred_update: why the deferred flag is cleared only after the device wait (update() refuses to run while an update is deferred, and the wait is what tells the caller whether the stage hid behind the step's last micro-batch).
Text only. The stagger test's docstring named the shared CPU machine its
first runs used; a cost comment cited an internal write-up for the
siv-sk1r numbers, which come from that run's TRAINER_STEP_TIMING lines;
and test docstrings and comments cited internal review and design item
ids ("Rule 7", "review P1/P2", "review finding F1", "diag-L3 R1/R3/R5").
They now say what the tests do. No code change: the module ASTs are
identical apart from docstrings.
…omment The --pipeline_step_boundary help text and the comment on StandardRLProgram.pipeline_step_boundary said step k's metrics and logging run while step k+1's first micro-batch "is on its way". The loop waits for that train call to return first (the SENT_AHEAD block in _train_loop), so they run once it has returned, while the device computes it. The help text also names what the send waits for: the checkpoint, the weight sync (only its staging with --async_weight_sync) and the commit. Help text and comment only; no behaviour change.
TRAINER_TUNED_FLAGS (true or false) applies the trainer settings measured on the RL trainer bench (AI-Hypercomputer/maxtext#5701) at the recipes' 128-chip trainer cell. They are the settings of runs 59 and 62 in AI-Hypercomputer/maxtext#5409 that the recipe does not use yet: - the trainer libtpu flag --xla_tpu_enable_offloading_copy_to_sparsecore=false, appended to TRAINER_LIBTPU_INIT_ARGS before the trainer env, the Pathways workers and the Pathways proxy copy it (a value already there is kept); - seven MaxText keys: context_parallel_strategy=all_gather sa_fuse_reciprocal=true sa_use_base2_exp=false gdn_cp_mode=seq gdn_cp_state=device moe_route=device, appended last (MaxText keeps the last value of a repeated key), and num_moe_token_chunks=2 through NUM_MOE_TOKEN_CHUNKS. The switch is on in mlperf_gbs1024_submission.sh. It is off in mlperf_397b_256_v7x.sh and in the recipes that source it, whose renders are byte-identical to before; in the submission recipe, TRAINER_TUNED_FLAGS=false renders the trainer as before. With the replay layout (a finite RAGGED_BUFFER_FACTOR with RETRY_WHEN_TOKENS_DROPPED=true), where it was measured, the switch sets two MoE token chunks instead of four. The fully dropless layout (RAGGED_BUFFER_FACTOR=-1) keeps four, and an explicit NUM_MOE_TOKEN_CHUNKS wins. Bench, 128 v7x chips, ragged_buffer_factor 2.75: 130.8 s/step for the team cell; 126.4 with the libtpu flag (another pod); 121.7 with the flag and the seven keys (same pod as 126.4). Tests: nine render tests in mlperf_submission_launcher_test.py. The existing test that pins four chunks for the replay layout now renders with the switch off.
TRAINER_TUNED_FLAGS=true now also appends gdn_cp_matmul_precision=high and cast_grads_after_all_reduce=true to the trainer's MaxText flags, in the same block as the seven keys: - gdn_cp_matmul_precision=high: under gdn_cp_mode=seq, the float32 matmuls that compose the GDN state across the sequence shards run 3 bf16 passes instead of 6. - cast_grads_after_all_reduce=true: a gradient that the forward/backward pass sums by an all-reduce alone, and that is not already float32, leaves it in the parameter's dtype and is cast to float32 after the all-reduce. Both keys have these meanings at this branch's MaxText pin (60425bf) and in the bench image's MaxText (3792fdef8). Bench (AI-Hypercomputer/maxtext#5701, 128 v7x chips), measured with the sharded log-prob head of #2873 on in both arms: 107.59 -> 106.24 s/step in one pod, and 107.71 -> 106.41 in a second pod that trained on real gradients. In that pod the gradient norm moved 1.47e-5 relative (two runs with the keys differ by 1.92e-5), the forward log-prob metric moved 1.1e-6 (3.6e-7), and the loss was identical as printed. Four MoE settings measured alongside them (prefuse_moe_weights=false, moe_accumulate_wi_dlhs=true, moe_accumulate_chunk_wgrad=true, decoder_layer_input=device) moved the gradient norm by -0.80% and the forward metric by 9.7e-4, so they are left out. Tests: the exact-set test expects nine keys, and _MAXTEXT_DEFAULTS has the two keys' MaxText defaults.
…tegration branch
…ranch One conflict, in tunix/experimental/orchestrator/rl_program.py, resolved by keeping both sides: StandardRLProgram.__init__ takes the over-generation arguments from #2633 (num_generations_to_dispatch, prefer_valid_rollouts) followed by #2877's pipeline_step_boundary, stage_next_microbatch and stage_before_update.
…ep boundary in the GBS-1,024 submission recipe - mlperf_gbs1024_submission.sh: TUNIX_LOGPS_SHARDED_HEAD=local (#2873), PIPELINE_STEP_BOUNDARY=1 and STAGE_NEXT_MICROBATCH=1 (#2877, the env vars that mlperf_base.sh and k8s_launcher.sh read). STAGE_NEXT_MICROBATCH follows PIPELINE_STEP_BOUNDARY unless set, so PIPELINE_STEP_BOUNDARY=false alone turns both off. - k8s_launcher.sh: passes TUNIX_LOGPS_SHARDED_HEAD to the trainer process when it is set. TRAINER_EXTRA_ENV cannot carry it from this recipe: a value set here replaces mlperf_397b_256_v7x.sh's default TRAINER_EXTRA_ENV, including the trainer's LIBTPU_INIT_ARGS. Unset, nothing is rendered, so the other recipes render as before. - mlperf_submission_launcher_test.py: the boundary is on by default in this recipe, PIPELINE_STEP_BOUNDARY=false turns both flags off, the refusal test sets PIPELINE_STEP_BOUNDARY=false, the head reaches the trainer only and 0 turns it off, and the other 397B v7x recipes keep both off.
This was referenced Oct 14, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Verification build, not for merging. This branch stacks every tunix change of the 128-chip best run on
atwigg/mlperf, so one image can reproduce that run. Land the component PRs instead.Companion: AI-Hypercomputer/maxtext#5723, the MaxText side of the same run, opened at the same time.
What is in it
Base:
atwigg/mlperfat4c7823f20. Merged with--no-ff, in this order:parambole/logps-sharded-headata29991d3f5d3f29b75TUNIX_LOGPS_SHARDED_HEAD)parambole/ours275-trainer-configat59ad7c8db998f587c6parambole/l3-convergeat42ada6cd08b7719e94PIPELINE_STEP_BOUNDARY,STAGE_NEXT_MICROBATCH)One conflict, in
rl_program.py: #2877 and #2633 (already on the base) both added keyword arguments toStandardRLProgram.__init__. I kept both. #2633 also adds a 4-line wait before the weight sync in_train_loop. It merged cleanly and only runs with over-generation, which this recipe leaves off.Extra commit
69962eb7a(integration only). Turns the opt-in pieces on inmlperf_gbs1024_submission.shonly:TUNIX_LOGPS_SHARDED_HEAD=localfor the trainer.k8s_launcher.shgets a one-line passthrough, rendered only when the variable is set. SettingTRAINER_EXTRA_ENVin this recipe instead would replace the 256 recipe's default list, including the trainer'sLIBTPU_INIT_ARGS.PIPELINE_STEP_BOUNDARY=1andSTAGE_NEXT_MICROBATCH=1, through [DO NOT MERGE] perf(orchestrator): opt-in pipelined step boundary that stages the next micro-batch before the deferred update #2877's env vars.Measured result this reproduces
Vaibhav's RL trainer bench (AI-Hypercomputer/maxtext#5701): 128 v7x chips, fsdp32 x context4 x expert2, 64 rows x 65,536 tokens per micro-batch, 9 micro-batches per optimizer step.
Per change: recipe keys + libtpu flag -8.80 s, sharded log-prob head -14.41 s, 2 precision flags -1.35 s, staged step boundary -7.66 s.
Correctness: the head and the precision flags passed a real-gradient TPU A/B. Real-gradient checks for the recipe keys and the staged boundary are scheduled for today.
How to verify
MICRO_STEPS=9. For the staged-boundary number, turn on the bench driver's opt-inSTAGE_BEFORE_UPDATEswitch.Caveats
py_compilepasses on all 16 changed.pyfiles andbash -non all 4 changed.shfiles. The CPU test suites were not run on this build.