Repository navigation
Conversation
A standalone timing harness for the RL trainer program of a Tunix GRPO e2e run: the same MaxText and libtpu flags, mesh, GRPO loss and chunked log-prob pass on one trainer slice, driven through Tunix's TrainerWorker with a synthetic micro-batch whose MoE routing load is calibrated to a production run, and without the orchestrator or rollout engines. render_e2e.py turns the DRY_RUN output of the e2e launcher into the bench JobSet, keeping the trainer JobSet's scheduling fields and Pathways containers and replacing only names, maxRestarts and the client command, with knobs for MaxText flags, libtpu flags, trainer and orchestrator argv, trainer slice and profiler window. Also included: an AOT memory check, a structured render diff, a MaxText config diff, CPU checks of the payload and the numerics knobs, a static HLO collective-overlap summary, and CPU unit tests of the renderer (on a hand-written fixture) and the payload builder.
…e DRY_RUN context switch
The bench is standalone; the e2e names only meant "derived from the e2e
launch's DRY_RUN render". Renames, no behaviour change:
- E2E_RENDER -> LAUNCH_RENDER; --e2e-render -> --launch-render
(render_bench.py and check_numerics_cpu.py)
- render_e2e.{sh,py} -> render_bench.{sh,py}
- e2e_driver.py -> bench_driver.py (also /tmp/bench_driver.py in the pod)
- E2E_TRAINER_ARGV / E2E_ORCH_ARGV -> TRAINER_ARGV / ORCH_ARGV
- aot_e2e.sh -> aot_bench.sh (default log logs/aot-bench-<time>.log)
- tests/render_e2e_test.py -> tests/render_bench_test.py;
tests/testdata/e2e_dryrun_fixture.yaml -> launch_dryrun_fixture.yaml
- heredoc E2E_DRIVER_EOF -> BENCH_DRIVER_EOF; log tag [E2EBench] ->
[TrainerBench]; RenderE2ETest -> RenderBenchTest
README states that the bench is standalone and that LAUNCH_RENDER is the
full launcher's DRY_RUN output. Prose "e2e" meaning the full RL run is kept.
…ts fwd_bwd prepare_fwd_bwd_input(payload) maps a micro-batch to loss inputs and device_puts the traced half onto the shardings fwd_bwd compiles them onto (_batch_data_shardings), leaving the closed-over half untouched; fwd_bwd(PreparedFwdBwdInput) then dispatches without moving data. Tunix's TrainerWorker.stage_fwd_bwd calls it when the RL loop sends the next step's first micro-batch while the device still computes the previous step's last micro-batch and update. Costs one micro-batch of inputs in device memory until it is trained (~20 MB per device for a [64 x 65,536] router-replay batch). Tests: prepared micro-batches train bit-identically (weights and metrics, compiled and eager, no recompile), also when staged through a real Tunix TrainerWorker; a one-token change is seen; preparing changes no engine state.
wait_for_device() blocks until every computation dispatched so far has finished and returns the seconds it waited. Tunix's TrainerWorker calls it after staging the next step's first micro-batch (prepare_fwd_bwd_input): a wait of ~0 means the device was already idle, so the transfer did not hide behind the step's last micro-batch and update, and the RL loop stops staging. It is a blocking wait rather than an is_ready() poll because under Pathways a readiness check is an RPC to the IFRT proxy, so a poll can read "not ready" for work that has finished. InflightThrottler.block_until_ready() does the waiting without dequeuing and without writing metrics, so the next wait_for_next / wait_for_all, and the metrics it writes, are what they would have been. Tests: training with a wait after each prepared micro-batch is bit-identical to training without (weights, metrics, compile count, metric writes in the same order); the wait leaves the in-flight queue as it was and returns with every queued leaf ready; block_until_ready writes nothing and a drain after it writes the same metrics. The review's 4-device test with Tunix's real model input (a nested RLTrainerPayload, batch split four ways) is ported, with the wait added.
… the fused update Brings prepare_fwd_bwd_input / wait_for_device up to the current engine (per-micro-batch dropless replay, the final micro-batch's dropless program, Tunix's fused final update) and makes their tests hold there: - maxtext_engine_prepared_input_dropless_test.py (new): on the one-layer Qwen3.5 MoE rig of the dropless tests, a prepared micro-batch that overflows is replayed from its prepared inputs, the step's last micro-batch runs the final program, and weights, losses, gradient norms and replay counts are bit-identical to the unprepared run, whether it is prepared before or after the update; also when the run's first micro-batch is prepared. A control that prepares the wrong micro-batch is caught. - maxtext_engine_prepared_input_test.py: prepares before and after the update (the second is where Tunix's stage lands once the update runs inside the last train call); the TrainerWorker test drives fwd_bwd(apply_optimizer=True, wait_for_update=False) then stage_fwd_bwd, and skips when the installed Tunix has no staging; wait_for_device is checked to wait on every queued computation, the update's gradient norm included, by recording what it blocked on rather than by is_ready() after the fact. - maxtext_engine_prepared_input_sharded_test.py: re-execs itself with four host devices, as maxtext_engine_data_parallel_test.py does; setting XLA_FLAGS at import time left one device under pytest and skipped green. - README: the two methods and the tests.
…update Tunix's staged step boundary can stage the next step's first micro-batch between a step's last fwd_bwd and its update: fwd_bwd(last, apply_optimizer=True), prepare_fwd_bwd_input(next), wait_for_device(), update(). Under retry_when_tokens_dropped without retry_dropless_final_micro_batch (the MLPerf recipes) the last micro-batch can overflow; its flag is still pending at the prepare and its dropless rerun runs inside update(), after it. maxtext_engine_stage_before_update_test.py, on the one-layer MoE rig of the prepared-input dropless test with the final dropless program off: the staged order trains bit-identically to fwd_bwd(last), update(), fwd_bwd(next payload) (weights, losses, gradient norms, replay counts) with the last micro-batch's rerun inside update() at two boundaries and the prepared micro-batch itself replayed at one; the flag is still pending when the next micro-batch is prepared and wait_for_device resolves nothing; no kernel is rebuilt. Also with the final dropless program on. Two controls are caught: a rerun that runs on the prepared micro-batch instead of its own, and preparing the wrong micro-batch.
…TrainerWorker The same order as the engine-level test, through a real TrainerWorker: fwd_bwd(last, apply_optimizer=True, defer_update=True), stage_fwd_bwd(next), apply_deferred_update(), fwd_bwd_staged(next), against the fused calls. It trains bit-identically (weights, losses, gradient norms) with the last micro-batch's rerun inside the deferred update, defer_update never reaches the engine, and the stage reports no device wait (the update call does). Skipped when the installed Tunix has no apply_deferred_update.
… alone Docstrings and README only, no code change. Tunix stages a step's first micro-batch between the previous step's last fwd_bwd and its update() when it defers the update. Under retry_when_tokens_dropped the last micro-batch's overflow flag is still pending at that point. State the contract this order relies on, which maxtext_engine_stage_before_update_test.py pins: - prepare_fwd_bwd_input neither reads nor resolves the pending check. The dropless rerun that update() may run reads the inputs that check kept, never the prepared ones. The mapping reads only the payload, so the prepared input is the same whether the update runs before or after it. - wait_for_device leaves a pending overflow check, like the in-flight queue and its metrics, to the next fwd_bwd or update(). - README: the same note under "One optimizer step", and the test file in the tests table. - Two test docstrings no longer say the MLPerf recipes set retry_dropless_final_micro_batch, or disagree about it: the team's trainer configuration sets retry_when_tokens_dropped, and the final dropless program is an option on top of it.
MaxTextEngine.prepare_fwd_bwd_input and wait_for_device: the engine hooks that Tunix's staged step boundary and the RL trainer bench's STAGE_BEFORE_UPDATE switch call.
The RL trainer bench under src/maxtext/experimental/rl_trainer_bench that produced the measured step times.
…driver (calls the #5718 engine hooks) Ports onto bench_driver.py the 7 blocks that the driver used by the measured 128-chip runs inserts into #5701's e2e_driver.py at d6db22f (md5 e78c082f -> 4e893e74; insertions only, no changed or removed lines). The result equals that driver except for #5701's 3 renamed lines (log tag, TRAINER_ARGV, ORCH_ARGV). STAGE_BEFORE_UPDATE defaults to off, which runs the original loop.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a standalone benchmark harness for the RL trainer program of a Tunix GRPO end-to-end run, adding tools for rendering, structured diffing, ahead-of-time compilation, and CPU-only validation. It also extends the MaxText training engine to support preparing forward-backward inputs early and waiting for device computations. The code review feedback is highly constructive, pointing out several opportunities to improve robustness and cross-platform compatibility. Key recommendations include using with context managers for file operations to prevent resource leaks, explicitly specifying utf-8 encoding to avoid platform-dependent decoding issues, and defensively handling potential ZeroDivisionError, AttributeError, or KeyError exceptions.
| costs = [ast.literal_eval(m) for m in re.findall(r"Cost analysis: (\{.*\})", s)] | ||
| mems = re.findall(r"Memory analysis: CompiledMemoryStats\(([^)]*)\)", s) | ||
| for name, c, m in zip(("fwd_bwd", "accumulate", "update"), costs, mems): | ||
| d = dict(kv.split("=") for kv in m.split(", ")) |
There was a problem hiding this comment.
Using kv.split('=') can raise a ValueError if any value in the CompiledMemoryStats string contains an = character. Using kv.split('=', 1) is safer and more robust.
Suggested change
| d = dict(kv.split("=") for kv in m.split(", ")) | |
| d = dict(kv.split("=", 1) for kv in m.split(", ")) |
Comment on lines
+25
to
+31
| def load(path): | ||
| out = {} | ||
| for line in open(path, errors="replace"): | ||
| m = re.search(r"Config param ([A-Za-z0-9_]+): (.*)$", line.rstrip("\n")) | ||
| if m: | ||
| out[m.group(1)] = m.group(2).strip() | ||
| return out |
There was a problem hiding this comment.
The file is opened without a with context manager, which can leave the file descriptor open. Additionally, specifying an explicit encoding (like utf-8) prevents platform-dependent decoding issues.
Suggested change
| def load(path): | |
| out = {} | |
| for line in open(path, errors="replace"): | |
| m = re.search(r"Config param ([A-Za-z0-9_]+): (.*)$", line.rstrip("\n")) | |
| if m: | |
| out[m.group(1)] = m.group(2).strip() | |
| return out | |
| def load(path): | |
| out = {} | |
| with open(path, encoding="utf-8", errors="replace") as f: | |
| for line in f: | |
| m = re.search(r"Config param ([A-Za-z0-9_]+): (.*)$", line.rstrip("\n")) | |
| if m: | |
| out[m.group(1)] = m.group(2).strip() | |
| return out |
Comment on lines
+63
to
+79
| def parse(path): | ||
| comps = {} | ||
| cur = None | ||
| for line in open(path): | ||
| if cur is None: | ||
| m = re.match(r"^(ENTRY )?%([\w.\-]+) .*\{\s*$", line) # headers start at column 0, instructions are indented | ||
| if m: | ||
| cur = m.group(2) | ||
| comps[cur] = [] | ||
| continue | ||
| if line.rstrip() == "}": # multi-line backend configs also start lines with "}}" | ||
| cur = None | ||
| continue | ||
| m = re.match(r"^\s*(ROOT )?%([\w.\-]+) = (.*)$", line) | ||
| if m: | ||
| comps[cur].append((m.group(2), m.group(3))) | ||
| return comps |
There was a problem hiding this comment.
The file is opened without a with context manager, which can leave the file descriptor open. It is recommended to use with open(path, encoding='utf-8') as f: to ensure proper resource cleanup and platform-independent decoding.
Suggested change
| def parse(path): | |
| comps = {} | |
| cur = None | |
| for line in open(path): | |
| if cur is None: | |
| m = re.match(r"^(ENTRY )?%([\w.\-]+) .*\{\s*$", line) # headers start at column 0, instructions are indented | |
| if m: | |
| cur = m.group(2) | |
| comps[cur] = [] | |
| continue | |
| if line.rstrip() == "}": # multi-line backend configs also start lines with "}}" | |
| cur = None | |
| continue | |
| m = re.match(r"^\s*(ROOT )?%([\w.\-]+) = (.*)$", line) | |
| if m: | |
| comps[cur].append((m.group(2), m.group(3))) | |
| return comps | |
| def parse(path): | |
| comps = {} | |
| cur = None | |
| with open(path, encoding="utf-8") as f: | |
| for line in f: | |
| if cur is None: | |
| m = re.match(r"^(ENTRY )?%([\w.\-]+) .*\{\s*$", line) # headers start at column 0, instructions are indented | |
| if m: | |
| cur = m.group(2) | |
| comps[cur] = [] | |
| continue | |
| if line.rstrip() == "}": # multi-line backend configs also start lines with "}}" | |
| cur = None | |
| continue | |
| m = re.match(r"^\s*(ROOT )?%([\w.\-]+) = (.*)$", line) | |
| if m: | |
| comps[cur].append((m.group(2), m.group(3))) | |
| return comps |
Comment on lines
+259
to
+264
| balanced = n_tok * top_k / ep | ||
| for l in range(num_layers): | ||
| e = block[:, :, l, :].ravel() | ||
| e = e[e >= 0] | ||
| counts = np.bincount(e // per_shard, minlength=ep) | ||
| best[l] = max(best[l], counts.max() / balanced) |
There was a problem hiding this comment.
If balanced is 0 (which can happen if chunk_len is 0 due to integer division when budget is small or chunks is large), a ZeroDivisionError will be raised. Adding a guard or handling the zero case prevents potential runtime crashes.
Suggested change
| balanced = n_tok * top_k / ep | |
| for l in range(num_layers): | |
| e = block[:, :, l, :].ravel() | |
| e = e[e >= 0] | |
| counts = np.bincount(e // per_shard, minlength=ep) | |
| best[l] = max(best[l], counts.max() / balanced) | |
| balanced = n_tok * top_k / ep | |
| if balanced > 0: | |
| for l in range(num_layers): | |
| e = block[:, :, l, :].ravel() | |
| e = e[e >= 0] | |
| counts = np.bincount(e // per_shard, minlength=ep) | |
| best[l] = max(best[l], counts.max() / balanced) |
| batch_config=batch_config, | ||
| num_generations=o_args.num_generations, | ||
| mini_batch_size=o_args.mini_batch_size, | ||
| train_micro_batch_size=algo.train_micro_batch_size, |
There was a problem hiding this comment.
On line 415, train_micro_batch_size is accessed defensively using getattr(algo, "train_micro_batch_size", 1). However, on line 438, it is accessed directly as algo.train_micro_batch_size. If algo does not have this attribute, it will raise an AttributeError. It is safer to be consistent and use getattr or a local variable.
Suggested change
| train_micro_batch_size=algo.train_micro_batch_size, | |
| train_micro_batch_size=getattr(algo, "train_micro_batch_size", 1), |
| """The JobSet-like documents of a DRY_RUN output: YAML documents split on `---`; chunks that do not parse as a | ||
| mapping (launcher log lines, kubectl output) are skipped.""" | ||
| docs = [] | ||
| with open(path) as f: |
There was a problem hiding this comment.
Opening files without specifying an explicit encoding can lead to platform-dependent decoding errors (e.g., on Windows systems where the default encoding might not be UTF-8). It is recommended to specify encoding="utf-8".
Suggested change
| with open(path) as f: | |
| with open(path, encoding="utf-8") as f: |
| md.setdefault("labels", {})["scheduling.x-k8s.io/gang-id"] = a.name | ||
| js["spec"].setdefault("failurePolicy", {})["maxRestarts"] = 0 | ||
|
|
||
| with open(a.driver) as f: |
There was a problem hiding this comment.
Specify encoding="utf-8" when opening the driver file to ensure consistent behavior across different platforms.
Suggested change
| with open(a.driver) as f: | |
| with open(a.driver, encoding="utf-8") as f: |
| extras = [f for f in env["value"].split() if f.split("=", 1)[0] not in base_keys] | ||
| env["value"] = " ".join(libtpu + extras) | ||
|
|
||
| with open(a.out, "w") as f: |
There was a problem hiding this comment.
Specify encoding="utf-8" when writing the output YAML file to ensure it is written consistently across different platforms.
Suggested change
| with open(a.out, "w") as f: | |
| with open(a.out, "w", encoding="utf-8") as f: |
|
|
||
|
|
||
| def parse(path): | ||
| with open(path) as f: |
There was a problem hiding this comment.
Specify encoding="utf-8" when opening the YAML file to ensure consistent behavior across different platforms.
Suggested change
| with open(path) as f: | |
| with open(path, encoding="utf-8") as f: |
|
|
||
|
|
||
| def summarize(path, cap_gib): | ||
| with open(path, errors="replace") as f: |
There was a problem hiding this comment.
Specify encoding="utf-8" when opening the log file to ensure consistent behavior across different platforms.
Suggested change
| with open(path, errors="replace") as f: | |
| with open(path, encoding="utf-8", errors="replace") as f: |
parambole
changed the title
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. Land #5718 itself.
This branch is
atwigg/mlperfplus #5718 (prepare_fwd_bwd_inputandwait_for_device, the engine hooks for Tunix's staged step boundary), merged atc79bcfd16. Its tree is identical to that merge. That is the only MaxText change the end-to-end run needs: every MaxText key the 128-chip best configuration sets already exists onatwigg/mlperf(num_moe_token_chunks,context_parallel_strategy,sa_fuse_reciprocal,sa_use_base2_exp,gdn_cp_mode,gdn_cp_state,moe_route,gdn_cp_matmul_precision,cast_grads_after_all_reduce); their values are set in the tunix recipe.Companion: google/tunix#2887 (tunix side: #2873 + #2878 + #2877 + the recipe switches).
Measured
On the RL trainer bench (#5701, the measuring tool, not part of this branch), 128 v7x, fsdp 32 x context 4 x expert 2, 64 rows x 65,536 tokens per micro-batch, 9 micro-batches per optimizer step:
Caveats
PIPELINE_STEP_BOUNDARY,STAGE_NEXT_MICROBATCH), which have not run end to end yet.