From 9b706a398188b78ce88a880ce875119de6df8a51 Mon Sep 17 00:00:00 2001
From: Vaibhav Singh <sivaibhav@google.com>
Date: Tue, 13 Oct 2026 18:37:23 +0000
Subject: [PATCH 1/4] Add an RL trainer bench under
 experimental/rl_trainer_bench

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.
---
 .../experimental/rl_trainer_bench/README.md   | 260 +++++++++
 .../experimental/rl_trainer_bench/__init__.py |  14 +
 .../experimental/rl_trainer_bench/aot_e2e.sh  |  46 ++
 .../rl_trainer_bench/aot_summarize.py         |  61 +++
 .../rl_trainer_bench/aot_trainer.py           | 125 +++++
 .../rl_trainer_bench/check_numerics_cpu.py    | 198 +++++++
 .../rl_trainer_bench/check_payload_routing.py | 112 ++++
 .../rl_trainer_bench/diff_config.py           |  43 ++
 .../rl_trainer_bench/e2e_driver.py            | 504 ++++++++++++++++++
 .../hlo_collective_schedule.py                | 142 +++++
 .../rl_trainer_bench/render_e2e.py            | 468 ++++++++++++++++
 .../rl_trainer_bench/render_e2e.sh            |  61 +++
 .../rl_trainer_bench/structdiff_render.py     |  95 ++++
 .../rl_trainer_bench/tests/__init__.py        |  14 +
 .../rl_trainer_bench/tests/payload_test.py    | 218 ++++++++
 .../rl_trainer_bench/tests/render_e2e_test.py | 204 +++++++
 .../tests/testdata/e2e_dryrun_fixture.yaml    | 167 ++++++
 17 files changed, 2732 insertions(+)
 create mode 100644 src/maxtext/experimental/rl_trainer_bench/README.md
 create mode 100644 src/maxtext/experimental/rl_trainer_bench/__init__.py
 create mode 100755 src/maxtext/experimental/rl_trainer_bench/aot_e2e.sh
 create mode 100644 src/maxtext/experimental/rl_trainer_bench/aot_summarize.py
 create mode 100644 src/maxtext/experimental/rl_trainer_bench/aot_trainer.py
 create mode 100644 src/maxtext/experimental/rl_trainer_bench/check_numerics_cpu.py
 create mode 100644 src/maxtext/experimental/rl_trainer_bench/check_payload_routing.py
 create mode 100644 src/maxtext/experimental/rl_trainer_bench/diff_config.py
 create mode 100644 src/maxtext/experimental/rl_trainer_bench/e2e_driver.py
 create mode 100644 src/maxtext/experimental/rl_trainer_bench/hlo_collective_schedule.py
 create mode 100644 src/maxtext/experimental/rl_trainer_bench/render_e2e.py
 create mode 100755 src/maxtext/experimental/rl_trainer_bench/render_e2e.sh
 create mode 100644 src/maxtext/experimental/rl_trainer_bench/structdiff_render.py
 create mode 100644 src/maxtext/experimental/rl_trainer_bench/tests/__init__.py
 create mode 100644 src/maxtext/experimental/rl_trainer_bench/tests/payload_test.py
 create mode 100644 src/maxtext/experimental/rl_trainer_bench/tests/render_e2e_test.py
 create mode 100644 src/maxtext/experimental/rl_trainer_bench/tests/testdata/e2e_dryrun_fixture.yaml

diff --git a/src/maxtext/experimental/rl_trainer_bench/README.md b/src/maxtext/experimental/rl_trainer_bench/README.md
new file mode 100644
index 0000000000..a8d468a70e
--- /dev/null
+++ b/src/maxtext/experimental/rl_trainer_bench/README.md
@@ -0,0 +1,260 @@
+<!--
+ Copyright 2026 Google LLC
+
+ Licensed under the Apache License, Version 2.0 (the "License");
+ you may not use this file except in compliance with the License.
+ You may obtain a copy of the License at
+
+    https://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
+-->
+
+# RL trainer bench
+
+This directory times the RL trainer program of a Tunix GRPO e2e run on its own. The bench runs the exact trainer
+program the e2e run compiles (the same MaxText flags, libtpu flags, mesh, GRPO loss and chunked log-prob pass, hence
+the same compile key) on one trainer slice, under the same Pathways setup, with a synthetic micro-batch whose MoE
+routing load is calibrated to a production run. There is no orchestrator, no rollout engines and no weight sync. A
+variant (a MaxText flag, a libtpu flag, a slice size) can therefore be A/B tested in one short job on one slice, and
+its profile read, without the rest of the RL system.
+
+## Why a separate harness
+
+MaxText's standalone engine script (`python3 -m maxtext.experimental.maxtext_engine.train`) runs MaxText's
+pre-training cross-entropy on MaxText's synthetic loader. That is a different program from the RL trainer: the GRPO
+loss and its vocabulary-chunked `compute_per_token_logps` pass are missing, the batch is laid out differently, and the
+optimizer step is driven by MaxText rather than by Tunix's `TrainerWorker`. Timings from it do not transfer to the RL
+step, and a variant that compiles there may not compile in the RL program (or the other way round).
+
+The bench therefore drives the trainer through Tunix's own code paths. `e2e_driver.py` builds the trainer worker with
+`run_trainer_node._parse_args` and `_create_maxtext_trainer_factory` from the e2e trainer's argv, builds the
+`GRPOAdapter`, `BatchConfig` and batch assembler with `run_deepswe_dist._parse_args` and `_build_algo` from the e2e
+orchestrator's argv, applies `configure_worker`'s loss and model-input functions, compiles against Tunix's
+TRAINER_WARMUP_COMPILE payload, and then runs `STEPS` optimizer steps of `MICRO_STEPS` `fwd_bwd` calls plus one
+`update`, as `DistributedRLEngine.train_step` drives them.
+
+## Requirements
+
+The driver imports Tunix (`tunix.experimental.examples.common.run_trainer_node`,
+`tunix.experimental.examples.deepswe_dist.run_deepswe_dist`, `trainer_warmup`, `batch_assembly`, `trainer_worker`)
+and MaxText, so it must run in an image whose Tunix and MaxText match the e2e run you want to reproduce; by default
+the bench keeps the client image of the e2e render. The renderer and the unit tests need only Python 3 with PyYAML
+and NumPy. Since most of what the driver touches is Tunix code, it may belong in Tunix in the long run; we leave that
+to the team.
+
+## Files
+
+| file | what it does |
+|---|---|
+| `e2e_driver.py` | the bench itself; runs as the Pathways client of the trainer slice |
+| `render_e2e.py`, `render_e2e.sh` | turn an e2e DRY_RUN render into the bench JobSet, with variant knobs |
+| `structdiff_render.py` | structured diff of two bench renders (YAML leaves, MaxText flags, libtpu flags, argv) |
+| `diff_config.py` | diff of the resolved MaxText `Config param` lines of two logs (e2e trainer vs bench) |
+| `aot_trainer.py`, `aot_e2e.sh`, `aot_summarize.py` | ahead-of-time compile and memory report on a CPU host |
+| `check_payload_routing.py` | CPU check of the realistic payload at the 397B production shapes (needs Tunix) |
+| `check_numerics_cpu.py` | CPU check of the numerics knobs with Tunix's GRPO loss (needs Tunix and MaxText) |
+| `hlo_collective_schedule.py` | static overlap summary of the collectives in a scheduled HLO dump |
+| `tests/` | CPU unit tests of the renderer and the payload builder |
+
+## How to use it
+
+### 1. Render your own e2e launch with DRY_RUN
+
+Start from the launch you want to reproduce, with the environment you would launch it with, and ask the team launcher
+for a dry run. For the MLPerf recipe in Tunix this is
+
+```bash
+DRY_RUN=true bash tunix/experimental/examples/recipes/mlperf_gbs1024_submission.sh start > e2e-dryrun.yaml
+```
+
+The output is a stream of JobSets (orchestrator, trainer, one per rollout engine), possibly interleaved with
+launcher log lines. The renderer picks the trainer and the orchestrator JobSet by the `--process_main=` of their
+`main` container and ignores everything else. If you already have the two JobSets as separate files, pass them with
+`--train-yaml` and `--orch-yaml` instead. The DRY_RUN output describes your environment (namespace, images,
+buckets), so keep it local.
+
+### 2. Render the bench JobSet
+
+```bash
+E2E_RENDER=e2e-dryrun.yaml NAME=tb-base K8S_NAMESPACE=<namespace> OUTPUT_PREFIX=<output-prefix> \
+  bash render_e2e.sh                     # writes rendered/tb-base.yaml and nothing else
+```
+
+The render keeps every scheduling field of the e2e trainer JobSet (Kueue labels and annotations, node selectors,
+tolerations, reservations, priority class, service account, the Pathways resource manager, proxy, worker and sidecar
+containers with their images, args and env). It changes the names, sets `failurePolicy.maxRestarts` to 0 (so a
+failure does not rerun the compile), and replaces the client command: the same inline env prefix as the e2e trainer
+(MAXTEXT_EXTRA_FLAGS, LIBTPU_INIT_ARGS and the rest) runs `e2e_driver.py`, written into the pod by the command itself,
+with the e2e trainer argv and orchestrator argv as `E2E_TRAINER_ARGV` and `E2E_ORCH_ARGV`. In the trainer argv,
+`--maxtext_ckpt_path` is cleared (random initialization is the one intended difference), `--maxtext_output_directory`
+becomes `<output-prefix>/<NAME>`, `--worker_id` becomes `NAME`, and the gRPC and discovery flags are dropped. The
+Tunix fail-fast wrapper is dropped as well, since it rewrites a clean exit into a restart.
+
+The knobs, all environment variables of `render_e2e.sh` (the Python script takes the same as flags):
+
+`IMAGE` replaces the client image; the Pathways containers keep the e2e render's. `STEPS` (default 3) is the number of
+optimizer steps and `MICRO_STEPS` the `fwd_bwd` calls per step. Set `MICRO_STEPS` to what your e2e run does per
+optimizer step (the orchestrator logs one train call per micro-batch); the default of 9 (5 with
+`TRAINER_SLICE=4x8x8`) is an example calibration from the 397B recipe, where about 1,152 trajectories per step packed
+into 9 micro-batches of 64 rows on a 4x4x8 trainer.
+
+`EXTRA_MAXTEXT_FLAGS="k=v ..."` appends MaxText flags (MaxText keys are last-wins), and
+`REPLACE_MAXTEXT_FLAGS="k=v ..."` replaces an existing key in place (it fails unless the key is there exactly once),
+which keeps one value per key in the rendered flag set, e.g. `REPLACE_MAXTEXT_FLAGS="context=remat"`.
+`SET_TRAINER_FLAGS` and `SET_ORCH_FLAGS` do the same for the trainer and orchestrator argv, e.g.
+`train_micro_batch_size=128` in both. `COMPUTE_LOGPS_CHUNK_SIZE` sets the trainer's `--compute_logps_chunk_size`.
+
+`LIBTPU_EXTRA="--flag=v ..."` adds or overrides libtpu flags and `LIBTPU_DROP="--flag ..."` removes them by name. Both
+are written to the three places the e2e render carries libtpu flags: the client's LIBTPU_INIT_ARGS, the
+`pathways-proxy` `--xla_*` args (where the programs are compiled) and the `pw-node` workers' LIBTPU_INIT_ARGS (which
+keep their worker-only extras).
+
+`TRAINER_SLICE=<topology>` (e.g. `4x8x8`) re-targets the trainer at another slice of the same accelerator and fills
+it with FSDP, keeping expert, context and tensor parallelism. It changes the resource manager's `--instance_type`,
+the `pw-node` completions and parallelism (hosts = chips / `google.com/tpu` per pod), the slice-topology annotation,
+the trainer's `--mesh_fsdp` and the orchestrator's `--trainer_fsdp`. The last one sets the packed micro-batch rows
+(fsdp x expert), so the number of micro-steps for the same token mass changes too; set `MICRO_STEPS` and
+`FINAL_PAD_FRACTION` to keep the real rows per step comparable.
+
+`PROFILER_SKIP=<n> PROFILER_STEPS=<m>` sets the trainer's `--skip_first_n_profiler_steps` and `--profiler_steps`,
+which Tunix maps to MaxText's `profiler=xplane`, `skip_first_n_steps_for_profiler` and `profiler_steps`. MaxText's
+`training_engine/micro_step_profiler.py` counts `fwd_bwd` calls from 0 over the whole run (the compile does not
+count), opens the trace before call `n` and closes it after call `n + m - 1`. So with `MICRO_STEPS=9`,
+`PROFILER_SKIP=16 PROFILER_STEPS=3` traces calls 16, 17 and 18: the second-to-last and last micro-batch of step 1,
+step 1's `update`, and the first micro-batch of step 2. That is the window to use for the last micro-batch and the
+optimizer update; `STEPS` must be large enough for the window to close. Profiling options (chips and SparseCores to
+trace) are MaxText flags and go in `EXTRA_MAXTEXT_FLAGS`.
+
+`DRIVER_ENV="K=V ..."` passes knobs to the driver; see the next section.
+
+To compare two renders, or a render with the e2e trainer JobSet it came from:
+
+```bash
+python3 structdiff_render.py rendered/tb-base.yaml rendered/tb-variant.yaml
+```
+
+It prints only differences, by YAML leaf and by flag (MaxText flags, libtpu flags, trainer argv, orchestrator argv,
+driver env), so it is quick to confirm that a variant differs in the intended fields only.
+
+### 3. Apply
+
+```bash
+kubectl --context <context> -n <namespace> apply --dry-run=server -f rendered/tb-base.yaml   # optional check
+kubectl --context <context> -n <namespace> apply -f rendered/tb-base.yaml
+kubectl --context <context> -n <namespace> logs -f -c main \
+  -l jobset.sigs.k8s.io/jobset-name=tb-base,jobset.sigs.k8s.io/replicatedjob-name=proc 2>&1 | tee -a tb-base.log
+kubectl --context <context> -n <namespace> delete jobset tb-base                              # when done
+```
+
+### 4. Read the result
+
+The client log carries one line per optimizer step:
+
+```
+completed step: <n>, seconds: <s>, micro_steps: <k>, seconds/micro_step: <s/k>, fwd_bwd dispatch s: [...]
+```
+
+The step time runs from the first `fwd_bwd` dispatch until the step's device work has drained (the driver waits on
+the engine's throttler), and includes the `update`. Step 0 includes first execution, and any step that overlaps the
+profiler window carries the profiler's start and stop cost, so compare steps that are neither. After each step the
+driver also logs `STEP_LOSS` (pooled and per-micro-batch loss and denominators, gradient norm) and `STEP_METRICS`
+(every trainer metric, per-micro-batch ones as lists), read after the step's timing so that they do not affect it.
+
+`python3 diff_config.py <e2e trainer log> <bench log>` compares the resolved MaxText `Config param` lines of the two
+runs and prints the keys that differ; we expect only the run name, the output directory and the checkpoint path to
+differ.
+
+The profile goes under `<output-prefix>/<NAME>/tensorboard/` in the usual `plugins/profile/<timestamp>/` layout of
+`jax.profiler`.
+
+## The payload
+
+Tunix's TRAINER_WARMUP_COMPILE payload fixes the shapes, dtypes, fields and per-row segment count of the packed
+micro-batch, which is what the compile needs. As a timing payload it is degenerate for a MoE model: at the 397B shapes
+it is about 97% padding, every token carries the same id and every routed expert is UNSET, so with router replay every
+token falls back to the router with identical inputs and lands on the same experts.
+
+`PAYLOAD=realistic` (the default) keeps the warm-up payload's structure and rebuilds its contents with
+`realistic_payload`. Real tokens get seeded ids in [256, vocab). Each row's tail is padding, a fraction
+`PAD_FRACTION` (default 0.07) of the row, written as the packer writes it (segment id 0, pad id, zero masks, UNSET
+routing). Real tokens get replayed top-k routing with distinct experts, so MaxText replays them rather than routing
+them. With `ROUTING=probe` (the default) each layer draws from a pool whose max expert-shard load ratio (under the
+EP=8 cp-as-ep layout) is calibrated by bisection to a per-layer target, with a seeded hot shard per layer.
+
+The default targets, `PROBE_REGULAR_MEAN_RATIO` in `e2e_driver.py`, were calibrated from a Qwen3.5-397B-A17B GRPO
+production run (the per-layer mean over its 65 regular micro-batches). For another model or another phase of
+training, pass your own table with `ROUTING_TARGETS="r0,r1,..."` (one ratio per MoE layer, in the same probe order),
+or stretch the default with `ROUTING_SCALE` (targets become `1 + scale * (target - 1)`). `ROUTING=uniform` draws
+uniform top-k instead. `EP_SIZE` (default 8) sets the expert-shard count the calibration assumes.
+
+The step's last micro-batch gets the same treatment plus `FINAL_PAD_FRACTION` (default 0.25) of its rows entirely
+padded, since the e2e run's last micro-batch of a step is a partially filled tail, and its `fwd_bwd` is called with
+`apply_optimizer=True` as in `DistributedRLEngine.train_step`. `PAYLOAD_SEED` seeds the regular payload and seed + 1
+the final one. Before the steps the driver logs one `PAYLOAD_ROUTING` line per payload with the offline per-layer
+ratio of the replayed tokens; the runtime ratio, which includes the padding's router-chosen experts, is in the
+engine's `MOE_DROP_AUDIT` lines with the MaxText flag `log_moe_drops=true`. Building the two payloads at the 397B
+shapes takes a couple of minutes of host time before the first step. `PAYLOAD=warmup` restores the warm-up payload.
+
+## Numerics mode
+
+With the e2e orchestrator flags, the GRPO loss gates sequences on the gap between the rollout log-probs and the
+trainer's (`--seq_logprob_error_threshold`, `--truncated_importance_sampling_*`). The payload's rollout log-probs are
+a synthetic constant, far from what a randomly initialized trainer assigns, so both gates drop every sequence, the
+loss denominator is 0 and the gradient is exactly 0 (with `step_skipped` 0). This does not affect timing, but it hides
+numerics differences. `LOGP_GATES=off` strips the gate flags from the orchestrator argv, so every sequence stays in
+the loss, and `ADVANTAGES=random` gives each segment a seeded N(0, 1) advantage while keeping token ids, segment cuts
+and routing byte-identical. With these two, the gradient norm of two variants (e.g. a MoE padding flag on and off) is
+a meaningful comparison; the loss value itself is a function of the payload alone when the ratio is 1.
+
+`check_numerics_cpu.py --e2e-render e2e-dryrun.yaml` reproduces this on CPU with Tunix's GRPO loss and a
+one-parameter stand-in for the model. The stand-in replaces `compute_per_token_logps`, so it is tied to that
+function's signature in the Tunix version it runs against.
+
+## Ahead-of-time memory check
+
+```bash
+E2E_RENDER=e2e-dryrun.yaml AOT_IMAGE=<image> [TRAINER_SLICE=4x8x8] bash aot_e2e.sh
+python3 aot_summarize.py logs/aot-e2e-<time>.log
+```
+
+`aot_e2e.sh` takes the trainer's MaxText flags, `--xla_*` libtpu flags and mesh from the e2e render
+(`render_e2e.py --print-aot-env`, with the same variant knobs as `render_e2e.sh`) and runs `aot_trainer.py` in the
+image on a CPU host against a virtual topology. `aot_trainer.py` builds the MaxText argv as the Tunix trainer node does
+and calls `training_engine/maxtext_engine_compile.py` with the GRPO loss, packed micro-batches and router replay; its
+recipe constants (sequence lengths, optimizer, GRPO config) are the 397B recipe's and need editing for another
+recipe. The report gives each kernel's temporaries, which is what XLA checks against `--xla_tpu_max_hbm_size_mib`.
+A static compile cannot see runtime fragmentation, staging copies made by weight sync, or kernels that only compile
+at runtime, so treat the result as a screen before spending a slice.
+
+## Tests
+
+```bash
+python3 -m pytest src/maxtext/experimental/rl_trainer_bench/tests
+```
+
+The tests run on CPU with PyYAML and NumPy only: the renderer on a hand-written DRY_RUN fixture
+(`tests/testdata/e2e_dryrun_fixture.yaml`), and the payload builder and metric helpers on a small stand-in payload.
+They live next to the code, so the repository's default `testpaths = tests` does not collect them; run them by path.
+`check_payload_routing.py` and `check_numerics_cpu.py` are the heavier checks at the production shapes and need a
+matching Tunix.
+
+## Caveats
+
+The weights are randomly initialized, so router decisions that the payload does not replay (the padding's, and those
+of fully padded rows) follow random weights rather than trained ones. The payload is synthetic but calibrated: its
+replayed routing follows a measured per-layer load profile, while its segment lengths, prompt and completion split and
+padding are chosen by the knobs above. A variant whose effect depends on the load imbalance of real rollouts needs a
+confirming e2e run.
+
+The profiler counts `fwd_bwd` calls, so its window is in micro-steps; a window that spans an optimizer step boundary
+includes the `update`. After each `update` the bench drains the step and reads the metrics on the host before
+dispatching the next step, so a trace across the boundary holds a host-side gap that the e2e trainer does not have.
+
+Weight sync to the rollout engines and checkpoint saves do not run, so their memory and time are not part of the bench.
+The AOT check is static (see above) and compiles some kernels with CPU fallbacks, so its peak is an estimate best used
+to compare variants with each other.
diff --git a/src/maxtext/experimental/rl_trainer_bench/__init__.py b/src/maxtext/experimental/rl_trainer_bench/__init__.py
new file mode 100644
index 0000000000..f611561ac8
--- /dev/null
+++ b/src/maxtext/experimental/rl_trainer_bench/__init__.py
@@ -0,0 +1,14 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
diff --git a/src/maxtext/experimental/rl_trainer_bench/aot_e2e.sh b/src/maxtext/experimental/rl_trainer_bench/aot_e2e.sh
new file mode 100755
index 0000000000..83d11343a4
--- /dev/null
+++ b/src/maxtext/experimental/rl_trainer_bench/aot_e2e.sh
@@ -0,0 +1,46 @@
+#!/usr/bin/env bash
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# AOT memory check (CPU, no TPU) of the e2e trainer program (Tunix GRPO loss, packed micro-batch with router replay,
+# fwd_bwd / accumulate / update) with aot_trainer.py, inside an image whose MaxText and Tunix match the e2e run.
+#
+#   E2E_RENDER=<file> AOT_IMAGE=<image> [TRAINER_SLICE=4x8x8] [AOT_COMPILE_TOPOLOGY=tpu7x-512]
+#   [EXTRA_MAXTEXT_FLAGS=... REPLACE_MAXTEXT_FLAGS=... LIBTPU_EXTRA=... LIBTPU_DROP=... COMPUTE_LOGPS_CHUNK_SIZE=...]
+#   [LOG=<file>] bash aot_e2e.sh [extra MaxText key=value ...]
+#
+# The MaxText flags, --xla_* libtpu flags and mesh come from the e2e render (render_e2e.py --print-aot-env, with the
+# same variant knobs as render_e2e.sh), so the compile matches what the bench would run. AOT_COMPILE_TOPOLOGY must
+# name the slice: tpu7x-256 for 4x4x8 (the default), tpu7x-512 for 4x8x8. Summarize with aot_summarize.py <log>.
+set -euo pipefail
+HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+: "${E2E_RENDER:?set E2E_RENDER to the DRY_RUN output of your e2e launch}"
+: "${AOT_IMAGE:?set AOT_IMAGE to an image with the MaxText and Tunix of the e2e run}"
+eval "$(python3 "$HERE/render_e2e.py" --print-aot-env --e2e-render "$E2E_RENDER" \
+  --extra-maxtext-flags "${EXTRA_MAXTEXT_FLAGS:-}" --replace-maxtext-flags "${REPLACE_MAXTEXT_FLAGS:-}" \
+  --libtpu-extra="${LIBTPU_EXTRA:-}" --libtpu-drop="${LIBTPU_DROP:-}" --trainer-slice "${TRAINER_SLICE:-}" \
+  --compute-logps-chunk-size "${COMPUTE_LOGPS_CHUNK_SIZE:-}")"
+case "${TRAINER_SLICE:-}" in 4x8x8) _topo=tpu7x-512 ;; *) _topo=tpu7x-256 ;; esac
+TOPO=${AOT_COMPILE_TOPOLOGY:-$_topo}
+LOG=${LOG:-logs/aot-e2e-$(date -u +%Y%m%dT%H%M%SZ).log}
+mkdir -p "$(dirname "$LOG")"
+echo "log: $LOG"
+{ echo "=== $(date -u -Is) aot_e2e TOPO=$TOPO FSDP=$AOT_MESH_FSDP TP=$AOT_MESH_TP EXPERT=$AOT_MESH_EXPERT CONTEXT=$AOT_MESH_CONTEXT IMAGE=$AOT_IMAGE"
+  echo "MAXTEXT_EXTRA_FLAGS=$MAXTEXT_EXTRA_FLAGS"; echo "LIBTPU_INIT_ARGS=$LIBTPU_INIT_ARGS"; } | tee -a "$LOG"
+docker run --rm --entrypoint python -v "$HERE":/bench:ro \
+  -e JAX_PLATFORMS=cpu,tpu -e LIBTPU_INIT_ARGS="$LIBTPU_INIT_ARGS" -e MAXTEXT_EXTRA_FLAGS="$MAXTEXT_EXTRA_FLAGS" \
+  -e AOT_MESH_FSDP="$AOT_MESH_FSDP" -e AOT_MESH_TP="$AOT_MESH_TP" -e AOT_MESH_EXPERT="$AOT_MESH_EXPERT" \
+  -e AOT_MESH_CONTEXT="$AOT_MESH_CONTEXT" -e AOT_COMPILE_TOPOLOGY="$TOPO" \
+  -e AOT_LOGPS_CHUNK_SIZE="${AOT_LOGPS_CHUNK_SIZE:-512}" -e TF_CPP_MIN_LOG_LEVEL=1 \
+  "$AOT_IMAGE" /bench/aot_trainer.py "$@" 2>&1 | tee -a "$LOG"
diff --git a/src/maxtext/experimental/rl_trainer_bench/aot_summarize.py b/src/maxtext/experimental/rl_trainer_bench/aot_summarize.py
new file mode 100644
index 0000000000..92b5f0b5a4
--- /dev/null
+++ b/src/maxtext/experimental/rl_trainer_bench/aot_summarize.py
@@ -0,0 +1,61 @@
+#!/usr/bin/env python3
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Summarizes aot_e2e.sh logs: per kernel (fwd_bwd, accumulate, update) the CompiledMemoryStats and cost analysis,
+and the tool's TRAIN-KERNEL DEVICE PEAK against an HBM cap.
+
+  python3 aot_summarize.py <log> [<log> ...] [--cap-mib 84992]
+"""
+import argparse
+import ast
+import re
+
+
+def summarize(path, cap_gib):
+  with open(path, errors="replace") as f:
+    s = f.read()
+  print("==", path)
+  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(", "))
+    print(
+        f"  {name:10s} flops={c.get('flops', 0):.6e} bytes={c.get('bytes accessed', 0):.6e} "
+        f"temp={int(d['temp_size_in_bytes']):,} arg={int(d['argument_size_in_bytes']):,} "
+        f"out={int(d['output_size_in_bytes']):,} code={int(d['generated_code_size_in_bytes']):,} "
+        f"host_temp={int(d['host_temp_size_in_bytes']):,}"
+    )
+  peak = re.search(r"TRAIN-KERNEL DEVICE PEAK: ([0-9.]+) GiB", s)
+  if peak:
+    p = float(peak.group(1))
+    print(f"  PEAK {p:.2f} GiB vs cap {cap_gib:.2f} GiB -> headroom {cap_gib - p:+.2f} GiB")
+  err = re.findall(r"^(\w*Error[^\n]*|.*RESOURCE_EXHAUSTED[^\n]*)", s, re.M)
+  if err and not peak:
+    print("  FIRST ERROR:", err[0][:400])
+  if "Finished training_engine" not in s and not err:
+    print("  (running or incomplete)")
+
+
+def main():
+  ap = argparse.ArgumentParser()
+  ap.add_argument("logs", nargs="+")
+  ap.add_argument("--cap-mib", type=float, default=84992, help="the trainer's --xla_tpu_max_hbm_size_mib")
+  a = ap.parse_args()
+  for path in a.logs:
+    summarize(path, a.cap_mib / 1024)
+
+
+if __name__ == "__main__":
+  main()
diff --git a/src/maxtext/experimental/rl_trainer_bench/aot_trainer.py b/src/maxtext/experimental/rl_trainer_bench/aot_trainer.py
new file mode 100644
index 0000000000..ef7d98a05f
--- /dev/null
+++ b/src/maxtext/experimental/rl_trainer_bench/aot_trainer.py
@@ -0,0 +1,125 @@
+#!/usr/bin/env python3
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""AOT compile (CPU, no TPU) of the Tunix RL MaxText trainer program for a virtual TPU topology.
+
+Builds the MaxText argv as the Tunix trainer node does (`tunix.utils.maxtext_utils.build_maxtext_config`, with
+MAXTEXT_EXTRA_FLAGS from the environment), appends the `compile_*` keys and calls MaxText's
+`training_engine/maxtext_engine_compile.py`, which compiles fwd_bwd, accumulate and update with Tunix's GRPO loss on a
+sequence-packed RLTrainerPayload with router replay, and reports each kernel's memory.
+
+The recipe constants below are those of the Qwen3.5-397B-A17B GRPO recipe the bench was built for (prompt and
+response lengths, optimizer, GRPO config). The mesh, the logps chunk and the topology come from the environment
+(`render_e2e.py --print-aot-env` derives the first two from an e2e DRY_RUN render); edit the rest for another
+recipe so that the compile matches your trainer. Run with `aot_e2e.sh`, inside an image whose MaxText and Tunix
+match the e2e run.
+
+Environment: MAXTEXT_EXTRA_FLAGS (required), LIBTPU_INIT_ARGS (--xla_* flags), AOT_MESH_FSDP (32), AOT_MESH_TP (1),
+AOT_MESH_EXPERT (2), AOT_MESH_CONTEXT (4), AOT_NUM_DEVICES (fsdp x tp x expert x context), AOT_COMPILE_TOPOLOGY
+(tpu7x-256), AOT_LOGPS_CHUNK_SIZE (512), AOT_TRAIN_MICRO_BATCH_SIZE (64), AOT_GLOBAL_BATCH (1024).
+Extra command-line arguments are appended to the MaxText argv as key=value overrides.
+"""
+
+import os
+import sys
+
+
+def main():
+  extra = sys.argv[1:]
+  os.environ.setdefault("TRAINER_MAXTEXT_ATTENTION", "flash")
+  os.environ.setdefault("ENABLE_PATHWAYS_PERSISTENCE", "1")
+  os.environ.setdefault("CHECKPOINT_ASYNC", "true")
+  assert os.environ.get("MAXTEXT_EXTRA_FLAGS"), "MAXTEXT_EXTRA_FLAGS must be set"
+
+  from maxtext.configs import pyconfig
+  from tunix.utils import maxtext_utils as tunix_mu
+
+  captured = {}
+  real_initialize = pyconfig.initialize
+
+  def capture(argv):
+    captured["argv"] = list(argv)
+    return None
+
+  pyconfig.initialize = capture
+  gbs = int(os.environ.get("AOT_GLOBAL_BATCH", str(64 * 16)))  # prompts x generations per optimizer step
+  micro = int(os.environ.get("AOT_TRAIN_MICRO_BATCH_SIZE", "64"))
+  fsdp, tp, ep, cp = (int(os.environ.get(f"AOT_MESH_{k}", d)) for k, d in
+                      (("FSDP", "32"), ("TP", "1"), ("EXPERT", "2"), ("CONTEXT", "4")))
+  try:
+    tunix_mu.build_maxtext_config(
+        model_name="qwen3.5-397b-a17b",
+        worker_id="aot",
+        train_micro_batch_size=micro,
+        mesh_fsdp=fsdp, mesh_tp=tp, mesh_expert=ep, mesh_context=cp,
+        num_devices=int(os.environ.get("AOT_NUM_DEVICES", str(fsdp * tp * ep * cp))),
+        max_prompt_length=8192, max_response_length=57344,
+        learning_rate=2e-6,
+        warmup_steps_fraction=0.0,
+        load_parameters_path="",
+        padded_moe_mlp_dim=0,
+        base_output_directory="/tmp/aot-maxtext-out",
+        gradient_accumulation_steps=gbs // micro,
+        checkpointing_options=type("C", (), {"save_interval_steps": 0, "max_to_keep": 10})(),
+        rollout_mesh_tp=1,
+        prefuse_moe_weights=True,
+        use_weight_converter=True,
+        max_seq_token_per_tpu=65536,
+        trainable_parameters_mask="^(?!.*routed_experts/gate/kernel).*",
+        base_num_kv_heads=2,
+        attention="flash",
+        remat_policy="custom",
+        learning_rate_final_fraction=1.0,
+        adam_b1=0.9, adam_b2=0.999, adam_eps=1e-8, adam_weight_decay=0.0,
+        gradient_clipping_threshold=0.0625,
+        skip_step_on_spikes=False, skip_step_on_nan=True,
+        float32_gate_logits=True, float32_logits=True,
+    )
+  finally:
+    pyconfig.initialize = real_initialize
+
+  argv = captured["argv"]
+  grpo = (
+      "{num_generations: 16, beta: 0.0, epsilon: 0.2, epsilon_high: 0.28, use_rollout_logps: false,"
+      " exact_token_continuity: true, loss_agg_mode: token-mean, advantage_estimator: grpo-loo,"
+      " overlong_loss_masking: true, seq_logprob_error_threshold: 2.0,"
+      " truncated_importance_sampling_type: seq-mask-tis,"
+      " truncated_importance_sampling_ratio_min: 0.999, truncated_importance_sampling_ratio: 1.002,"
+      " temperature: 1.0}"
+  )
+  argv += [
+      # tpu7x-256 is a 4x4x8 slice (128 chips, 256 devices); tpu7x-512 is 4x8x8 (MaxText accelerator_to_spec_map).
+      f"compile_topology={os.environ.get('AOT_COMPILE_TOPOLOGY', 'tpu7x-256')}",
+      "compile_topology_num_slices=1",
+      "compile_engine_loss=grpo",
+      f"compile_engine_grpo_config={grpo}",
+      # Match the trainer's --compute_logps_chunk_size.
+      f"compile_engine_logps_chunk_size={int(os.environ.get('AOT_LOGPS_CHUNK_SIZE', '512'))}",
+      "compile_engine_max_seq_token_per_tpu=65536",
+      "compile_engine_max_segments_per_packed_row=16",
+      "compile_engine_router_replay=true",
+      "max_prefill_predict_length=8192",
+      "decode_sampling_temperature=1.0",
+  ] + extra
+  print("AOT argv:", argv, flush=True)
+  import maxtext
+  print("maxtext from:", maxtext.__file__, flush=True)
+
+  from maxtext.training_engine import maxtext_engine_compile
+  maxtext_engine_compile.main(argv)
+
+
+if __name__ == "__main__":
+  sys.exit(main())
diff --git a/src/maxtext/experimental/rl_trainer_bench/check_numerics_cpu.py b/src/maxtext/experimental/rl_trainer_bench/check_numerics_cpu.py
new file mode 100644
index 0000000000..d4b56d0001
--- /dev/null
+++ b/src/maxtext/experimental/rl_trainer_bench/check_numerics_cpu.py
@@ -0,0 +1,198 @@
+#!/usr/bin/env python3
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""CPU check of the numerics knobs of e2e_driver.py (no TPU, no cluster): why gradient_norm was 0, and the fix.
+
+Builds the TRAINER_WARMUP_COMPILE payload with Tunix's own builder from the e2e orchestrator argv of a DRY_RUN
+render (--e2e-render; exactly as the driver does: `run_deepswe_dist._parse_args`, `_build_algo`,
+`BatchConfig`, routed_experts left out to keep it light), then the realistic payload, then runs Tunix's GRPO
+loss (`algo.loss_fn()`, the function `configure_worker` hands the trainer) on it with a one-parameter stand-in for
+the model: `compute_per_token_logps` is replaced by `theta + base_logp` per token, so everything after the forward
+pass (gates, masks, aggregation, aux) is the real code and d loss / d theta is the gradient signal.
+
+Checks, per configuration:
+  1. advantages of the warm-up and realistic payloads (ADVANTAGES=inherit|random);
+  2. gates on (production argv), trainer logps far from the payload's synthetic rollout logps (-0.7): every
+     sequence dropped, loss denominator 0, gradient 0 (the engine's `_update_kernel` turns a 0 denominator into
+     zero gradients, hence gradient_norm=0 with step_skipped=0);
+  3. gates on, trainer logps equal to the rollout logps: sequences kept, gradient non-zero (the gate is the cause);
+  4. LOGP_GATES=off (`strip_logp_gates`), far logps: sequences kept, gradient non-zero;
+  5. `summarize_metrics` / `metric_log_lines` on a hand-built MaxText MetricsBuffer with a zero-denominator entry.
+With --fsdp 64 it also reports the packed rows per micro-batch of the 256-chip variant.
+
+Run with a Tunix and MaxText that match the e2e run, e.g. inside the trainer image:
+  docker run --rm --entrypoint python -v $PWD:/tb:ro -v <dir of render>:/r:ro -e JAX_PLATFORMS=cpu <image> \
+    /tb/check_numerics_cpu.py --e2e-render /r/<e2e-dryrun>.yaml
+"""
+
+import argparse
+import dataclasses
+import os
+import re
+import shlex
+import sys
+
+import numpy as np
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, HERE)
+import e2e_driver as d  # pylint: disable=g-import-not-at-top
+
+
+def orch_argv_from_render(path, fsdp=None):
+  from render_e2e import _proc_main, load_documents, select_jobsets, split_command  # pylint: disable=g-import-not-at-top
+  _, orch = select_jobsets(load_documents(path))
+  _, flags = split_command(_proc_main(orch)[1]["command"][-1])
+  if fsdp is not None:
+    flags = [t for t in flags if not t.startswith("--trainer_fsdp=")] + [f"--trainer_fsdp={fsdp}"]
+  return flags
+
+
+def build(orch_argv):
+  from tunix.experimental.examples.deepswe_dist import run_deepswe_dist  # pylint: disable=g-import-not-at-top
+  from tunix.experimental.examples.deepswe_dist import trainer_warmup  # pylint: disable=g-import-not-at-top
+  from tunix.experimental.orchestrator import batch_assembly  # pylint: disable=g-import-not-at-top
+  o = run_deepswe_dist._parse_args(orch_argv)  # pylint: disable=protected-access
+  algo = run_deepswe_dist._build_algo(o)  # pylint: disable=protected-access
+  bc = batch_assembly.BatchConfig(
+      pad_id=0, max_prompt_length=o.max_prompt_length, max_response_length=o.max_response_length,
+      max_seq_token_per_tpu=o.max_seq_token_per_tpu, max_segments_per_packed_row=o.max_segments_per_packed_row,
+      trainer_fsdp=o.trainer_fsdp, trainer_dp=o.trainer_dp, trainer_expert=o.trainer_expert)
+  warm = trainer_warmup.build_warmup_payload(
+      algo=algo, batch_config=bc, num_generations=o.num_generations, mini_batch_size=o.mini_batch_size,
+      train_micro_batch_size=algo.train_micro_batch_size, routed_experts_shape=None)
+  return o, algo, warm
+
+
+def adv_stats(p):
+  a = np.asarray(p.advantages)[np.asarray(p.completion_mask) > 0]
+  return dict(scored=int(a.size), nonzero_frac=round(float((a != 0).mean()), 4), uniq=np.unique(a)[:6].tolist()
+              if np.unique(a).size <= 6 else f"{np.unique(a).size} distinct", mean=round(float(a.mean()), 5),
+              std=round(float(a.std()), 5))
+
+
+def run_loss(algo, payload, base_logp):
+  """(loss value, loss denominator, d loss/d theta, kept_frac) of Tunix's GRPO loss with logps = theta + base."""
+  import jax  # pylint: disable=g-import-not-at-top
+  import jax.numpy as jnp  # pylint: disable=g-import-not-at-top
+  from flax import nnx  # pylint: disable=g-import-not-at-top
+  from tunix.rl import algo_core  # pylint: disable=g-import-not-at-top
+
+  class Stand(nnx.Module):
+    def __init__(self):
+      self.theta = nnx.Param(jnp.zeros((), jnp.float32))
+
+  def fake_logps(graphdef, state, *, completion_tokens, return_moe_audit=False, **_):
+    theta = nnx.merge(graphdef, state).theta.value
+    lp = theta + jnp.full(completion_tokens.shape, base_logp, jnp.float32)
+    out = (lp, jnp.zeros_like(lp), jnp.array(False))
+    return out + (None,) if return_moe_audit else out
+
+  orig = algo_core.common.compute_per_token_logps
+  algo_core.common.compute_per_token_logps = fake_logps
+  try:
+    ex = jax.tree.map(lambda x: jnp.asarray(x) if isinstance(x, np.ndarray) else x, payload)
+    gin = algo.build_gen_model_input_fn(pad_id=0, eos_id=0)(ex)
+    loss_fn = algo.loss_fn()
+    gdef, st = nnx.split(Stand())
+
+    def f(s):
+      out = loss_fn(nnx.merge(gdef, s), **gin)
+      return out.primary_loss.unreduced_sum, out
+    (_, out), g = jax.value_and_grad(f, has_aux=True)(st)
+  finally:
+    algo_core.common.compute_per_token_logps = orig
+  kept = out.aux_metrics.get("sample_mask/kept_frac")
+  return (float(out.primary_loss.compute()), float(out.primary_loss.denominator), float(g["theta"].value),
+          float(kept.compute()) if kept is not None else None)
+
+
+def check_summary():
+  import jax.numpy as jnp  # pylint: disable=g-import-not-at-top
+  from maxtext.training_engine import abstract_engine  # pylint: disable=g-import-not-at-top
+  wm = abstract_engine.WeightedMetric
+  buf = abstract_engine.MetricsBuffer(
+      id=0,
+      weighted_metrics={"loss": wm(jnp.array([2.0, 0.0, -3.0]), jnp.array([4.0, 0.0, 6.0])),
+                        "sample_mask/kept_frac": wm(jnp.array([1.0, 0.0, 2.0]), jnp.array([1.0, 0.0, 2.0]),
+                                                    min_denom=1.0)},
+      scalar_metrics={"gradient_norm": jnp.array([0.125]), "is_ratio/mean": jnp.array([1.0, 1.0, 1.0])})
+  s = d.summarize_metrics(buf)
+  assert abs(s["loss"] - (-1.0 / 10.0)) < 1e-12, s
+  assert s["micro_denominator"] == [4.0, 0.0, 6.0] and np.isnan(s["micro_loss"][1]), s
+  assert abs(s["micro_loss"][0] - 0.5) < 1e-12 and abs(s["micro_loss"][2] + 0.5) < 1e-12, s
+  assert s["scalar"]["gradient_norm"] == 0.125 and s["scalar"]["is_ratio/mean"] == [1.0, 1.0, 1.0], s
+  for line in d.metric_log_lines(0, s):
+    print(line)
+  return True
+
+
+def main():
+  ap = argparse.ArgumentParser()
+  ap.add_argument("--e2e-render", required=True, help="DRY_RUN output of the e2e launcher (see README)")
+  ap.add_argument("--fsdp", type=int, default=None, help="override --trainer_fsdp (64 for the 256-chip variant)")
+  ap.add_argument("--rows", type=int, default=4, help="rows of the micro-batch fed to the loss (CPU time)")
+  ap.add_argument("--far-logp", type=float, default=-12.0,
+                  help="stand-in trainer logp per token; ~ -ln(vocab) for a near-uniform random-init model")
+  a = ap.parse_args()
+
+  argv = orch_argv_from_render(a.e2e_render, a.fsdp)
+  gates = [t for t in argv if t.split("=", 1)[0] in d.LOGP_GATE_FLAGS]
+  print("orch gate flags:", gates)
+  o, algo_on, warm = build(argv)
+  rows, budget = np.asarray(warm.completion_ids).shape
+  print(f"trainer_fsdp={o.trainer_fsdp} trainer_expert={o.trainer_expert}: warm-up micro-batch rows={rows} "
+        f"budget={budget} segments/row={[len(np.unique(r[r > 0])) for r in np.asarray(warm.segment_ids)[:4]]}")
+  print("warm-up advantages:", adv_stats(warm))
+  print("warm-up overlong:", None if warm.overlong is None else np.unique(np.asarray(warm.overlong)).tolist(),
+        "| old_per_token_logps is None:", warm.old_per_token_logps is None,
+        "| rollout_per_token_logps unique:", np.unique(np.asarray(warm.rollout_per_token_logps)).tolist()[:4])
+
+  kw = dict(num_experts=512, ep=8, pad_id=0, vocab_size=248_000, pad_fraction=0.07, seed=0)
+  inh = d.realistic_payload(warm, **kw)
+  rnd = d.realistic_payload(warm, advantages="random", **kw)
+  for f in ("completion_ids", "segment_ids", "segment_positions", "completion_mask", "rollout_per_token_logps"):
+    assert np.array_equal(np.asarray(getattr(inh, f)), np.asarray(getattr(rnd, f))), f
+  print("realistic ADVANTAGES=inherit:", adv_stats(inh))
+  print("realistic ADVANTAGES=random: ", adv_stats(rnd), "(ids/segments/masks/rollout logps identical to inherit)")
+  print("realistic rollout_per_token_logps on scored tokens:",
+        np.unique(np.asarray(rnd.rollout_per_token_logps)[np.asarray(rnd.completion_mask) > 0]).tolist())
+
+  _, algo_off, _ = build(d.strip_logp_gates(argv))
+  print("gates off: algo_config seq_logprob_error_threshold=%s truncated_importance_sampling_type=%s" % (
+      algo_off.algo_config.seq_logprob_error_threshold, algo_off.algo_config.truncated_importance_sampling_type))
+
+  small = dataclasses.replace(rnd, **{f.name: np.asarray(getattr(rnd, f.name))[:a.rows]
+                                      for f in dataclasses.fields(rnd)
+                                      if isinstance(getattr(rnd, f.name), np.ndarray)
+                                      and np.asarray(getattr(rnd, f.name)).shape[:1] == (rows,)})
+  results = {}
+  for label, algo, base in (("gates on, trainer logp %.1f" % a.far_logp, algo_on, a.far_logp),
+                            ("gates on, trainer logp -0.7 (= rollout)", algo_on, -0.7),
+                            ("gates off, trainer logp %.1f" % a.far_logp, algo_off, a.far_logp)):
+    loss, den, grad, kept = run_loss(algo, small, base)
+    results[label] = (loss, den, grad, kept)
+    print(f"GRPO loss [{label}] rows={a.rows}: loss={loss:.6g} denominator={den:.6g} dloss/dtheta={grad:.6g} "
+          f"kept_frac={kept}")
+  on_far, on_eq, off_far = results.values()
+  assert on_far[1] == 0.0 and on_far[2] == 0.0, on_far
+  assert on_eq[1] > 0 and on_eq[2] != 0.0, on_eq
+  assert off_far[1] > 0 and off_far[2] != 0.0, off_far
+  check_summary()
+  print("NUMERICS_CPU_CHECK_PASSED")
+
+
+if __name__ == "__main__":
+  main()
diff --git a/src/maxtext/experimental/rl_trainer_bench/check_payload_routing.py b/src/maxtext/experimental/rl_trainer_bench/check_payload_routing.py
new file mode 100644
index 0000000000..ed2f07c4d6
--- /dev/null
+++ b/src/maxtext/experimental/rl_trainer_bench/check_payload_routing.py
@@ -0,0 +1,112 @@
+#!/usr/bin/env python3
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""CPU check of e2e_driver.realistic_payload at the 397B production shapes (no TPU, no cluster; needs Tunix).
+
+Builds the TRAINER_WARMUP_COMPILE payload with Tunix's own builder at the 397B e2e run's BatchConfig (64 rows x 65,536
+tokens, 16 segments per row at most, trainer_fsdp=32, trainer_expert=2, routed_experts [60 layers, top-10]), then
+the realistic regular and final micro-batches, and checks them under the cp-as-ep EP=8 layout
+(`payload_shard_load_ratio`: EP group = 2 rows x 4 context slices, 4 token chunks per slice, 512 experts in 8
+shards of 64):
+
+  * shapes, dtypes and fields equal the warm-up payload's; padding has pad ids and UNSET routing; every real
+    token has top-k distinct experts in range (so MaxText replays it rather than falling back to the router);
+  * regular, routing=probe: each layer's max shard-load ratio is within 0.10 of its target in the default
+    calibration table, the max over layers is in that run's regular range [1.8, 2.45], and the hottest layer is
+    the table's (model layer 59);
+  * ROUTING_SCALE stretches it: scale 1.46 puts the max at about 2.45, the highest regular micro-batch of that run;
+  * routing=uniform: max ratio below 1.2;
+  * final: the pad fraction is FINAL_PAD_FRACTION of whole rows plus PAD_FRACTION of the rest.
+
+Run with a Tunix that matches the e2e run, e.g. inside the trainer image:
+  docker run --rm --entrypoint python -v $PWD:/tb:ro -e JAX_PLATFORMS=cpu <image> /tb/check_payload_routing.py
+The shapes and thresholds are the 397B recipe's; edit the constants below for another model.
+"""
+
+import os
+import sys
+import time
+
+import numpy as np
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+import e2e_driver as d  # pylint: disable=g-import-not-at-top
+
+from tunix.experimental.common import datatypes  # pylint: disable=g-import-not-at-top
+from tunix.experimental.examples.deepswe_dist import trainer_warmup  # pylint: disable=g-import-not-at-top
+from tunix.experimental.orchestrator import algorithm_adapter  # pylint: disable=g-import-not-at-top
+from tunix.experimental.orchestrator import batch_assembly  # pylint: disable=g-import-not-at-top
+from tunix.rl import algorithm_config  # pylint: disable=g-import-not-at-top
+
+NUM_EXPERTS, EP, LAYERS, TOP_K = 512, 8, 60, 10
+
+
+def main():
+  assert datatypes.UNSET_ROUTED_EXPERT == d.UNSET_ROUTED_EXPERT
+  algo = algorithm_adapter.GRPOAdapter(
+      algo_config=algorithm_config.GRPOConfig(num_generations=16, beta=0.0, temperature=1.0),
+      mini_batch_size=72, train_micro_batch_size=64, max_packed_len=65536, max_response_length=57344)
+  bc = batch_assembly.BatchConfig(pad_id=0, max_prompt_length=8192, max_response_length=57344,
+                                  max_seq_token_per_tpu=65536, max_segments_per_packed_row=16,
+                                  trainer_fsdp=32, trainer_dp=1, trainer_expert=2)
+  t = time.monotonic()
+  warm = trainer_warmup.build_warmup_payload(algo=algo, batch_config=bc, num_generations=16, mini_batch_size=72,
+                                             train_micro_batch_size=64, routed_experts_shape=(LAYERS, TOP_K))
+  print(f"warm-up payload {time.monotonic() - t:.1f}s: {batch_assembly.describe_payload(warm)}", flush=True)
+  w_ratio, w_pad = d.payload_shard_load_ratio(warm, num_experts=NUM_EXPERTS, ep=EP)
+  print(f"warm-up: pad {w_pad:.3f}, replayed max ratio {w_ratio.max():.3f} (all routing UNSET: the router decides)")
+
+  kw = dict(num_experts=NUM_EXPERTS, ep=EP, pad_id=0, vocab_size=248_000)
+  targets = d.layer_targets(LAYERS)
+  results = {}
+  for name, extra in (("regular", {}), ("scale1.46", {"scale": 1.46}), ("uniform", {"routing": "uniform"}),
+                      ("final", {"pad_rows_fraction": 0.25, "seed": 1})):
+    t = time.monotonic()
+    p = d.realistic_payload(warm, **kw, **extra)
+    gen = time.monotonic() - t
+    ratio, pad = d.payload_shard_load_ratio(p, num_experts=NUM_EXPERTS, ep=EP)
+    top = np.argsort(-ratio)[:5]
+    print(f"{name}: built {gen:.1f}s, pad {pad:.4f}, max ratio {ratio.max():.3f}, hottest "
+          f"{[(int(l), round(float(ratio[l]), 3)) for l in top]}, min {ratio.min():.3f}", flush=True)
+    results[name] = (p, ratio, pad)
+
+  p, ratio, pad = results["regular"]
+  for f in ("completion_ids", "completion_mask", "advantages", "segment_ids", "segment_positions",
+            "old_per_token_logps", "rollout_per_token_logps", "overlong", "routed_experts"):
+    a, b = np.asarray(getattr(p, f)), np.asarray(getattr(warm, f))
+    assert a.shape == b.shape and a.dtype == b.dtype, (f, a.shape, b.shape, a.dtype, b.dtype)
+  assert p.num_segments == warm.num_segments
+  routed = np.asarray(p.routed_experts)
+  real = np.asarray(p.segment_ids) > 0
+  assert (routed[~real] == d.UNSET_ROUTED_EXPERT).all() and (np.asarray(p.completion_ids)[~real] == 0).all()
+  r = routed[real]  # [n, L, K]
+  assert r.min() >= 0 and r.max() < NUM_EXPERTS
+  srt = np.sort(r[:200_000], axis=-1)
+  assert not (srt[..., 1:] == srt[..., :-1]).any(), "duplicate experts would fall back to the router"
+  assert len(np.unique(np.asarray(p.completion_ids)[real][:100_000])) > 50_000, "token ids not varied"
+  assert abs(pad - 0.07) < 0.005, pad
+  dev = np.abs(ratio - targets)
+  print(f"regular: |ratio - target| max {dev.max():.3f}, mean {dev.mean():.3f}")
+  assert dev.max() < 0.10, dev.max()
+  assert 1.8 <= ratio.max() <= 2.45 and int(np.argmax(ratio)) == 59, (ratio.max(), np.argmax(ratio))
+  assert 2.35 <= results["scale1.46"][1].max() <= 2.55, results["scale1.46"][1].max()
+  assert results["uniform"][1].max() < 1.2, results["uniform"][1].max()
+  fpad = results["final"][2]
+  assert abs(fpad - (0.25 + 0.75 * 0.07)) < 0.01, fpad
+  print("PAYLOAD_ROUTING_CHECK_PASSED", flush=True)
+
+
+if __name__ == "__main__":
+  main()
diff --git a/src/maxtext/experimental/rl_trainer_bench/diff_config.py b/src/maxtext/experimental/rl_trainer_bench/diff_config.py
new file mode 100644
index 0000000000..f55e4271b7
--- /dev/null
+++ b/src/maxtext/experimental/rl_trainer_bench/diff_config.py
@@ -0,0 +1,43 @@
+#!/usr/bin/env python3
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Diffs the resolved MaxText "Config param" dumps of two logs (e.g. e2e trainer vs harness client).
+
+  python3 diff_config.py <e2e_cfg_or_log> <bench_log>
+Lines are matched on "Config param <key>: <value>" anywhere in the line; prints only differing keys.
+"""
+import re
+import sys
+
+
+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 main():
+  a, b = load(sys.argv[1]), load(sys.argv[2])
+  print(f"{len(a)} keys in {sys.argv[1]}, {len(b)} keys in {sys.argv[2]}")
+  for k in sorted(set(a) | set(b)):
+    if a.get(k) != b.get(k):
+      print(f"{k}\n  e2e:   {a.get(k)}\n  bench: {b.get(k)}")
+
+
+if __name__ == "__main__":
+  main()
diff --git a/src/maxtext/experimental/rl_trainer_bench/e2e_driver.py b/src/maxtext/experimental/rl_trainer_bench/e2e_driver.py
new file mode 100644
index 0000000000..cb5d828954
--- /dev/null
+++ b/src/maxtext/experimental/rl_trainer_bench/e2e_driver.py
@@ -0,0 +1,504 @@
+#!/usr/bin/env python3
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Standalone driver for the e2e RL trainer program (no orchestrator, no rollout engines).
+
+Runs as the Pathways client in the trainer pod and drives the same objects the e2e run drives, through the
+Tunix code paths themselves:
+
+  * trainer: `run_trainer_node._parse_args` + `_create_maxtext_trainer_factory` + `TrainerWorker`, i.e. the
+    MaxText training engine wrapped in the Tunix adapter, built from the e2e trainer's own argv (and the
+    MAXTEXT_EXTRA_FLAGS env it reads);
+  * algorithm and inputs: `run_deepswe_dist._parse_args` + `_build_algo` (GRPOAdapter/GRPOConfig) from the e2e
+    orchestrator's own argv, the program's `BatchConfig` and `create_batch_assembler` as `StandardRLProgram`
+    builds them, and the TRAINER_WARMUP_COMPILE payload (`trainer_warmup.build_warmup_payload`), which runs
+    synthetic trajectories through `GRPOAdapter.create_trainer_payloads` and the assembler;
+  * configure_worker: the body of `DistributedRLEngine.configure_worker` for the ACTOR role (loss fn with
+    has_aux=True, gen_model_input_fn with the assembler's pad/eos ids), then `TrainerWorker.compile` as
+    `RLProgram.run_async` does for the warm-up;
+  * steps: STEPS optimizer steps of MICRO_STEPS `fwd_bwd` calls on the payload, then `update`, as
+    `DistributedRLEngine.train_step` drives them.
+
+Payload (PAYLOAD=realistic, the default): the warm-up payload only fixes shapes, dtypes and the per-row segment
+count; the micro-batches the steps run are rebuilt from it by `realistic_payload` below, with seeded varied token
+ids, replayed routing whose per-layer max expert-shard load ratio follows a calibration table
+(PROBE_REGULAR_MEAN_RATIO, measured on a 397B production run; ROUTING_TARGETS overrides it), and tail padding
+(identical pad ids, UNSET routing, as the packer writes it). The step's last micro-batch additionally
+has FINAL_PAD_FRACTION of its rows fully padded, which reproduces the final micro-batch's load spike, and is called
+with apply_optimizer=True as the orchestrator does. PAYLOAD=warmup restores the old behaviour (every token the same
+id, every routed expert UNSET, so the router sends every token to the same experts: a degenerate MoE load).
+
+Numerics knobs (moe_exclude_padding_from_routing on/off A/B, see README "Numerics check"): ADVANTAGES=random gives
+every segment a seeded N(0, 1) advantage (token ids and routing unchanged); LOGP_GATES=off strips the orchestrator's
+seq_logprob_error_threshold / truncated_importance_sampling_* flags. With them on, the payload's synthetic rollout
+log-probs (-0.7) are far from a random-init trainer's, both gates drop every sequence and the gradient is exactly 0.
+After each step the driver logs `STEP_LOSS step=N loss= grad_norm= micro_loss=[..] micro_denominator=[..]` and
+`STEP_METRICS step=N ...` (every trainer metric; per-micro-batch ones as lists).
+
+The intended difference from e2e: no checkpoint is loaded (random init; the trainer argv's --maxtext_ckpt_path is
+cleared by the launcher), so weights and routing values differ while shapes, shardings and the program match.
+"""
+
+import logging
+import os
+import shlex
+import sys
+import time
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - [E2EBench] %(message)s", force=True)
+
+import dataclasses
+
+import numpy as np
+
+UNSET_ROUTED_EXPERT = -1  # tunix.experimental.common.datatypes.UNSET_ROUTED_EXPERT
+
+# Default routing calibration. Per-layer max expert-shard load ratio (EP=8) of the regular (non-final) micro-batches
+# of a Qwen3.5-397B-A17B GRPO production run, mean over 65 regular micro-batches, read from the engine's per-layer
+# load-ratio log lines (per micro-batch the max over layers ranged 1.71-2.45, median 2.00). Probe order:
+# position-major over the scanned 4-layer blocks, i.e. index k is model layer 4 * (k % 15) + k // 15 (see
+# `probe_index_to_layer`). For another model or a later training phase, pass your own table in the same order with
+# ROUTING_TARGETS="r0,r1,..." (one ratio per layer), or stretch this one with ROUTING_SCALE.
+PROBE_REGULAR_MEAN_RATIO = (
+    1.721, 1.499, 1.276, 1.515, 1.541, 1.797, 1.466, 1.607, 1.743, 1.341, 1.504, 1.917, 1.570, 1.506, 1.450,
+    1.269, 1.823, 1.241, 1.496, 1.417, 1.388, 1.316, 1.665, 1.820, 1.353, 1.576, 1.336, 1.926, 1.355, 1.532,
+    1.773, 1.518, 1.224, 1.733, 1.507, 1.250, 1.579, 1.466, 1.376, 1.507, 1.558, 1.309, 1.534, 1.733, 1.543,
+    1.280, 1.519, 1.475, 1.497, 1.262, 1.511, 1.463, 1.698, 1.694, 1.746, 1.664, 1.407, 1.524, 1.721, 1.995,
+)
+
+
+def probe_index_to_layer(k, num_layers=60, cycle=4):
+  """Model layer of index k of the probe's per-layer vector (position-major over `cycle`-layer scanned blocks)."""
+  blocks = num_layers // cycle
+  return cycle * (k % blocks) + k // blocks
+
+
+def parse_routing_targets(text, num_layers):
+  """ROUTING_TARGETS="r0,r1,...": one max shard-load ratio per layer in probe order; validated against the depth."""
+  vals = [float(x) for x in text.replace(" ", "").split(",") if x]
+  if len(vals) != num_layers:
+    raise ValueError(f"ROUTING_TARGETS has {len(vals)} values, the model has {num_layers} MoE layers")
+  if min(vals) < 1.0:
+    raise ValueError("ROUTING_TARGETS values are max/balanced load ratios and must be >= 1.0")
+  return tuple(vals)
+
+
+def layer_targets(num_layers, scale=1.0, cycle=4, probe=None):
+  """Per model layer target max shard-load ratio: the probe table (default PROBE_REGULAR_MEAN_RATIO, in probe
+  order), stretched by `scale` around 1.0."""
+  probe = np.asarray(PROBE_REGULAR_MEAN_RATIO if probe is None else probe, dtype=np.float64)
+  if num_layers != probe.size:
+    # Another depth: cycle the probe's model-ordered profile.
+    by_layer = np.empty_like(probe)
+    by_layer[[probe_index_to_layer(k, probe.size, cycle) for k in range(probe.size)]] = probe
+    by_layer = np.resize(by_layer, num_layers)
+  else:
+    by_layer = np.empty_like(probe)
+    by_layer[[probe_index_to_layer(k, num_layers, cycle) for k in range(num_layers)]] = probe
+  return 1.0 + scale * (by_layer - 1.0)
+
+
+def _gumbel_topk(log_w, n, k, rng):
+  """n draws of k distinct experts, without replacement, with probability proportional to exp(log_w)."""
+  g = rng.gumbel(size=(n, log_w.size)) + log_w[None, :]
+  return np.argpartition(-g, k - 1, axis=1)[:, :k]
+
+
+def _pool_ratio(pool, num_experts, ep):
+  """Max over expert shards of the pool's assignments to the shard, over the balanced share."""
+  counts = np.bincount((pool // (num_experts // ep)).ravel(), minlength=ep)
+  return counts.max() * ep / pool.size
+
+
+def routing_pool(target, num_experts, ep, top_k, pool_size, rng, hot_shard):
+  """`pool_size` distinct-expert top-k rows whose max shard-load ratio is ~`target`: `hot_shard`'s experts get a
+  weight multiplier found by bisection, the other experts are uniform."""
+  per_shard = num_experts // ep
+  hot = np.zeros(num_experts, dtype=bool)
+  hot[hot_shard * per_shard:(hot_shard + 1) * per_shard] = True
+  def draw(log_m, n):
+    sub = np.random.default_rng(rng.integers(1 << 62))
+    return _gumbel_topk(np.where(hot, log_m, 0.0), n, top_k, sub)
+  lo, hi = 0.0, 12.0
+  if target <= 1.0:
+    return draw(0.0, pool_size)
+  calib = min(pool_size, 4096)
+  for _ in range(20):
+    mid = 0.5 * (lo + hi)
+    if _pool_ratio(draw(mid, calib), num_experts, ep) < target:
+      lo = mid
+    else:
+      hi = mid
+  return draw(0.5 * (lo + hi), pool_size)
+
+
+def realistic_payload(base, *, num_experts, ep=8, routing="probe", scale=1.0, pad_fraction=0.07,
+                      pad_rows_fraction=0.0, pad_id=0, vocab_size=151_000, seed=0, pool_size=8192, cycle=4,
+                      layer_targets_override=None, advantages="inherit", probe_targets=None):
+  """`base` (the warm-up payload) with realistic tokens, routing and padding; same shapes, dtypes, fields.
+
+  Per row: the real tokens fill the first (1 - pad_fraction) of the row, split into the row's own number of
+  segments in `base` (each segment a prompt part, completion_mask 0, then a completion part, mask 1, in the
+  prompt/completion proportion of `base`); the tail is padding: segment id 0, pad_id, every mask 0 and UNSET
+  routing, as the packer writes it. `pad_rows_fraction` of the rows (the last ones) are entirely padding.
+  Real tokens get seeded ids in [256, vocab_size) and replayed routing: routing="probe" draws each (token, layer)
+  from a per-layer pool calibrated to the probe's per-layer max shard-load ratio (scaled by `scale`), with a
+  seeded hot shard per layer (`probe_targets` replaces the default table, same order); routing="uniform" draws
+  top-k uniformly without replacement.
+  Advantages are one value per segment, constant over its completion tokens (as `create_trainer_payloads` writes a
+  trajectory's advantage): advantages="inherit" draws it from the warm-up payload's non-zero values (+-1 if it has
+  none); advantages="random" draws it standard normal from a separate generator seeded [seed, 1], so token ids,
+  segment cuts and routing stay byte-identical to the "inherit" payload of the same seed.
+  """
+  if advantages not in ("inherit", "random"):
+    raise ValueError(f"advantages must be inherit or random, got {advantages!r}")
+  adv_rng = np.random.default_rng([seed, 1])
+  rng = np.random.default_rng(seed)
+  ids0 = np.asarray(base.completion_ids)
+  seg0 = np.asarray(base.segment_ids)
+  cm0 = np.asarray(base.completion_mask)
+  adv0 = np.asarray(base.advantages)
+  routed0 = np.asarray(base.routed_experts) if base.routed_experts is not None else None
+  rows, budget = ids0.shape
+  real0 = seg0 > 0
+  prompt_frac = 1.0 - (cm0[real0].mean() if real0.any() else 0.5)
+  adv_vals = np.unique(adv0[adv0 != 0]) if (adv0 != 0).any() else np.array([1.0, -1.0])
+
+  ids = np.full_like(ids0, pad_id)
+  seg = np.zeros_like(seg0)
+  pos = np.zeros_like(np.asarray(base.segment_positions))
+  cm = np.zeros_like(cm0)
+  adv = np.zeros_like(adv0)
+  n_pad_rows = int(round(pad_rows_fraction * rows))
+  real_len = int(round((1.0 - pad_fraction) * budget))
+  real_mask = np.zeros((rows, budget), dtype=bool)
+  for r in range(rows - n_pad_rows):
+    n_seg = max(1, len(np.unique(seg0[r][seg0[r] > 0])))
+    cuts = np.sort(rng.choice(np.arange(64, real_len, 64), size=n_seg - 1, replace=False)) if n_seg > 1 else []
+    bounds = [0, *cuts, real_len]
+    for s in range(n_seg):
+      a, b = bounds[s], bounds[s + 1]
+      seg[r, a:b] = s + 1
+      pos[r, a:b] = np.arange(b - a)
+      p_end = a + int(prompt_frac * (b - a))
+      cm[r, p_end:b] = 1
+      a_val = adv_vals[rng.integers(adv_vals.size)]  # always drawn, so `rng` advances the same in both modes
+      adv[r, p_end:b] = adv_rng.standard_normal() if advantages == "random" else a_val
+    pos[r, real_len:] = np.arange(budget - real_len)
+    real_mask[r, :real_len] = True
+  for r in range(rows - n_pad_rows, rows):
+    pos[r] = np.arange(budget)
+  ids[real_mask] = rng.integers(256, vocab_size, size=int(real_mask.sum()), dtype=ids.dtype)
+
+  out = dict(completion_ids=ids, segment_ids=seg, segment_positions=pos, completion_mask=cm, advantages=adv)
+  for name in ("old_per_token_logps", "rollout_per_token_logps"):
+    v = getattr(base, name, None)
+    if v is not None:
+      out[name] = np.where(cm > 0, np.float32(-0.7), np.float32(0.0)).astype(np.asarray(v).dtype)
+  if routed0 is not None:
+    num_layers, top_k = routed0.shape[2], routed0.shape[3]
+    routed = np.full(routed0.shape, UNSET_ROUTED_EXPERT, dtype=routed0.dtype)
+    targets = (layer_targets_override if layer_targets_override is not None
+               else layer_targets(num_layers, scale, cycle, probe_targets))
+    flat_real = np.flatnonzero(real_mask.ravel())
+    routed_flat = routed.reshape(rows * budget, num_layers, top_k)
+    for l in range(num_layers):
+      if routing == "uniform":
+        pool = _gumbel_topk(np.zeros(num_experts), pool_size, top_k, rng)
+      else:
+        pool = routing_pool(float(targets[l]), num_experts, ep, top_k, pool_size, rng, int(rng.integers(ep)))
+      routed_flat[flat_real, l, :] = pool[rng.integers(pool_size, size=flat_real.size)].astype(routed.dtype)
+    out["routed_experts"] = routed
+  return dataclasses.replace(base, **out)
+
+
+def payload_shard_load_ratio(payload, *, num_experts, ep=8, ep_rows=2, context=4, chunks=4):
+  """Per layer, the max over EP groups, token chunks and shards of the replayed (real-token) load over the
+  balanced load, laid out as the ring-of-experts layer sees it under cp-as-ep: an EP group is `ep_rows` rows
+  (expert axis) x `context` sequence slices, each slice split into `chunks` token chunks, and a chunk all-gathers
+  its slice-chunk from every member. The balanced load counts every token, padding included, as the layer's
+  buffer sizing does; padding's own (router-chosen) experts are unknown here and count for nothing, so this is the
+  replayed tokens' share. Returns (per_layer_ratio[L], pad_token_fraction).
+  """
+  routed = np.asarray(payload.routed_experts)
+  rows, budget, num_layers, top_k = routed.shape
+  per_shard = num_experts // ep
+  slice_len = budget // context
+  chunk_len = slice_len // chunks
+  best = np.zeros(num_layers)
+  for g in range(rows // ep_rows):
+    rsel = slice(g * ep_rows, (g + 1) * ep_rows)
+    for j in range(chunks):
+      cols = np.concatenate([np.arange(c * slice_len + j * chunk_len, c * slice_len + (j + 1) * chunk_len)
+                             for c in range(context)])
+      block = routed[rsel][:, cols]  # [ep_rows, tokens, L, K]
+      n_tok = block.shape[0] * block.shape[1]
+      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)
+  pad = float((np.asarray(payload.segment_ids) == 0).mean())
+  return best, pad
+
+
+# Orchestrator flags of the sampler-vs-trainer log-prob gates (GRPO loss: `sequence_loss_mask` and
+# `truncated_importance_weights`). LOGP_GATES=off strips them, so they take their argparse default None (gate off).
+LOGP_GATE_FLAGS = ("--seq_logprob_error_threshold", "--truncated_importance_sampling_type",
+                   "--truncated_importance_sampling_ratio_min", "--truncated_importance_sampling_ratio")
+
+
+def strip_logp_gates(argv):
+  """`argv` without the log-prob gate flags (both `--flag=v` and `--flag v` forms)."""
+  out, skip = [], False
+  for t in argv:
+    if skip:
+      skip = False
+      continue
+    name = t.split("=", 1)[0]
+    if name in LOGP_GATE_FLAGS:
+      skip = "=" not in t
+      continue
+    out.append(t)
+  return out
+
+
+def summarize_metrics(buf):
+  """Host-side summary of one step's MetricsBuffer (after `jax.device_get`).
+
+  The engine appends one entry per micro-batch to every metric recorded in `fwd_bwd` (`loss` and the loss's aux, as
+  WeightedMetric sums/denominators or scalar arrays) and one entry per step to those recorded in `update`
+  (gradient_norm, learning_rate, step_skipped, moe_* counters). Returns a dict with:
+    loss: the step's token-weighted loss, sum of the micro-batch sums over the sum of their denominators (the
+      denominator the engine normalizes the accumulated gradients by), NaN when that denominator is 0;
+    micro_loss / micro_denominator: per micro-batch sum/denominator (NaN where 0) and denominator;
+    weighted: name -> pooled value (sum over sums / max(sum over denominators, min_denom), NaN when 0);
+    scalar: name -> float for one entry, else the list of per-micro-batch values.
+  """
+  out = {"loss": float("nan"), "micro_loss": [], "micro_denominator": [], "weighted": {}, "scalar": {}}
+  for name, wm in sorted((getattr(buf, "weighted_metrics", None) or {}).items()):
+    sums = np.asarray(wm.unreduced_sum, dtype=np.float64).ravel()
+    dens = np.asarray(wm.denominator, dtype=np.float64).ravel()
+    tot = dens.sum() + (wm.eps or 0.0)
+    if wm.min_denom is not None:
+      tot = max(tot, wm.min_denom)
+    pooled = float(sums.sum() / tot) if tot != 0 else float("nan")
+    out["weighted"][name] = pooled
+    if name == "loss":
+      out["loss"] = pooled
+      with np.errstate(divide="ignore", invalid="ignore"):
+        out["micro_loss"] = [float(x) for x in np.where(dens != 0, sums / np.where(dens != 0, dens, 1), np.nan)]
+      out["micro_denominator"] = [float(x) for x in dens]
+  for name, v in sorted((getattr(buf, "scalar_metrics", None) or {}).items()):
+    arr = np.asarray(v, dtype=np.float64).ravel()
+    out["scalar"][name] = float(arr[0]) if arr.size == 1 else [float(x) for x in arr]
+  return out
+
+
+def metric_log_lines(step, summary):
+  """`STEP_LOSS` (loss, per-micro-batch loss and denominators, grad norm) and `STEP_METRICS` (every metric; a
+  per-micro-batch metric as its list) lines for one step."""
+  fmt = lambda x: f"{x:.9g}"
+  gn = summary["scalar"].get("gradient_norm", float("nan"))
+  loss_line = (f"STEP_LOSS step={step} loss={fmt(summary['loss'])} grad_norm={fmt(gn) if isinstance(gn, float) else gn} "
+               f"micro_loss=[{','.join(fmt(x) for x in summary['micro_loss'])}] "
+               f"micro_denominator=[{','.join(fmt(x) for x in summary['micro_denominator'])}]")
+  parts = [f"{k}={fmt(v)}" for k, v in summary["weighted"].items()]
+  for k, v in summary["scalar"].items():
+    parts.append(f"{k}={fmt(v)}" if isinstance(v, float) else f"{k}=[{','.join(fmt(x) for x in v)}]")
+  return loss_line, f"STEP_METRICS step={step} " + " ".join(parts)
+
+
+def main() -> int:
+  trainer_argv = shlex.split(os.environ["E2E_TRAINER_ARGV"])
+  orch_argv = shlex.split(os.environ["E2E_ORCH_ARGV"])
+  # LOGP_GATES=off (numerics A/B): the payload's rollout log-probs are synthetic (-0.7 per completion token), so
+  # against a random-init trainer the seq_logprob_error_threshold and seq-mask-tis gates drop every sequence, the
+  # loss mask is empty and the gradient is exactly 0. Stripping the gate flags keeps every sequence in the loss.
+  logp_gates = os.environ.get("LOGP_GATES", "on")
+  if logp_gates == "off":
+    orch_argv = strip_logp_gates(orch_argv)
+  elif logp_gates != "on":
+    raise ValueError(f"LOGP_GATES must be on or off, got {logp_gates!r}")
+  steps = int(os.environ.get("STEPS", "3"))
+  micro_steps = int(os.environ.get("MICRO_STEPS", "4"))
+
+  # Same JAX init as K8sJaxContext.initialize() under JAX_PLATFORMS=proxy.
+  import pathwaysutils  # pylint: disable=g-import-not-at-top
+  pathwaysutils.initialize()
+  import jax  # pylint: disable=g-import-not-at-top
+
+  from tunix.experimental.examples.common import run_trainer_node  # pylint: disable=g-import-not-at-top
+  from tunix.experimental.examples.deepswe_dist import run_deepswe_dist  # pylint: disable=g-import-not-at-top
+  from tunix.experimental.examples.deepswe_dist import trainer_warmup  # pylint: disable=g-import-not-at-top
+  from tunix.experimental.orchestrator import batch_assembly  # pylint: disable=g-import-not-at-top
+  from tunix.experimental.common import datatypes  # pylint: disable=g-import-not-at-top
+  from tunix.experimental.worker import trainer_worker  # pylint: disable=g-import-not-at-top
+  from transformers import AutoTokenizer  # pylint: disable=g-import-not-at-top
+
+  t_args = run_trainer_node._parse_args(trainer_argv)  # pylint: disable=protected-access
+  o_args = run_deepswe_dist._parse_args(orch_argv)  # pylint: disable=protected-access
+  logging.info("Trainer args: %s", t_args)
+  logging.info("Orchestrator args: %s", o_args)
+
+  # Trainer worker, as run_trainer_node.main builds it.
+  trainer_factory, mesh = run_trainer_node._create_maxtext_trainer_factory(t_args)  # pylint: disable=protected-access
+  worker = trainer_worker.TrainerWorker(
+      trainer_factory=trainer_factory,
+      worker_id=t_args.worker_id,
+      logps_chunk_size=t_args.compute_logps_chunk_size,
+      logps_micro_batch_size=t_args.compute_logps_micro_batch_size,
+      execution_context=mesh,
+  )
+  t0 = time.monotonic()
+  worker.initialize()
+  logging.info("Trainer worker initialized in %.1fs on %d devices.", time.monotonic() - t0, jax.device_count())
+
+  # Tokenizer ids, as run_deepswe_dist.main derives them.
+  tokenizer_path = o_args.tokenizer_path or os.getenv("MODEL_DIR") or o_args.model_id
+  tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, trust_remote_code=True)
+  if tokenizer.pad_token_id is None and tokenizer.eos_token is not None:
+    tokenizer.pad_token = tokenizer.eos_token
+  pad_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0
+
+  # Algorithm, BatchConfig and assembler, as run_deepswe_dist.main / StandardRLProgram build them.
+  algo = run_deepswe_dist._build_algo(o_args)  # pylint: disable=protected-access
+  batch_config = batch_assembly.BatchConfig(
+      pad_id=pad_id,
+      max_prompt_length=o_args.max_prompt_length,
+      max_response_length=o_args.max_response_length,
+      max_seq_token_per_tpu=o_args.max_seq_token_per_tpu,
+      max_segments_per_packed_row=o_args.max_segments_per_packed_row,
+      trainer_fsdp=o_args.trainer_fsdp,
+      trainer_dp=o_args.trainer_dp,
+      trainer_expert=o_args.trainer_expert,
+  )
+  assembler = batch_assembly.create_batch_assembler(
+      num_generations=o_args.num_generations,
+      mini_batch_size=o_args.mini_batch_size,
+      train_micro_batch_size=getattr(algo, "train_micro_batch_size", 1),
+      batch_config=batch_config,
+  )
+
+  # configure_worker (DistributedRLEngine.configure_worker, ACTOR role).
+  a_pad = getattr(assembler, "pad_id", pad_id)
+  a_eos = getattr(assembler, "eos_id", a_pad)
+  logging.info("configure_worker: loss fn + gen_model_input_fn (pad_id=%s, eos_id=%s).", a_pad, a_eos)
+  worker.with_loss_fn(algo.loss_fn(), has_aux=True)
+  worker.with_gen_model_input_fn(algo.build_gen_model_input_fn(pad_id=a_pad, eos_id=a_eos))
+
+  # TRAINER_WARMUP_COMPILE payload, as run_deepswe_dist.main builds it.
+  routed_shape = None
+  if o_args.trainer_warmup_routed_experts:
+    routed_shape = trainer_warmup.routed_experts_shape_from_config(tokenizer_path)
+    if routed_shape is None and o_args.trainer_warmup_routed_layers and o_args.trainer_warmup_routed_top_k:
+      routed_shape = (o_args.trainer_warmup_routed_layers, o_args.trainer_warmup_routed_top_k)
+  logging.info("routed_experts [num_layers, top_k] = %s", routed_shape)
+  payload = trainer_warmup.build_warmup_payload(
+      algo=algo,
+      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,
+      routed_experts_shape=routed_shape,
+  )
+  logging.info("Bench micro-batch: %s", batch_assembly.describe_payload(payload))
+
+  t0 = time.monotonic()
+  worker.compile(dummy_data=payload)
+  logging.info("Trainer compile finished in %.1fs.", time.monotonic() - t0)
+
+  # The micro-batches the steps run (see the module docstring).
+  mode = os.environ.get("PAYLOAD", "realistic")
+  regular = final = payload
+  if mode == "realistic":
+    cfg = worker._trainer._config  # pylint: disable=protected-access
+    knobs = dict(
+        num_experts=int(cfg.num_experts),
+        ep=int(os.environ.get("EP_SIZE", "8")),
+        routing=os.environ.get("ROUTING", "probe"),
+        scale=float(os.environ.get("ROUTING_SCALE", "1.0")),
+        pad_id=a_pad,
+        vocab_size=int(getattr(tokenizer, "vocab_size", 151_000)),
+        cycle=int(getattr(cfg, "inhomogeneous_layer_cycle_interval", 4) or 4),
+        advantages=os.environ.get("ADVANTAGES", "inherit"),
+    )
+    if os.environ.get("ROUTING_TARGETS"):
+      n_layers = int(np.asarray(payload.routed_experts).shape[2]) if payload.routed_experts is not None else 0
+      knobs["probe_targets"] = parse_routing_targets(os.environ["ROUTING_TARGETS"], n_layers)
+    seed = int(os.environ.get("PAYLOAD_SEED", "0"))
+    tg = time.monotonic()
+    regular = realistic_payload(payload, pad_fraction=float(os.environ.get("PAD_FRACTION", "0.07")), seed=seed, **knobs)
+    final = realistic_payload(payload, pad_fraction=float(os.environ.get("PAD_FRACTION", "0.07")),
+                              pad_rows_fraction=float(os.environ.get("FINAL_PAD_FRACTION", "0.25")), seed=seed + 1,
+                              **knobs)
+    logging.info("Realistic payloads built in %.1fs (%s).", time.monotonic() - tg, knobs)
+    for name, p in (("regular", regular), ("final", final)):
+      a = np.asarray(p.advantages)[np.asarray(p.completion_mask) > 0]
+      logging.info("PAYLOAD_ADVANTAGES %s rows=%d scored_tokens=%d nonzero_frac=%.4f mean=%.6g std=%.6g min=%.6g "
+                   "max=%.6g", name, np.asarray(p.completion_ids).shape[0], a.size,
+                   float((a != 0).mean()) if a.size else 0.0, a.mean() if a.size else 0.0,
+                   a.std() if a.size else 0.0, a.min() if a.size else 0.0, a.max() if a.size else 0.0)
+    for name, p in (("regular", regular), ("final", final)):
+      if p.routed_experts is None:
+        continue
+      ratios, pad = payload_shard_load_ratio(p, num_experts=knobs["num_experts"], ep=knobs["ep"])
+      top = np.argsort(-ratios)[:5]
+      logging.info(
+          "PAYLOAD_ROUTING %s pad_fraction=%.4f replayed_max_ratio=%.4f hottest_layers=%s ratios=[%s]",
+          name, pad, ratios.max(), [(int(l), round(float(ratios[l]), 3)) for l in top],
+          ",".join(f"{v:.3f}" for v in ratios),
+      )
+  elif mode != "warmup":
+    raise ValueError(f"PAYLOAD must be realistic or warmup, got {mode!r}")
+
+  for step in range(steps):
+    t0 = time.monotonic()
+    times = []
+    for m in range(micro_steps):
+      tm = time.monotonic()
+      last = m == micro_steps - 1
+      req = datatypes.TrainRequest(request_id=f"bench_{step}_{m}", payload=final if last else regular, metadata={})
+      # As DistributedRLEngine.train_step: the step's last micro-batch is the one with apply_optimizer=True.
+      # Per-micro-step runtime load ratios: MOE_DROP_AUDIT lines with the MaxText flag log_moe_drops=true.
+      worker.fwd_bwd(request=req, apply_optimizer=last)
+      times.append(time.monotonic() - tm)
+    worker.update()
+    # Drain the step's device work before reading the clock, as maxtext_engine/train.py does.
+    throttler = getattr(worker._trainer, "_throttler", None)  # pylint: disable=protected-access
+    if throttler is not None:
+      throttler.wait_for_all()
+    else:
+      jax.block_until_ready(worker.get_metrics())
+    dt = time.monotonic() - t0
+    logging.info(
+        "completed step: %d, seconds: %.3f, micro_steps: %d, seconds/micro_step: %.3f, fwd_bwd dispatch s: %s",
+        step, dt, micro_steps, dt / micro_steps, [round(x, 3) for x in times],
+    )
+    # Numerics comparison (e.g. moe_exclude_padding_from_routing on vs off): the step's loss (per micro-batch and
+    # pooled), gradient norm and every other trainer metric, read after the step's timing so they do not affect it.
+    try:
+      loss_line, metrics_line = metric_log_lines(step, summarize_metrics(jax.device_get(worker.get_metrics())))
+      logging.info("%s", loss_line)
+      logging.info("%s", metrics_line)
+    except Exception as e:  # pylint: disable=broad-except
+      logging.info("STEP_METRICS step=%d unavailable: %r", step, e)
+  worker.stop()
+  return 0
+
+
+if __name__ == "__main__":
+  sys.exit(main())
diff --git a/src/maxtext/experimental/rl_trainer_bench/hlo_collective_schedule.py b/src/maxtext/experimental/rl_trainer_bench/hlo_collective_schedule.py
new file mode 100644
index 0000000000..d6130d08ca
--- /dev/null
+++ b/src/maxtext/experimental/rl_trainer_bench/hlo_collective_schedule.py
@@ -0,0 +1,142 @@
+#!/usr/bin/env python3
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Static schedule summary of the collectives in a scheduled (after_optimizations) XLA TPU HLO text dump.
+
+  python3 hlo_collective_schedule.py <module.after_optimizations.txt> [--min-mib 1] [--detail]
+
+For every async-start whose wrapped op is a collective (all-gather, reduce-scatter, all-reduce, all-to-all,
+collective-permute; recognised by the start's name, e.g. all-gather.2222.cloned.1.call-start), it finds the matching
+async-done in the same computation and counts what the schedule places between them: instructions, and "compute"
+instructions (fusion / convolution / custom-call / dot). The instruction order in a scheduled dump is the TensorCore
+program order, so "compute between start and done" is the static overlap the latency-hiding scheduler granted. It
+also counts synchronous collectives (all-gather(...) etc. issued directly in a scheduled computation, not wrapped).
+Sizes are from the done's result shape (per device). This is a static view: it cannot see runtime latency,
+SparseCore queueing behind other SC work, or ICI contention.
+"""
+import argparse
+import collections
+import re
+import statistics
+import sys
+
+DT = {"bf16": 2, "f32": 4, "s32": 4, "u32": 4, "f16": 2, "s8": 1, "u8": 1, "pred": 1, "f8e4m3fn": 1, "s16": 2,
+      "u16": 2, "f64": 8, "s64": 8, "u64": 8, "f8e5m2": 1}
+KINDS = ("all-gather", "reduce-scatter", "all-reduce", "all-to-all", "collective-permute")
+COMPUTE = re.compile(r"\b(fusion|convolution|custom-call|dot)\(")
+SHAPE = re.compile(r"\b([a-z0-9]+)\[([0-9,]*)\]")
+
+
+def nbytes(shape_text):
+  tot = 0
+  for dt, dims in SHAPE.findall(shape_text):
+    if dt not in DT:
+      continue
+    n = 1
+    for d in dims.split(","):
+      if d:
+        n *= int(d)
+    tot += n * DT[dt]
+  return tot
+
+
+def kind_of(name):
+  base = name.lstrip("%")
+  for k in KINDS:
+    if base.startswith(k):
+      return k
+  return None
+
+
+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 main():
+  ap = argparse.ArgumentParser()
+  ap.add_argument("path")
+  ap.add_argument("--min-mib", type=float, default=0.0)
+  ap.add_argument("--detail", action="store_true")
+  a = ap.parse_args()
+  comps = parse(a.path)
+  rows = []
+  sync = collections.Counter()
+  sync_bytes = collections.Counter()
+  for cname, instrs in comps.items():
+    # Only scheduled computations that hold async-starts or wrapped collectives matter; computations that are
+    # the async bodies themselves (async_computation.*, wrapped_*) are skipped as places.
+    pos = {n: i for i, (n, _) in enumerate(instrs)}
+    for i, (n, rhs) in enumerate(instrs):
+      native = re.search(r"\) (all-gather|reduce-scatter|all-reduce|all-to-all|collective-permute)-start\(", rhs)
+      if " async-start(" in rhs or native:
+        k = native.group(1) if native else kind_of(n)
+        if not k:
+          continue
+        thread = re.search(r'async_execution_thread="(\w+)"', rhs)
+        thread = thread.group(1) if thread else "tensorcore"
+        done_op = f" {k}-done(" if native else " async-done("
+        done = None
+        for j in range(i + 1, len(instrs)):
+          if done_op in instrs[j][1] and f"%{n})" in instrs[j][1]:
+            done = j
+            break
+        if done is None:
+          continue
+        between = instrs[i + 1:done]
+        ncomp = sum(1 for _, r in between if COMPUTE.search(r.split(" metadata=")[0]))
+        size = nbytes(instrs[done][1].split(done_op)[0])
+        rows.append((k, thread, "entry" if cname.startswith("main.") else "loop", size, done - i, ncomp,
+                     cname, n))
+      else:
+        m = re.search(r"[\]\}\)] (all-gather|reduce-scatter|all-reduce|all-to-all|collective-permute)\(", rhs)
+        if m and any(" async-start(" in r for _, r in instrs):  # only scheduled computations, not reducers
+          sync[(m.group(1), "entry" if cname.startswith("main.") else "loop")] += 1
+          sync_bytes[(m.group(1), "entry" if cname.startswith("main.") else "loop")] += nbytes(rhs.split(m.group(0))[0])
+  rows = [r for r in rows if r[3] >= a.min_mib * 2**20]
+  agg = collections.defaultdict(list)
+  for r in rows:
+    agg[(r[0], r[1], r[2])].append(r)
+  print(f"{'kind':16s} {'thread':10s} {'where':5s} {'n':>5s} {'MiB med':>8s} {'MiB tot':>9s} {'instr med':>9s} "
+        f"{'compute med':>11s} "
+        f"{'compute=0':>9s}")
+  for (k, t, w), rs in sorted(agg.items()):
+    print(f"{k:16s} {t:10s} {w:5s} {len(rs):5d} {statistics.median(r[3] for r in rs) / 2**20:8.2f} "
+          f"{sum(r[3] for r in rs) / 2**20:9.1f} {statistics.median(r[4] for r in rs):9.0f} "
+          f"{statistics.median(r[5] for r in rs):11.0f} {sum(1 for r in rs if r[5] == 0):9d}")
+  print("synchronous collectives (no start/done; the TensorCore waits on them) in scheduled computations:")
+  for key in sorted(sync):
+    print(f"  {key[0]:16s} {key[1]:5s} n={sync[key]:4d} MiB tot={sync_bytes[key] / 2**20:9.1f}")
+  if a.detail:
+    for r in sorted(rows, key=lambda r: (r[0], r[6], r[7])):
+      print(f"  {r[0]:15s} {r[1]:10s} {r[3] / 2**20:8.2f} MiB instr={r[4]:5d} compute={r[5]:4d} {r[6][:40]} {r[7]}")
+
+
+if __name__ == "__main__":
+  sys.exit(main())
diff --git a/src/maxtext/experimental/rl_trainer_bench/render_e2e.py b/src/maxtext/experimental/rl_trainer_bench/render_e2e.py
new file mode 100644
index 0000000000..56a62f4988
--- /dev/null
+++ b/src/maxtext/experimental/rl_trainer_bench/render_e2e.py
@@ -0,0 +1,468 @@
+#!/usr/bin/env python3
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Renders the RL trainer benchmark JobSet from a DRY_RUN render of a Tunix GRPO e2e run.
+
+Input: the output of the e2e launcher with DRY_RUN=true (e.g. `DRY_RUN=true mlperf_gbs1024_submission.sh start`),
+saved to a file (`--e2e-render`). That output is a stream of YAML documents (one JobSet per orchestrator, trainer
+and rollout engine) that may be interleaved with launcher log lines; the trainer JobSet and the orchestrator JobSet
+are picked from it by the `--process_main=` of their `main` container. Alternatively pass the two JobSets as
+separate files (`--train-yaml`, `--orch-yaml`).
+
+The trainer JobSet is kept as is (Pathways resource manager, proxy and worker containers and their args, env,
+images, node selectors, tolerations, reservations, Kueue labels and annotations) except:
+
+  * names: JobSet name, namespace, gang-id label; failurePolicy.maxRestarts=0 (a failed bench should not rerun
+    the compile);
+  * the client command: the same inline env prefix (MAXTEXT_EXTRA_FLAGS, LIBTPU_INIT_ARGS, ...) in front of
+    e2e_driver.py instead of the Tunix runtime main, with the e2e trainer's flags (E2E_TRAINER_ARGV) and the e2e
+    orchestrator's flags (E2E_ORCH_ARGV); the Tunix fail-fast wrapper and the gRPC port are dropped;
+  * trainer flags: --maxtext_ckpt_path is cleared (random init, the one intended difference),
+    --maxtext_output_directory points at `--output-prefix/<name>` and --worker_id is the bench name;
+  * optional knobs for A/B variants: --libtpu-extra / --libtpu-drop (written to every place the e2e render carries
+    libtpu flags: the client's LIBTPU_INIT_ARGS, the pathways-proxy --xla_* args and the pw-node workers'
+    LIBTPU_INIT_ARGS), --extra-maxtext-flags (appended to MAXTEXT_EXTRA_FLAGS), --replace-maxtext-flags (replace an
+    existing MAXTEXT_EXTRA_FLAGS key in place), --set-trainer-flags / --set-orch-flags (replace an existing trainer or
+    orchestrator flag in place), --compute-logps-chunk-size, --profiler-skip / --profiler-steps;
+  * optional --trainer-slice (e.g. 4x8x8): the trainer slice and the FSDP size that fills it, with expert, context
+    and tp kept. It changes the Pathways resource manager's --instance_type, the pw-node completions/parallelism
+    (hosts = chips / chips per host), the pw-node gke-tpu-slice-topology annotation, the trainer's --mesh_fsdp and
+    the orchestrator's --trainer_fsdp (which sets the assembler's pack size, fsdp x expert rows).
+
+Writes a file and nothing else; it never talks to a cluster. With --print-aot-env it prints shell exports of the
+trainer's MaxText and XLA flags and mesh for aot_e2e.sh instead of rendering.
+"""
+
+import argparse
+import copy
+import os
+import re
+import shlex
+import sys
+
+import yaml
+
+RUNTIME_MAIN = "python -m tunix.experimental.distributed.runtime.main"
+TRAINER_MAIN = "tunix.experimental.examples.common.run_trainer_node.main"
+ORCH_MAIN = "tunix.experimental.examples.deepswe_dist.run_deepswe_dist.main"
+MAX_NAME_LEN = 30  # pod hostnames <name>-pw-node-0-<idx>-<5 chars> must stay within 63 characters.
+TOPOLOGY_ANNOTATION = "cloud.google.com/gke-tpu-slice-topology"
+
+
+class _LiteralDumper(yaml.SafeDumper):
+  """Dumps multi-line strings (the client script) as literal blocks so the rendered YAML is readable."""
+
+
+def _str_representer(dumper, data):
+  style = "|" if "\n" in data else None
+  return dumper.represent_scalar("tag:yaml.org,2002:str", data, style=style)
+
+
+_LiteralDumper.add_representer(str, _str_representer)
+
+
+def find(items, name):
+  for it in items:
+    if it["name"] == name:
+      return it
+  raise KeyError(name)
+
+
+def compose_libtpu(base: str, drop: str, extra: str) -> list[str]:
+  """Returns base flags minus any whose name is in `drop`, plus `extra` (extra overrides same-named flags)."""
+
+  def key(f):
+    return f.split("=", 1)[0]
+
+  drop_keys = {key(f) for f in drop.split()}
+  extra_flags = extra.split()
+  extra_keys = {key(f) for f in extra_flags}
+  out = [f for f in base.split() if key(f) not in drop_keys and key(f) not in extra_keys]
+  return out + extra_flags
+
+
+def _proc_main(js):
+  proc = find(js["spec"]["replicatedJobs"], "proc")["template"]["spec"]["template"]["spec"]
+  return proc, find(proc["containers"], "main")
+
+
+def load_documents(path: str) -> list[dict]:
+  """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:
+    text = f.read()
+  for chunk in re.split(r"(?m)^---\s*$", text):
+    try:
+      d = yaml.safe_load(chunk)
+    except yaml.YAMLError:
+      continue
+    if isinstance(d, dict) and d.get("kind") and d.get("spec"):
+      docs.append(d)
+  return docs
+
+
+def process_main_of(js) -> str | None:
+  """The `--process_main=` value of a JobSet's proc/main container command, or None."""
+  try:
+    _, main_c = _proc_main(js)
+  except (KeyError, TypeError):
+    return None
+  m = re.search(r"--process_main=(\S+)", " ".join(main_c.get("command", [])))
+  return m.group(1).rstrip("\\") if m else None
+
+
+def select_jobsets(docs: list[dict], trainer_main: str = TRAINER_MAIN, orch_main: str = ORCH_MAIN):
+  """(trainer JobSet, orchestrator JobSet) of a DRY_RUN render; fails unless each matches exactly once."""
+  out = []
+  for want in (trainer_main, orch_main):
+    hits = [d for d in docs if process_main_of(d) == want]
+    if len(hits) != 1:
+      names = [d.get("metadata", {}).get("name") for d in hits]
+      raise ValueError(f"expected one JobSet with --process_main={want}, found {len(hits)}: {names}")
+    out.append(hits[0])
+  return out[0], out[1]
+
+
+def split_command(cmd: str) -> tuple[str, list[str]]:
+  """Returns (inline env prefix, flags after --process_main=...) of a Tunix runtime launch command."""
+  i = cmd.index("(\n") + 2
+  j = cmd.index(RUNTIME_MAIN)
+  k = cmd.index("\n)", j)
+  env_prefix = cmd[i:j]
+  # Backslash-newline is a line continuation in bash; shlex would keep the newline as a token.
+  tokens = shlex.split(cmd[j + len(RUNTIME_MAIN) : k].replace("\\\n", " "))
+  for n, t in enumerate(tokens):
+    if t.startswith("--process_main="):
+      return env_prefix, tokens[n + 1 :]
+  raise ValueError("no --process_main in command")
+
+
+def set_flag(tokens: list[str], name: str, value: str) -> list[str]:
+  """Removes every --name= and appends --name=value."""
+  out = [t for t in tokens if not t.startswith(f"--{name}=")]
+  return out + [f"--{name}={value}"]
+
+
+def set_in_place(tokens: list[str], kvs: list[str], what: str = "trainer") -> list[str]:
+  """Replaces existing `--k=` flags in place (keeps argv order); fails unless each key is there exactly once."""
+  tokens = list(tokens)
+  for kv in kvs:
+    k, _, v = kv.partition("=")
+    hits = [i for i, t in enumerate(tokens) if t.startswith(f"--{k}=")]
+    if len(hits) != 1:
+      raise ValueError(f"--set-{what}-flags {kv!r}: --{k}= found {len(hits)} times in the {what} argv")
+    tokens[hits[0]] = f"--{k}={v}"
+  return tokens
+
+
+def upsert_flag(tokens: list[str], name: str, value: str) -> list[str]:
+  """Replaces --name= in place if present once, appends it if absent, fails if present more than once."""
+  hits = [i for i, t in enumerate(tokens) if t.startswith(f"--{name}=")]
+  if len(hits) > 1:
+    raise ValueError(f"--{name}= found {len(hits)} times")
+  if not hits:
+    return list(tokens) + [f"--{name}={value}"]
+  return set_in_place(tokens, [f"{name}={value}"])
+
+
+def maxtext_extra_flags(prefix: str) -> str:
+  m = re.findall(r'MAXTEXT_EXTRA_FLAGS="([^"]*)"', prefix)
+  if len(m) != 1:
+    raise ValueError(f"expected one MAXTEXT_EXTRA_FLAGS in the env prefix, found {len(m)}")
+  return m[0]
+
+
+def libtpu_init_args(prefix: str) -> str:
+  m = re.findall(r"LIBTPU_INIT_ARGS='([^']*)'", prefix)
+  if len(m) != 1:
+    raise ValueError(f"expected one LIBTPU_INIT_ARGS='...' in the env prefix, found {len(m)}")
+  return m[0]
+
+
+def rewrite_env_prefix(prefix: str, libtpu: list[str], extra_maxtext: str, replace_maxtext: str = "") -> str:
+  # LIBTPU_INIT_ARGS='...' (single-quoted in the e2e render).
+  prefix, n = re.subn(
+      r"LIBTPU_INIT_ARGS='[^']*'", lambda _: "LIBTPU_INIT_ARGS=" + shlex.quote(" ".join(libtpu)), prefix
+  )
+  if n != 1:
+    raise ValueError(f"expected one LIBTPU_INIT_ARGS in the env prefix, found {n}")
+  for kv in replace_maxtext.split():
+    # Replaces an existing key=value of MAXTEXT_EXTRA_FLAGS in place (e.g. context=remat for context=offload), so the
+    # rendered flag set carries one value per key. Fails if the key is not there exactly once.
+    key = kv.split("=", 1)[0]
+    m = re.search(r'MAXTEXT_EXTRA_FLAGS="([^"]*)"', prefix)
+    if m is None:
+      raise ValueError("no MAXTEXT_EXTRA_FLAGS in the env prefix")
+    toks = m.group(1).split(" ")
+    hits = [i for i, t in enumerate(toks) if t.split("=", 1)[0] == key]
+    if len(hits) != 1:
+      raise ValueError(f"--replace-maxtext-flags {kv!r}: key found {len(hits)} times in MAXTEXT_EXTRA_FLAGS")
+    toks[hits[0]] = kv
+    prefix = prefix[: m.start(1)] + " ".join(toks) + prefix[m.end(1) :]
+  if extra_maxtext:
+    prefix, n = re.subn(
+        r'MAXTEXT_EXTRA_FLAGS="([^"]*)"', lambda m: f'MAXTEXT_EXTRA_FLAGS="{m.group(1)} {extra_maxtext}"', prefix
+    )
+    if n != 1:
+      raise ValueError(f"expected one MAXTEXT_EXTRA_FLAGS in the env prefix, found {n}")
+  return prefix
+
+
+def _flag_value(tokens: list[str], name: str) -> str:
+  vals = [t.split("=", 1)[1] for t in tokens if t.startswith(f"--{name}=")]
+  if len(vals) != 1:
+    raise ValueError(f"expected one --{name}= flag, found {vals}")
+  return vals[0]
+
+
+def current_slice(js: dict) -> tuple[str, str]:
+  """(accelerator, topology) of the trainer JobSet, e.g. ("tpu7x", "4x4x8"), from the resource manager's
+  --instance_type; checks it against the pw-node slice-topology annotation."""
+  proc, _ = _proc_main(js)
+  rm = find(proc["initContainers"], "pathways-rm")
+  it = [x.split("=", 1)[1] for x in rm["args"] if x.startswith("--instance_type=")]
+  if len(it) != 1 or ":" not in it[0]:
+    raise ValueError(f"expected one --instance_type=<accelerator>:<topology> in pathways-rm args, found {it}")
+  accel, topo = it[0].split(":", 1)
+  pw = find(js["spec"]["replicatedJobs"], "pw-node")["template"]["spec"]
+  ann = pw["template"]["metadata"]["annotations"]
+  if ann.get(TOPOLOGY_ANNOTATION) != topo:
+    raise ValueError(f"pathways-rm --instance_type {it[0]} disagrees with pw-node {TOPOLOGY_ANNOTATION}: {ann}")
+  return accel, topo
+
+
+def apply_trainer_slice(
+    js: dict, trainer_flags: list[str], orch_flags: list[str], topo: str, devices_per_chip: int = 2
+):
+  """Re-targets the trainer JobSet and flags at a `topo` slice (e.g. 4x8x8) of the same accelerator, filling it with
+  FSDP.
+
+  Returns (trainer_flags, orch_flags, summary). Checks the reference values first, so an unexpected input fails here
+  rather than rendering a half-converted JobSet. Chips per host come from the pw-node's google.com/tpu limit.
+  """
+  accel, ref_topo = current_slice(js)
+  dims = [int(x) for x in topo.split("x")]
+  chips = 1
+  for d in dims:
+    chips *= d
+  pw = find(js["spec"]["replicatedJobs"], "pw-node")["template"]["spec"]
+  worker = find(pw["template"]["spec"]["containers"], "pathways-worker")
+  chips_per_host = int(worker["resources"]["limits"]["google.com/tpu"])
+  if chips % chips_per_host:
+    raise ValueError(f"{chips} chips not divisible by {chips_per_host} chips per host")
+  hosts = chips // chips_per_host
+  devices = chips * devices_per_chip
+  tp, ep, cp = (int(_flag_value(trainer_flags, f"mesh_{k}")) for k in ("tp", "expert", "context"))
+  if devices % (tp * ep * cp):
+    raise ValueError(f"{devices} devices not divisible by tp*expert*context={tp * ep * cp}")
+  fsdp = devices // (tp * ep * cp)
+  if _flag_value(orch_flags, "trainer_expert") != str(ep):
+    raise ValueError("orchestrator --trainer_expert differs from the trainer's --mesh_expert")
+  if pw["completions"] != pw["parallelism"]:
+    raise ValueError(f"unexpected pw-node completions/parallelism: {pw['completions']}/{pw['parallelism']}")
+
+  proc, _ = _proc_main(js)
+  rm = find(proc["initContainers"], "pathways-rm")
+  it = [n for n, x in enumerate(rm["args"]) if x.startswith("--instance_type=")]
+  rm["args"][it[0]] = f"--instance_type={accel}:{topo}"
+  pw["template"]["metadata"]["annotations"][TOPOLOGY_ANNOTATION] = topo
+  pw["completions"] = pw["parallelism"] = hosts
+  trainer_flags = set_flag(trainer_flags, "mesh_fsdp", str(fsdp))
+  orch_flags = set_flag(orch_flags, "trainer_fsdp", str(fsdp))
+  summary = (
+      f"trainer slice {accel}:{ref_topo} -> {accel}:{topo}: {chips} chips, {hosts} hosts, {devices} devices; "
+      f"mesh fsdp={fsdp} tp={tp} expert={ep} context={cp}; pack size {fsdp * ep} rows per micro-batch"
+  )
+  return trainer_flags, orch_flags, summary
+
+
+def client_script(
+    env_prefix: str,
+    driver_src: str,
+    trainer_argv: list[str],
+    orch_argv: list[str],
+    steps: int,
+    micro_steps: int,
+    driver_env: str = "",
+) -> str:
+  """The client container's bash script: writes the driver to /tmp and runs it under the e2e env prefix."""
+  bench_env = (
+      f"E2E_TRAINER_ARGV={shlex.quote(shlex.join(trainer_argv))} "
+      f"E2E_ORCH_ARGV={shlex.quote(shlex.join(orch_argv))} STEPS={steps} MICRO_STEPS={micro_steps}"
+  )
+  for kv in shlex.split(driver_env):
+    k, _, v = kv.partition("=")
+    if not re.fullmatch(r"[A-Z_][A-Z0-9_]*", k):
+      raise ValueError(f"bad --driver-env entry {kv!r}")
+    bench_env += f" {k}={shlex.quote(v)}"
+  if "E2E_DRIVER_EOF" in driver_src:
+    raise ValueError("driver source contains the heredoc terminator")
+  return f"""# PID 1 wrapper: forwards SIGTERM to the benchmark so the grace period shuts it down cleanly.
+echo "Program Start: $(date)"
+cat > /tmp/e2e_driver.py <<'E2E_DRIVER_EOF'
+{driver_src}E2E_DRIVER_EOF
+_sigterm() {{ kill -SIGTERM "$!" 2>/dev/null; }}
+trap _sigterm SIGTERM
+(
+{env_prefix} {bench_env} exec python /tmp/e2e_driver.py
+) & PID=$!
+while kill -0 $PID 2>/dev/null; do
+  sleep 5
+done
+wait $PID
+EXIT_CODE=$?
+echo "Program End: $(date)"
+echo "EXIT_CODE=$EXIT_CODE"
+exit $EXIT_CODE
+"""
+
+
+def aot_env(env_prefix: str, libtpu: list[str], trainer_flags: list[str]) -> str:
+  """Shell exports for aot_e2e.sh: the trainer's MaxText flags, its --xla_* libtpu flags and its mesh."""
+  lines = [
+      f"export MAXTEXT_EXTRA_FLAGS={shlex.quote(maxtext_extra_flags(env_prefix))}",
+      f"export LIBTPU_INIT_ARGS={shlex.quote(' '.join(f for f in libtpu if f.startswith('--xla_')))}",
+  ]
+  for k in ("fsdp", "tp", "expert", "context"):
+    lines.append(f"export AOT_MESH_{k.upper()}={_flag_value(trainer_flags, f'mesh_{k}')}")
+  chunk = [t.split("=", 1)[1] for t in trainer_flags if t.startswith("--compute_logps_chunk_size=")]
+  if len(chunk) == 1:
+    lines.append(f"export AOT_LOGPS_CHUNK_SIZE={chunk[0]}")
+  return "\n".join(lines)
+
+
+def build_parser() -> argparse.ArgumentParser:
+  p = argparse.ArgumentParser(description=__doc__.split("\n", 1)[0])
+  p.add_argument("--e2e-render", default="", help="DRY_RUN output of the e2e launcher (multi-document YAML)")
+  p.add_argument("--train-yaml", default="", help="the e2e trainer JobSet alone (instead of --e2e-render)")
+  p.add_argument("--orch-yaml", default="", help="the e2e orchestrator JobSet alone (instead of --e2e-render)")
+  p.add_argument("--trainer-main", default=TRAINER_MAIN, help="--process_main of the trainer JobSet")
+  p.add_argument("--orch-main", default=ORCH_MAIN, help="--process_main of the orchestrator JobSet")
+  p.add_argument("--driver", default=os.path.join(os.path.dirname(os.path.abspath(__file__)), "e2e_driver.py"))
+  p.add_argument("--out", default="")
+  p.add_argument("--name", default="")
+  p.add_argument("--namespace", default="")
+  p.add_argument("--output-prefix", default="", help="the bench writes to <output-prefix>/<name>")
+  p.add_argument("--steps", type=int, default=3)
+  p.add_argument("--micro-steps", type=int, default=4)
+  p.add_argument("--libtpu-drop", default="")
+  p.add_argument("--libtpu-extra", default="")
+  p.add_argument("--extra-maxtext-flags", default="")
+  p.add_argument("--replace-maxtext-flags", default="", help='"k=v ...": replace existing keys in place')
+  p.add_argument("--set-trainer-flags", default="", help='"k=v ...": replace existing trainer flags in place')
+  p.add_argument("--set-orch-flags", default="", help='"k=v ...": replace existing orchestrator flags in place')
+  p.add_argument("--image", default="", help="client (main container) image; empty keeps the e2e render's")
+  # Driver knobs (e2e_driver.py), as "K=V K=V": PAYLOAD ROUTING ROUTING_SCALE ROUTING_TARGETS PAD_FRACTION
+  # FINAL_PAD_FRACTION PAYLOAD_SEED EP_SIZE ADVANTAGES LOGP_GATES.
+  p.add_argument("--driver-env", default="")
+  p.add_argument("--trainer-slice", default="", help="e.g. 4x8x8: FSDP fills it, expert/context/tp kept")
+  p.add_argument("--devices-per-chip", type=int, default=2, help="JAX devices per chip (tpu7x: 2)")
+  p.add_argument("--compute-logps-chunk-size", default="", help="trainer --compute_logps_chunk_size; empty keeps")
+  p.add_argument("--profiler-skip", default="", help="trainer --skip_first_n_profiler_steps (counts fwd_bwd calls)")
+  p.add_argument("--profiler-steps", default="", help="trainer --profiler_steps (fwd_bwd calls in the window)")
+  p.add_argument("--print-aot-env", action="store_true", help="print exports for aot_e2e.sh instead of rendering")
+  return p
+
+
+def main(argv=None):
+  a = build_parser().parse_args(argv)
+  if a.e2e_render:
+    if a.train_yaml or a.orch_yaml:
+      sys.exit("pass either --e2e-render or --train-yaml/--orch-yaml")
+    train, orch = select_jobsets(load_documents(a.e2e_render), a.trainer_main, a.orch_main)
+  elif a.train_yaml and a.orch_yaml:
+    with open(a.train_yaml) as f:
+      train = yaml.safe_load(f)
+    with open(a.orch_yaml) as f:
+      orch = yaml.safe_load(f)
+  else:
+    sys.exit("pass --e2e-render, or both --train-yaml and --orch-yaml")
+  if not a.print_aot_env:
+    missing = [k for k in ("out", "name", "namespace", "output_prefix") if not getattr(a, k)]
+    if missing:
+      sys.exit(f"missing required: {', '.join('--' + k.replace('_', '-') for k in missing)}")
+    if len(a.name) > MAX_NAME_LEN or not re.fullmatch(r"[a-z0-9-]+", a.name):
+      sys.exit(f"NAME must be lowercase alphanumeric/dashes, <= {MAX_NAME_LEN} chars: {a.name!r}")
+
+  js = copy.deepcopy(train)
+  _, orch_main = _proc_main(orch)
+  _, orch_flags = split_command(orch_main["command"][-1])
+
+  proc, main_c = _proc_main(js)
+  env_prefix, trainer_flags = split_command(main_c["command"][-1])
+  base_libtpu = libtpu_init_args(env_prefix)
+  libtpu = compose_libtpu(base_libtpu, a.libtpu_drop, a.libtpu_extra)
+  env_prefix = rewrite_env_prefix(env_prefix, libtpu, a.extra_maxtext_flags, a.replace_maxtext_flags)
+
+  trainer_flags = set_flag(trainer_flags, "maxtext_ckpt_path", "")
+  trainer_flags = set_flag(trainer_flags, "maxtext_output_directory", f"{a.output_prefix}/{a.name}")
+  trainer_flags = set_flag(trainer_flags, "worker_id", a.name)
+  if a.compute_logps_chunk_size:
+    trainer_flags = set_in_place(trainer_flags, [f"compute_logps_chunk_size={int(a.compute_logps_chunk_size)}"])
+  trainer_flags = [t for t in trainer_flags if not t.startswith(("--discovery_addrs=", "--port="))]
+  if a.trainer_slice and a.trainer_slice != current_slice(js)[1]:
+    trainer_flags, orch_flags, summary = apply_trainer_slice(
+        js, trainer_flags, orch_flags, a.trainer_slice, a.devices_per_chip
+    )
+    print(summary, file=sys.stderr)
+  # After --trainer-slice, so an explicit flag wins over the slice's (e.g. --set-orch-flags trainer_fsdp=...).
+  trainer_flags = set_in_place(trainer_flags, a.set_trainer_flags.split(), "trainer")
+  orch_flags = set_in_place(orch_flags, a.set_orch_flags.split(), "orch")
+  if a.profiler_steps:
+    trainer_flags = upsert_flag(trainer_flags, "profiler_steps", str(int(a.profiler_steps)))
+  if a.profiler_skip:
+    trainer_flags = upsert_flag(trainer_flags, "skip_first_n_profiler_steps", str(int(a.profiler_skip)))
+
+  if a.print_aot_env:
+    print(aot_env(env_prefix, libtpu, trainer_flags))
+    return 0
+
+  md = js["metadata"]
+  md["name"], md["namespace"] = a.name, a.namespace
+  md.setdefault("labels", {})["scheduling.x-k8s.io/gang-id"] = a.name
+  js["spec"].setdefault("failurePolicy", {})["maxRestarts"] = 0
+
+  with open(a.driver) as f:
+    driver_src = f.read()
+  main_c["command"] = [
+      "bash",
+      "-c",
+      client_script(env_prefix, driver_src, trainer_flags, orch_flags, a.steps, a.micro_steps, a.driver_env),
+  ]
+  main_c.pop("ports", None)
+  if a.image:
+    main_c["image"] = a.image
+
+  if a.libtpu_drop or a.libtpu_extra:
+    proxy = find(proc["initContainers"], "pathways-proxy")
+    orig = proxy["args"]
+    head = [x for x in orig if not x.startswith("--xla_") and not x.startswith("--sidecar_name")]
+    tail = [x for x in orig if x.startswith("--sidecar_name")]
+    proxy["args"] = head + [f for f in libtpu if f.startswith("--xla_")] + tail
+    worker_pod = find(js["spec"]["replicatedJobs"], "pw-node")["template"]["spec"]["template"]["spec"]
+    env = find(find(worker_pod["containers"], "pathways-worker")["env"], "LIBTPU_INIT_ARGS")
+    base_keys = {x.split("=", 1)[0] for x in base_libtpu.split()}
+    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:
+    yaml.dump(js, f, Dumper=_LiteralDumper, sort_keys=False, width=120)
+  print(a.out)
+  return 0
+
+
+if __name__ == "__main__":
+  sys.exit(main())
diff --git a/src/maxtext/experimental/rl_trainer_bench/render_e2e.sh b/src/maxtext/experimental/rl_trainer_bench/render_e2e.sh
new file mode 100755
index 0000000000..d3ae4e0409
--- /dev/null
+++ b/src/maxtext/experimental/rl_trainer_bench/render_e2e.sh
@@ -0,0 +1,61 @@
+#!/usr/bin/env bash
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Renders the RL trainer benchmark JobSet (e2e_driver.py on the trainer JobSet of an e2e DRY_RUN render).
+# Writes one YAML file and nothing else; it never talks to a cluster.
+#
+# Required:
+#   E2E_RENDER=<file>        DRY_RUN output of the e2e launcher (see README)
+#   NAME=<name>              JobSet name, worker id and output folder (lowercase, at most 30 characters)
+#   K8S_NAMESPACE=<namespace>
+#   OUTPUT_PREFIX=<prefix>   the bench writes to <prefix>/<NAME> (e.g. a GCS folder you own)
+# Optional:
+#   OUT=<file>               default ./rendered/<NAME>.yaml
+#   IMAGE=<image>            client image; default keeps the e2e render's
+#   STEPS=3                  optimizer steps
+#   MICRO_STEPS=<n>          fwd_bwd calls per step; default 9 (5 with TRAINER_SLICE=4x8x8), see README
+#   EXTRA_MAXTEXT_FLAGS="k=v ..."      appended to MAXTEXT_EXTRA_FLAGS (last wins)
+#   REPLACE_MAXTEXT_FLAGS="k=v ..."    replace existing MAXTEXT_EXTRA_FLAGS keys in place
+#   SET_TRAINER_FLAGS="k=v ..."        replace existing trainer flags in place (e.g. train_micro_batch_size=128)
+#   SET_ORCH_FLAGS="k=v ..."           replace existing orchestrator flags in place
+#   LIBTPU_EXTRA="--flag=v ..."        add or override libtpu flags (client, proxy and workers)
+#   LIBTPU_DROP="--flag ..."           remove libtpu flags by name
+#   COMPUTE_LOGPS_CHUNK_SIZE=<n>       trainer --compute_logps_chunk_size in place; empty keeps the render's
+#   TRAINER_SLICE=<topology>           e.g. 4x8x8: re-target the trainer slice, FSDP fills it
+#   PROFILER_SKIP=<n> PROFILER_STEPS=<m>  trace fwd_bwd calls n .. n+m-1 of the run (see README)
+#   DRIVER_ENV="K=V ..."               e2e_driver.py knobs: PAYLOAD ROUTING ROUTING_SCALE ROUTING_TARGETS
+#                                      PAD_FRACTION FINAL_PAD_FRACTION PAYLOAD_SEED EP_SIZE ADVANTAGES LOGP_GATES
+set -euo pipefail
+HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+: "${E2E_RENDER:?set E2E_RENDER to the DRY_RUN output of your e2e launch}"
+: "${NAME:?set NAME (lowercase, <=30 chars)}"
+: "${K8S_NAMESPACE:?set K8S_NAMESPACE}"
+: "${OUTPUT_PREFIX:?set OUTPUT_PREFIX (the bench writes to OUTPUT_PREFIX/NAME)}"
+OUT=${OUT:-rendered/$NAME.yaml}
+mkdir -p "$(dirname "$OUT")"
+# Default micro-steps: an example calibration from the 397B recipe, where one optimizer step of about 1,152
+# trajectories packed into 9 micro-batches of 64 rows on a 4x4x8 trainer, and the same token mass into 5 micro-batches
+# of 128 rows on 4x8x8. Set MICRO_STEPS to the fwd_bwd count per step of your own e2e run.
+case "${TRAINER_SLICE:-}" in 4x8x8) _ms=5 ;; *) _ms=9 ;; esac
+python3 "$HERE/render_e2e.py" \
+  --e2e-render "$E2E_RENDER" --driver "$HERE/e2e_driver.py" --out "$OUT" \
+  --name "$NAME" --namespace "$K8S_NAMESPACE" --output-prefix "$OUTPUT_PREFIX" \
+  --steps "${STEPS:-3}" --micro-steps "${MICRO_STEPS:-$_ms}" \
+  --extra-maxtext-flags "${EXTRA_MAXTEXT_FLAGS:-}" --replace-maxtext-flags "${REPLACE_MAXTEXT_FLAGS:-}" \
+  --set-trainer-flags "${SET_TRAINER_FLAGS:-}" --set-orch-flags "${SET_ORCH_FLAGS:-}" \
+  --libtpu-extra="${LIBTPU_EXTRA:-}" --libtpu-drop="${LIBTPU_DROP:-}" \
+  --image "${IMAGE:-}" --driver-env "${DRIVER_ENV:-}" --trainer-slice "${TRAINER_SLICE:-}" \
+  --compute-logps-chunk-size "${COMPUTE_LOGPS_CHUNK_SIZE:-}" \
+  --profiler-skip "${PROFILER_SKIP:-}" --profiler-steps "${PROFILER_STEPS:-}"
diff --git a/src/maxtext/experimental/rl_trainer_bench/structdiff_render.py b/src/maxtext/experimental/rl_trainer_bench/structdiff_render.py
new file mode 100644
index 0000000000..9ce85d4590
--- /dev/null
+++ b/src/maxtext/experimental/rl_trainer_bench/structdiff_render.py
@@ -0,0 +1,95 @@
+#!/usr/bin/env python3
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Structured diff of two render_e2e.py JobSets: every YAML leaf except the client command, plus the client command
+split into the env prefix (MAXTEXT_EXTRA_FLAGS tokens, LIBTPU_INIT_ARGS tokens, other env), the trainer argv
+(E2E_TRAINER_ARGV), the orchestrator argv (E2E_ORCH_ARGV) and the driver env. Prints only differences.
+
+  python3 structdiff_render.py a.yaml b.yaml
+"""
+import shlex
+import sys
+
+import yaml
+
+sys.path.insert(0, __import__("os").path.dirname(__import__("os").path.abspath(__file__)))
+from render_e2e import _proc_main  # pylint: disable=g-import-not-at-top
+
+
+def leaves(x, path=""):
+  if isinstance(x, dict):
+    for k, v in x.items():
+      yield from leaves(v, f"{path}.{k}")
+  elif isinstance(x, list):
+    for i, v in enumerate(x):
+      name = v.get("name") if isinstance(v, dict) and "name" in v else i
+      yield from leaves(v, f"{path}[{name}]")
+  else:
+    yield path, x
+
+
+def flagmap(tokens):
+  out = {}
+  for t in tokens:
+    k, _, v = t.partition("=")
+    out.setdefault(k, []).append(v)
+  return out
+
+
+def parse(path):
+  with open(path) as f:
+    js = yaml.safe_load(f)
+  _, main_c = _proc_main(js)
+  script = main_c["command"][-1]
+  main_c["command"] = ["<client script>"]
+  line = next(l for l in script.splitlines() if "exec python /tmp/e2e_driver.py" in l)
+  toks = shlex.split(line)
+  env = {}
+  for t in toks:
+    if t in ("exec", "python", "/tmp/e2e_driver.py"):
+      continue
+    k, _, v = t.partition("=")
+    env[k] = v
+  sec = {
+      "leaf": dict(leaves(js)),
+      "MAXTEXT_EXTRA_FLAGS": flagmap(env.pop("MAXTEXT_EXTRA_FLAGS", "").split()),
+      "LIBTPU_INIT_ARGS": flagmap(env.pop("LIBTPU_INIT_ARGS", "").split()),
+      "trainer_argv": flagmap(shlex.split(env.pop("E2E_TRAINER_ARGV", ""))),
+      "orch_argv": flagmap(shlex.split(env.pop("E2E_ORCH_ARGV", ""))),
+      "env": env,
+  }
+  return sec
+
+
+def main_args(path_a, path_b):
+  """Prints one line per difference between two renders and returns the number of differences."""
+  a, b = parse(path_a), parse(path_b)
+  n = 0
+  for s in a:
+    for k in sorted(set(a[s]) | set(b[s])):
+      va, vb = a[s].get(k, "<absent>"), b[s].get(k, "<absent>")
+      if va != vb:
+        n += 1
+        print(f"[{s}] {k}: {str(va)[:200]} -> {str(vb)[:200]}")
+  print(f"{n} differences")
+  return n
+
+
+def main():
+  main_args(sys.argv[1], sys.argv[2])
+
+
+if __name__ == "__main__":
+  main()
diff --git a/src/maxtext/experimental/rl_trainer_bench/tests/__init__.py b/src/maxtext/experimental/rl_trainer_bench/tests/__init__.py
new file mode 100644
index 0000000000..f611561ac8
--- /dev/null
+++ b/src/maxtext/experimental/rl_trainer_bench/tests/__init__.py
@@ -0,0 +1,14 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
diff --git a/src/maxtext/experimental/rl_trainer_bench/tests/payload_test.py b/src/maxtext/experimental/rl_trainer_bench/tests/payload_test.py
new file mode 100644
index 0000000000..c164581827
--- /dev/null
+++ b/src/maxtext/experimental/rl_trainer_bench/tests/payload_test.py
@@ -0,0 +1,218 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""CPU-only tests of e2e_driver.py's payload builder and metric helpers (numpy only; no Tunix, no JAX).
+
+The payload is a small stand-in for Tunix's packed RLTrainerPayload with the fields `realistic_payload` reads; it
+checks the same invariants as check_payload_routing.py does at the production shapes with Tunix's own builder.
+"""
+
+import dataclasses
+import os
+import sys
+import types
+import unittest
+
+import numpy as np
+import pytest
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, os.path.dirname(HERE))
+import e2e_driver as d  # pylint: disable=g-import-not-at-top,wrong-import-position
+
+pytestmark = pytest.mark.cpu_only
+
+ROWS, BUDGET, LAYERS, TOP_K, EXPERTS, EP = 4, 4096, 4, 2, 64, 8
+RATIO_KW = dict(ep_rows=2, context=4, chunks=4)
+
+
+@dataclasses.dataclass(frozen=True)
+class FakePayload:
+  """The fields of a packed RLTrainerPayload that realistic_payload reads or writes."""
+
+  completion_ids: np.ndarray
+  segment_ids: np.ndarray
+  segment_positions: np.ndarray
+  completion_mask: np.ndarray
+  advantages: np.ndarray
+  old_per_token_logps: np.ndarray | None
+  rollout_per_token_logps: np.ndarray | None
+  routed_experts: np.ndarray | None
+  num_segments: int = 3
+
+
+def warmup_like(routed=True):
+  """A warm-up-style payload: one short item per row at the start (3 segments), the rest padding, all routing
+  UNSET and every token the same id, as Tunix's TRAINER_WARMUP_COMPILE builder writes it."""
+  seg = np.zeros((ROWS, BUDGET), np.int32)
+  cm = np.zeros((ROWS, BUDGET), np.int32)
+  adv = np.zeros((ROWS, BUDGET), np.float32)
+  for s in range(3):
+    seg[:, s * 100 : (s + 1) * 100] = s + 1
+    cm[:, s * 100 + 40 : (s + 1) * 100] = 1
+    adv[:, s * 100 + 40 : (s + 1) * 100] = 1.0 if s % 2 else -1.0
+  return FakePayload(
+      completion_ids=np.where(seg > 0, 7, 0).astype(np.int32),
+      segment_ids=seg,
+      segment_positions=np.tile(np.arange(BUDGET, dtype=np.int32), (ROWS, 1)),
+      completion_mask=cm,
+      advantages=adv,
+      old_per_token_logps=None,
+      rollout_per_token_logps=np.zeros((ROWS, BUDGET), np.float32),
+      routed_experts=(np.full((ROWS, BUDGET, LAYERS, TOP_K), d.UNSET_ROUTED_EXPERT, np.int16) if routed else None),
+  )
+
+
+def build(**kw):
+  args = dict(num_experts=EXPERTS, ep=EP, pad_id=0, vocab_size=5000, seed=0, pool_size=4096,
+              layer_targets_override=np.array([1.0, 1.3, 1.6, 2.0]))
+  args.update(kw)
+  return d.realistic_payload(warmup_like(), **args)
+
+
+class PayloadTest(unittest.TestCase):
+
+  def test_shapes_dtypes_and_padding(self):
+    base = warmup_like()
+    p = build()
+    for f in dataclasses.fields(base):
+      a, b = getattr(p, f.name), getattr(base, f.name)
+      if isinstance(b, np.ndarray):
+        self.assertEqual((a.shape, a.dtype), (b.shape, b.dtype), f.name)
+    self.assertIsNone(p.old_per_token_logps)
+    self.assertEqual(p.num_segments, base.num_segments)
+    real = p.segment_ids > 0
+    self.assertAlmostEqual(1.0 - real.mean(), 0.07, delta=0.001)
+    self.assertTrue((p.completion_ids[~real] == 0).all())
+    self.assertTrue((p.routed_experts[~real] == d.UNSET_ROUTED_EXPERT).all())
+    self.assertTrue((p.completion_mask[~real] == 0).all() and (p.advantages[~real] == 0).all())
+    # Same number of segments per row as the base, contiguous ids starting at 1.
+    for r in range(ROWS):
+      self.assertEqual(sorted(np.unique(p.segment_ids[r][real[r]])), [1, 2, 3])
+    # Varied ids in range and replayed routing: top-k distinct, in range, on every real token.
+    ids = p.completion_ids[real]
+    self.assertTrue((ids >= 256).all() and (ids < 5000).all() and len(np.unique(ids)) > 1000)
+    r = p.routed_experts[real]
+    self.assertTrue((r >= 0).all() and (r < EXPERTS).all())
+    srt = np.sort(r, axis=-1)
+    self.assertFalse((srt[..., 1:] == srt[..., :-1]).any())
+    # Rollout logps: the synthetic constant on scored tokens, 0 elsewhere.
+    self.assertTrue((p.rollout_per_token_logps[p.completion_mask > 0] == np.float32(-0.7)).all())
+    self.assertTrue((p.rollout_per_token_logps[p.completion_mask == 0] == 0).all())
+
+  def test_routing_follows_targets(self):
+    targets = np.array([1.0, 1.3, 1.6, 2.0])
+    ratio, pad = d.payload_shard_load_ratio(build(), num_experts=EXPERTS, ep=EP, **RATIO_KW)
+    self.assertAlmostEqual(pad, 0.07, delta=0.001)
+    # The replayed share counts padding in the balanced load, so scale by the real fraction before comparing.
+    np.testing.assert_allclose(ratio / 0.93, targets, atol=0.25)
+    self.assertEqual(int(np.argmax(ratio)), 3)
+    uniform, _ = d.payload_shard_load_ratio(build(routing="uniform"), num_experts=EXPERTS, ep=EP, **RATIO_KW)
+    # Small-shape sampling noise lifts the max over chunks above 1; the hot layer still stands well clear of it.
+    self.assertLess(uniform.max() / 0.93, 1.35)
+    self.assertGreater(ratio[3], uniform.max() + 0.4)
+
+  def test_routing_pool_bisection(self):
+    rng = np.random.default_rng(0)
+    for target in (1.25, 1.75, 2.5):
+      pool = d.routing_pool(target, EXPERTS, EP, TOP_K, 8192, rng, hot_shard=3)
+      self.assertAlmostEqual(d._pool_ratio(pool, EXPERTS, EP), target, delta=0.1)  # pylint: disable=protected-access
+      counts = np.bincount((pool // (EXPERTS // EP)).ravel(), minlength=EP)
+      self.assertEqual(int(np.argmax(counts)), 3)
+
+  def test_final_micro_batch_pad_rows(self):
+    p = build(pad_rows_fraction=0.5, seed=1)
+    real = p.segment_ids > 0
+    self.assertFalse(real[2:].any())  # the last half of the rows are entirely padding
+    self.assertAlmostEqual(real[:2].mean(), 0.93, delta=0.001)
+    self.assertTrue((p.routed_experts[2:] == d.UNSET_ROUTED_EXPERT).all())
+
+  def test_advantages_random_keeps_everything_else(self):
+    a, b = build(), build(advantages="random")
+    for f in ("completion_ids", "segment_ids", "segment_positions", "completion_mask", "routed_experts",
+              "rollout_per_token_logps"):
+      np.testing.assert_array_equal(getattr(a, f), getattr(b, f), err_msg=f)
+    scored = a.completion_mask > 0
+    self.assertEqual(set(np.unique(a.advantages[scored])), {-1.0, 1.0})  # inherited from the base
+    self.assertGreater(len(np.unique(b.advantages[scored])), 3)
+    # One advantage per segment.
+    for r in range(ROWS):
+      for s in (1, 2, 3):
+        vals = np.unique(b.advantages[r][(b.segment_ids[r] == s) & scored[r]])
+        self.assertEqual(len(vals), 1)
+    with self.assertRaises(ValueError):
+      build(advantages="bogus")
+
+  def test_no_routed_experts(self):
+    p = d.realistic_payload(warmup_like(routed=False), num_experts=EXPERTS, ep=EP, seed=0, vocab_size=5000)
+    self.assertIsNone(p.routed_experts)
+
+  def test_seeded(self):
+    np.testing.assert_array_equal(build().routed_experts, build().routed_experts)
+    self.assertFalse(np.array_equal(build().completion_ids, build(seed=5).completion_ids))
+
+
+class CalibrationTest(unittest.TestCase):
+
+  def test_probe_index_to_layer_is_a_permutation(self):
+    layers = [d.probe_index_to_layer(k) for k in range(60)]
+    self.assertEqual(sorted(layers), list(range(60)))
+    self.assertEqual(d.probe_index_to_layer(59), 59)
+    self.assertEqual(d.probe_index_to_layer(1), 4)
+
+  def test_layer_targets(self):
+    t = d.layer_targets(60)
+    self.assertEqual(t.shape, (60,))
+    self.assertAlmostEqual(t.max(), max(d.PROBE_REGULAR_MEAN_RATIO))
+    np.testing.assert_allclose(d.layer_targets(60, scale=2.0) - 1.0, 2.0 * (t - 1.0))
+    self.assertEqual(d.layer_targets(8).shape, (8,))  # another depth cycles the profile
+    own = tuple(1.0 + 0.01 * k for k in range(60))
+    t_own = d.layer_targets(60, probe=own)
+    self.assertAlmostEqual(t_own[d.probe_index_to_layer(10)], 1.10)
+
+  def test_parse_routing_targets(self):
+    self.assertEqual(d.parse_routing_targets("1.5, 2.0,1.0", 3), (1.5, 2.0, 1.0))
+    with self.assertRaises(ValueError):
+      d.parse_routing_targets("1.5,2.0", 3)
+    with self.assertRaises(ValueError):
+      d.parse_routing_targets("0.5,2.0", 2)
+
+
+class HelpersTest(unittest.TestCase):
+
+  def test_strip_logp_gates(self):
+    argv = ["--a=1", "--seq_logprob_error_threshold=2.0", "--truncated_importance_sampling_type", "seq-mask-tis",
+            "--truncated_importance_sampling_ratio_min=0.999", "--truncated_importance_sampling_ratio=1.002", "--b=2"]
+    self.assertEqual(d.strip_logp_gates(argv), ["--a=1", "--b=2"])
+
+  def test_summarize_metrics(self):
+    wm = lambda s, den, min_denom=None: types.SimpleNamespace(  # pylint: disable=unnecessary-lambda-assignment
+        unreduced_sum=np.array(s), denominator=np.array(den), eps=None, min_denom=min_denom)
+    buf = types.SimpleNamespace(
+        weighted_metrics={"loss": wm([2.0, 0.0, -3.0], [4.0, 0.0, 6.0]), "kept": wm([0.0], [0.0], min_denom=1.0)},
+        scalar_metrics={"gradient_norm": np.array([0.125]), "is_ratio/mean": np.array([1.0, 1.0, 1.0])},
+    )
+    s = d.summarize_metrics(buf)
+    self.assertAlmostEqual(s["loss"], -0.1)
+    self.assertEqual(s["micro_denominator"], [4.0, 0.0, 6.0])
+    self.assertTrue(np.isnan(s["micro_loss"][1]))
+    self.assertEqual(s["weighted"]["kept"], 0.0)  # min_denom keeps an empty metric finite
+    self.assertEqual(s["scalar"], {"gradient_norm": 0.125, "is_ratio/mean": [1.0, 1.0, 1.0]})
+    loss_line, metrics_line = d.metric_log_lines(3, s)
+    self.assertTrue(loss_line.startswith("STEP_LOSS step=3 loss=-0.1 grad_norm=0.125 micro_loss=[0.5,nan,-0.5]"))
+    self.assertIn("is_ratio/mean=[1,1,1]", metrics_line)
+
+
+if __name__ == "__main__":
+  unittest.main()
diff --git a/src/maxtext/experimental/rl_trainer_bench/tests/render_e2e_test.py b/src/maxtext/experimental/rl_trainer_bench/tests/render_e2e_test.py
new file mode 100644
index 0000000000..24e2153ce1
--- /dev/null
+++ b/src/maxtext/experimental/rl_trainer_bench/tests/render_e2e_test.py
@@ -0,0 +1,204 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""CPU-only tests of render_e2e.py on a hand-written DRY_RUN fixture (no cluster, no Tunix, no JAX)."""
+
+import io
+import os
+import re
+import shlex
+import sys
+import tempfile
+import unittest
+from contextlib import redirect_stdout
+
+import pytest
+import yaml
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+PKG = os.path.dirname(HERE)
+sys.path.insert(0, PKG)
+import render_e2e as R  # pylint: disable=g-import-not-at-top,wrong-import-position
+import structdiff_render as SD  # pylint: disable=g-import-not-at-top,wrong-import-position
+
+FIXTURE = os.path.join(HERE, "testdata", "e2e_dryrun_fixture.yaml")
+DRIVER = os.path.join(PKG, "e2e_driver.py")
+
+pytestmark = pytest.mark.cpu_only
+
+
+def _client_env(js):
+  """The driver's inline env (E2E_TRAINER_ARGV, E2E_ORCH_ARGV, MAXTEXT_EXTRA_FLAGS, ...) of a rendered JobSet."""
+  _, main_c = R._proc_main(js)  # pylint: disable=protected-access
+  line = next(l for l in main_c["command"][-1].splitlines() if "exec python /tmp/e2e_driver.py" in l)
+  env = {}
+  for t in shlex.split(line):
+    k, eq, v = t.partition("=")
+    if eq:
+      env[k] = v
+  return env
+
+
+class RenderE2ETest(unittest.TestCase):
+
+  def setUp(self):
+    self.tmp = tempfile.TemporaryDirectory()  # pylint: disable=consider-using-with
+
+  def tearDown(self):
+    self.tmp.cleanup()
+
+  def render(self, *extra, name="tb-test"):
+    out = os.path.join(self.tmp.name, f"{name}.yaml")
+    argv = ["--e2e-render", FIXTURE, "--driver", DRIVER, "--out", out, "--name", name, "--namespace", "bench-ns",
+            "--output-prefix", "/bench-out", *extra]
+    with redirect_stdout(io.StringIO()):
+      self.assertEqual(R.main(argv), 0)
+    with open(out) as f:
+      return yaml.safe_load(f), out
+
+  def test_selects_trainer_and_orchestrator(self):
+    docs = R.load_documents(FIXTURE)
+    self.assertEqual([d["metadata"]["name"] for d in docs], ["example-orch", "example-train", "example-roll-0"])
+    train, orch = R.select_jobsets(docs)
+    self.assertEqual(train["metadata"]["name"], "example-train")
+    self.assertEqual(orch["metadata"]["name"], "example-orch")
+    with self.assertRaises(ValueError):
+      R.select_jobsets(docs + [train])  # two trainers
+
+  def test_base_render(self):
+    js, _ = self.render()
+    md = js["metadata"]
+    self.assertEqual((md["name"], md["namespace"]), ("tb-test", "bench-ns"))
+    self.assertEqual(md["labels"]["scheduling.x-k8s.io/gang-id"], "tb-test")
+    self.assertEqual(md["labels"]["kueue.x-k8s.io/queue-name"], "example-queue")  # scheduling fields kept
+    self.assertEqual(js["spec"]["failurePolicy"]["maxRestarts"], 0)
+    proc, main_c = R._proc_main(js)  # pylint: disable=protected-access
+    self.assertNotIn("ports", main_c)
+    self.assertEqual(main_c["image"], "example.invalid/tunix:tag")
+    script = main_c["command"][-1]
+    self.assertIn("cat > /tmp/e2e_driver.py <<'E2E_DRIVER_EOF'", script)
+    self.assertNotIn("FT_REGISTERED_MARKER", script)  # Tunix fail-fast wrapper dropped
+    with open(DRIVER) as f:
+      self.assertIn(f.read(), script)
+
+    env = _client_env(js)
+    targv = shlex.split(env["E2E_TRAINER_ARGV"])
+    self.assertIn("--maxtext_ckpt_path=", targv)
+    self.assertIn("--maxtext_output_directory=/bench-out/tb-test", targv)
+    self.assertIn("--worker_id=tb-test", targv)
+    self.assertFalse([t for t in targv if t.startswith(("--port=", "--discovery_addrs=", "--process_main="))])
+    self.assertEqual([t for t in targv if t.startswith("--mesh_fsdp=")], ["--mesh_fsdp=32"])
+    oargv = shlex.split(env["E2E_ORCH_ARGV"])
+    self.assertIn("--trainer_fsdp=32", oargv)
+    self.assertIn("--seq_logprob_error_threshold=2.0", oargv)
+    self.assertEqual((env["STEPS"], env["MICRO_STEPS"]), ("3", "4"))
+    self.assertEqual(env["MAXTEXT_EXTRA_FLAGS"], "context=offload num_moe_token_chunks=4 packing=True")
+    # Without libtpu knobs, the proxy and worker flags are untouched.
+    proxy = R.find(proc["initContainers"], "pathways-proxy")
+    self.assertEqual(proxy["args"][-1], "--sidecar_name=external")
+    self.assertIn("--xla_tpu_flag_b=4", proxy["args"])
+
+  def test_libtpu_knobs_reach_all_three_places(self):
+    js, _ = self.render("--libtpu-extra=--xla_tpu_flag_b=8 --xla_new=1", "--libtpu-drop=--xla_tpu_flag_a")
+    env = _client_env(js)
+    self.assertEqual(env["LIBTPU_INIT_ARGS"].split(), ["--RUNTIME_only_flag=true", "--xla_tpu_flag_b=8", "--xla_new=1"])
+    proc, _ = R._proc_main(js)  # pylint: disable=protected-access
+    proxy = R.find(proc["initContainers"], "pathways-proxy")["args"]
+    self.assertEqual([a for a in proxy if a.startswith("--xla_")], ["--xla_tpu_flag_b=8", "--xla_new=1"])
+    self.assertEqual(proxy[-1], "--sidecar_name=external")
+    self.assertFalse([a for a in proxy if a.startswith("--RUNTIME")])  # runtime-only flags stay off the proxy
+    pw = R.find(js["spec"]["replicatedJobs"], "pw-node")["template"]["spec"]["template"]["spec"]
+    wenv = R.find(R.find(pw["containers"], "pathways-worker")["env"], "LIBTPU_INIT_ARGS")["value"].split()
+    self.assertEqual(wenv, ["--RUNTIME_only_flag=true", "--xla_tpu_flag_b=8", "--xla_new=1", "--worker_only_flag=1"])
+
+  def test_maxtext_flag_knobs(self):
+    js, _ = self.render("--replace-maxtext-flags", "context=remat", "--extra-maxtext-flags", "log_moe_drops=true")
+    self.assertEqual(_client_env(js)["MAXTEXT_EXTRA_FLAGS"],
+                     "context=remat num_moe_token_chunks=4 packing=True log_moe_drops=true")
+    with self.assertRaises(ValueError):
+      self.render("--replace-maxtext-flags", "not_a_key=1")
+
+  def test_trainer_and_orch_flag_knobs(self):
+    js, _ = self.render("--set-trainer-flags", "train_micro_batch_size=128", "--set-orch-flags",
+                        "train_micro_batch_size=128", "--compute-logps-chunk-size", "1024",
+                        "--profiler-skip", "16", "--profiler-steps", "3")
+    env = _client_env(js)
+    targv = shlex.split(env["E2E_TRAINER_ARGV"])
+    for flag in ("--train_micro_batch_size=128", "--compute_logps_chunk_size=1024", "--profiler_steps=3",
+                 "--skip_first_n_profiler_steps=16"):
+      self.assertEqual(sum(t.split("=", 1)[0] == flag.split("=", 1)[0] for t in targv), 1, flag)
+      self.assertIn(flag, targv)
+    self.assertIn("--train_micro_batch_size=128", shlex.split(env["E2E_ORCH_ARGV"]))
+    with self.assertRaises(ValueError):
+      self.render("--set-trainer-flags", "no_such_flag=1")
+    with self.assertRaises(ValueError):
+      self.render("--set-orch-flags", "no_such_flag=1")
+
+  def test_trainer_slice(self):
+    base, base_path = self.render(name="tb-a")
+    js, path = self.render("--trainer-slice", "4x8x8", name="tb-b")
+    proc, _ = R._proc_main(js)  # pylint: disable=protected-access
+    rm = R.find(proc["initContainers"], "pathways-rm")["args"]
+    self.assertIn("--instance_type=tpu7x:4x8x8", rm)
+    pw = R.find(js["spec"]["replicatedJobs"], "pw-node")["template"]["spec"]
+    self.assertEqual((pw["completions"], pw["parallelism"]), (64, 64))
+    self.assertEqual(pw["template"]["metadata"]["annotations"]["cloud.google.com/gke-tpu-slice-topology"], "4x8x8")
+    env = _client_env(js)
+    self.assertIn("--mesh_fsdp=64", shlex.split(env["E2E_TRAINER_ARGV"]))
+    self.assertIn("--trainer_fsdp=64", shlex.split(env["E2E_ORCH_ARGV"]))
+    # Only the documented fields (plus the names) differ from the 4x4x8 render.
+    buf = io.StringIO()
+    with redirect_stdout(buf):
+      SD.main_args(base_path, path)
+    diffs = [l for l in buf.getvalue().splitlines() if l.startswith("[")]
+    keys = sorted(re.match(r"\[(\w+)\] (\S+?):", l).group(1, 2) for l in diffs)
+    self.assertEqual(keys, sorted([
+        ("leaf", ".metadata.name"), ("leaf", ".metadata.labels.scheduling.x-k8s.io/gang-id"),
+        ("leaf", ".spec.replicatedJobs[proc].template.spec.template.spec.initContainers[pathways-rm].args[3]"),
+        ("leaf", ".spec.replicatedJobs[pw-node].template.spec.completions"),
+        ("leaf", ".spec.replicatedJobs[pw-node].template.spec.parallelism"),
+        ("leaf", ".spec.replicatedJobs[pw-node].template.spec.template.metadata.annotations."
+                 "cloud.google.com/gke-tpu-slice-topology"),
+        ("trainer_argv", "--maxtext_output_directory"), ("trainer_argv", "--worker_id"),
+        ("trainer_argv", "--mesh_fsdp"), ("orch_argv", "--trainer_fsdp"),
+    ]))
+    del base
+
+  def test_driver_env_and_image(self):
+    js, _ = self.render("--driver-env", "PAYLOAD=realistic ROUTING_TARGETS=1.5,2.0", "--image", "example.invalid/x:y",
+                        "--steps", "2", "--micro-steps", "5")
+    env = _client_env(js)
+    self.assertEqual((env["PAYLOAD"], env["ROUTING_TARGETS"], env["STEPS"], env["MICRO_STEPS"]),
+                     ("realistic", "1.5,2.0", "2", "5"))
+    self.assertEqual(R._proc_main(js)[1]["image"], "example.invalid/x:y")  # pylint: disable=protected-access
+    with self.assertRaises(ValueError):
+      self.render("--driver-env", "lower=1")
+
+  def test_print_aot_env(self):
+    buf = io.StringIO()
+    with redirect_stdout(buf):
+      R.main(["--print-aot-env", "--e2e-render", FIXTURE, "--trainer-slice", "4x8x8", "--libtpu-extra=--xla_new=1"])
+    out = buf.getvalue()
+    self.assertIn("export MAXTEXT_EXTRA_FLAGS='context=offload num_moe_token_chunks=4 packing=True'", out)
+    self.assertIn("export LIBTPU_INIT_ARGS='--xla_tpu_flag_a=true --xla_tpu_flag_b=4 --xla_new=1'", out)
+    self.assertIn("export AOT_MESH_FSDP=64", out)
+    self.assertIn("export AOT_LOGPS_CHUNK_SIZE=512", out)
+
+  def test_name_validation(self):
+    with self.assertRaises(SystemExit):
+      self.render(name="Bad_Name")
+
+
+if __name__ == "__main__":
+  unittest.main()
diff --git a/src/maxtext/experimental/rl_trainer_bench/tests/testdata/e2e_dryrun_fixture.yaml b/src/maxtext/experimental/rl_trainer_bench/tests/testdata/e2e_dryrun_fixture.yaml
new file mode 100644
index 0000000000..4673b7a0f8
--- /dev/null
+++ b/src/maxtext/experimental/rl_trainer_bench/tests/testdata/e2e_dryrun_fixture.yaml
@@ -0,0 +1,167 @@
+launcher: this line and the next are log output that a DRY_RUN may interleave with the YAML documents
+kubectl delete jobset example-orch -n example-ns
+---
+# Hand-written fixture in the shape of a Tunix e2e DRY_RUN render (orchestrator, trainer, one rollout JobSet).
+apiVersion: jobset.x-k8s.io/v1alpha2
+kind: JobSet
+metadata:
+  name: example-orch
+  namespace: example-ns
+  labels:
+    kueue.x-k8s.io/queue-name: example-queue
+    scheduling.x-k8s.io/gang-id: "example"
+spec:
+  failurePolicy:
+    maxRestarts: 0
+  replicatedJobs:
+  - name: proc
+    replicas: 1
+    template:
+      spec:
+        completions: 1
+        parallelism: 1
+        template:
+          spec:
+            containers:
+            - name: main
+              image: example.invalid/tunix:tag
+              command:
+              - bash
+              - -c
+              - |
+                echo "Program Start: $(date)"
+                (
+                  PYTHONUNBUFFERED=1 python -m tunix.experimental.distributed.runtime.main \
+                    --process_executor=tunix.experimental.distributed.runtime.executor.K8sExecutor \
+                    --process_main=tunix.experimental.examples.deepswe_dist.run_deepswe_dist.main \
+                    --num_generations=16 \
+                    --mini_batch_size=72 \
+                    --train_micro_batch_size=64 \
+                    --trainer_fsdp=32 \
+                    --trainer_dp=1 \
+                    --trainer_expert=2 \
+                    --max_prompt_length=8192 \
+                    --max_response_length=57344 \
+                    --seq_logprob_error_threshold=2.0 \
+                    --truncated_importance_sampling_type=seq-mask-tis
+                ) & PID=$!
+                wait $PID
+---
+apiVersion: jobset.x-k8s.io/v1alpha2
+kind: JobSet
+metadata:
+  name: example-train
+  namespace: example-ns
+  labels:
+    kueue.x-k8s.io/queue-name: example-queue
+    scheduling.x-k8s.io/gang-id: "example"
+spec:
+  failurePolicy:
+    restartStrategy: Recreate
+    maxRestarts: 3
+  replicatedJobs:
+  - name: proc
+    replicas: 1
+    template:
+      spec:
+        completions: 1
+        parallelism: 1
+        template:
+          spec:
+            initContainers:
+            - name: pathways-rm
+              image: example.invalid/pathways-server:tag
+              restartPolicy: Always
+              args:
+              - --server_port=29001
+              - --node_type=resource_manager
+              - --instance_count=1
+              - --instance_type=tpu7x:4x4x8
+            - name: pathways-proxy
+              image: example.invalid/pathways-proxy:tag
+              restartPolicy: Always
+              args:
+              - --server_port=29000
+              - --resource_manager_address=localhost:29001
+              - --xla_tpu_flag_a=true
+              - --xla_tpu_flag_b=4
+              - --sidecar_name=external
+            containers:
+            - name: main
+              image: example.invalid/tunix:tag
+              ports:
+              - containerPort: 20002
+                name: grpc
+              command:
+              - bash
+              - -c
+              - |
+                echo "Program Start: $(date)"
+                export FT_REGISTERED_MARKER=/dev/shm/registered
+                (
+                  PYTHONUNBUFFERED=1 MAXTEXT_EXTRA_FLAGS="context=offload num_moe_token_chunks=4 packing=True" LIBTPU_INIT_ARGS='--RUNTIME_only_flag=true --xla_tpu_flag_a=true --xla_tpu_flag_b=4' python -m tunix.experimental.distributed.runtime.main \
+                    --discovery_addrs=example-orch:20000 \
+                    --process_executor=tunix.experimental.distributed.runtime.executor.K8sExecutor \
+                    --process_main=tunix.experimental.examples.common.run_trainer_node.main \
+                    --worker_id=example-train \
+                    --port=20002 \
+                    --mesh_fsdp=32 \
+                    --mesh_tp=1 \
+                    --mesh_expert=2 \
+                    --mesh_context=4 \
+                    --maxtext_ckpt_path=/ckpt/example \
+                    --maxtext_output_directory=/out/example \
+                    --compute_logps_chunk_size=512 \
+                    --train_micro_batch_size=64 \
+                    --profiler_steps=1 \
+                    --skip_first_n_profiler_steps=3
+                ) & PID=$!
+                wait $PID
+                if [ ! -f "$FT_REGISTERED_MARKER" ]; then exit 75; fi
+  - name: pw-node
+    replicas: 1
+    template:
+      spec:
+        completions: 32
+        parallelism: 32
+        template:
+          metadata:
+            annotations:
+              cloud.google.com/gke-tpu-slice-topology: 4x4x8
+          spec:
+            containers:
+            - name: pathways-worker
+              image: example.invalid/pathways-server:tag
+              args:
+              - --server_port=29005
+              resources:
+                limits:
+                  google.com/tpu: "4"
+              env:
+              - name: LIBTPU_INIT_ARGS
+                value: --RUNTIME_only_flag=true --xla_tpu_flag_a=true --xla_tpu_flag_b=4 --worker_only_flag=1
+---
+apiVersion: jobset.x-k8s.io/v1alpha2
+kind: JobSet
+metadata:
+  name: example-roll-0
+  namespace: example-ns
+spec:
+  replicatedJobs:
+  - name: proc
+    replicas: 1
+    template:
+      spec:
+        template:
+          spec:
+            containers:
+            - name: main
+              image: example.invalid/rollout:tag
+              command:
+              - bash
+              - -c
+              - |
+                (
+                  python -m tunix.experimental.distributed.runtime.main \
+                    --process_main=tunix.experimental.examples.common.run_rollout_node.main
+                )

From d6db22f99df91cdf53cbf6a8e9552feb7c7aa36e Mon Sep 17 00:00:00 2001
From: Vaibhav Singh <sivaibhav@google.com>
Date: Tue, 13 Oct 2026 18:51:05 +0000
Subject: [PATCH 2/4] rl_trainer_bench: label example numbers, fail fast in
 aot_e2e.sh, note DRY_RUN context switch

---
 src/maxtext/experimental/rl_trainer_bench/README.md     | 9 +++++----
 src/maxtext/experimental/rl_trainer_bench/aot_e2e.sh    | 5 +++--
 src/maxtext/experimental/rl_trainer_bench/render_e2e.py | 2 ++
 3 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/src/maxtext/experimental/rl_trainer_bench/README.md b/src/maxtext/experimental/rl_trainer_bench/README.md
index a8d468a70e..a13b8412fe 100644
--- a/src/maxtext/experimental/rl_trainer_bench/README.md
+++ b/src/maxtext/experimental/rl_trainer_bench/README.md
@@ -76,7 +76,8 @@ The output is a stream of JobSets (orchestrator, trainer, one per rollout engine
 launcher log lines. The renderer picks the trainer and the orchestrator JobSet by the `--process_main=` of their
 `main` container and ignores everything else. If you already have the two JobSets as separate files, pass them with
 `--train-yaml` and `--orch-yaml` instead. The DRY_RUN output describes your environment (namespace, images,
-buckets), so keep it local.
+buckets), so keep it local. In our runs the dry run still switched the current kubectl context and printed (without
+executing) the delete commands of a stop, so check your context afterwards.
 
 ### 2. Render the bench JobSet
 
@@ -100,7 +101,7 @@ The knobs, all environment variables of `render_e2e.sh` (the Python script takes
 `IMAGE` replaces the client image; the Pathways containers keep the e2e render's. `STEPS` (default 3) is the number of
 optimizer steps and `MICRO_STEPS` the `fwd_bwd` calls per step. Set `MICRO_STEPS` to what your e2e run does per
 optimizer step (the orchestrator logs one train call per micro-batch); the default of 9 (5 with
-`TRAINER_SLICE=4x8x8`) is an example calibration from the 397B recipe, where about 1,152 trajectories per step packed
+`TRAINER_SLICE=4x8x8`) is an example calibration from our 397B runs, where about 1,152 trajectories per step packed
 into 9 micro-batches of 64 rows on a 4x4x8 trainer.
 
 `EXTRA_MAXTEXT_FLAGS="k=v ..."` appends MaxText flags (MaxText keys are last-wins), and
@@ -176,7 +177,7 @@ The profile goes under `<output-prefix>/<NAME>/tensorboard/` in the usual `plugi
 
 Tunix's TRAINER_WARMUP_COMPILE payload fixes the shapes, dtypes, fields and per-row segment count of the packed
 micro-batch, which is what the compile needs. As a timing payload it is degenerate for a MoE model: at the 397B shapes
-it is about 97% padding, every token carries the same id and every routed expert is UNSET, so with router replay every
+it was about 97% padding in our runs (measured with `check_payload_routing.py`), every token carries the same id and every routed expert is UNSET, so with router replay every
 token falls back to the router with identical inputs and lands on the same experts.
 
 `PAYLOAD=realistic` (the default) keeps the warm-up payload's structure and rebuilds its contents with
@@ -198,7 +199,7 @@ padded, since the e2e run's last micro-batch of a step is a partially filled tai
 the final one. Before the steps the driver logs one `PAYLOAD_ROUTING` line per payload with the offline per-layer
 ratio of the replayed tokens; the runtime ratio, which includes the padding's router-chosen experts, is in the
 engine's `MOE_DROP_AUDIT` lines with the MaxText flag `log_moe_drops=true`. Building the two payloads at the 397B
-shapes takes a couple of minutes of host time before the first step. `PAYLOAD=warmup` restores the warm-up payload.
+shapes took about two minutes each of host time in our runs, before the first step. `PAYLOAD=warmup` restores the warm-up payload.
 
 ## Numerics mode
 
diff --git a/src/maxtext/experimental/rl_trainer_bench/aot_e2e.sh b/src/maxtext/experimental/rl_trainer_bench/aot_e2e.sh
index 83d11343a4..fa3f596224 100755
--- a/src/maxtext/experimental/rl_trainer_bench/aot_e2e.sh
+++ b/src/maxtext/experimental/rl_trainer_bench/aot_e2e.sh
@@ -27,10 +27,11 @@ set -euo pipefail
 HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
 : "${E2E_RENDER:?set E2E_RENDER to the DRY_RUN output of your e2e launch}"
 : "${AOT_IMAGE:?set AOT_IMAGE to an image with the MaxText and Tunix of the e2e run}"
-eval "$(python3 "$HERE/render_e2e.py" --print-aot-env --e2e-render "$E2E_RENDER" \
+_aot_env=$(python3 "$HERE/render_e2e.py" --print-aot-env --e2e-render "$E2E_RENDER" \
   --extra-maxtext-flags "${EXTRA_MAXTEXT_FLAGS:-}" --replace-maxtext-flags "${REPLACE_MAXTEXT_FLAGS:-}" \
   --libtpu-extra="${LIBTPU_EXTRA:-}" --libtpu-drop="${LIBTPU_DROP:-}" --trainer-slice "${TRAINER_SLICE:-}" \
-  --compute-logps-chunk-size "${COMPUTE_LOGPS_CHUNK_SIZE:-}")"
+  --compute-logps-chunk-size "${COMPUTE_LOGPS_CHUNK_SIZE:-}")
+eval "$_aot_env"
 case "${TRAINER_SLICE:-}" in 4x8x8) _topo=tpu7x-512 ;; *) _topo=tpu7x-256 ;; esac
 TOPO=${AOT_COMPILE_TOPOLOGY:-$_topo}
 LOG=${LOG:-logs/aot-e2e-$(date -u +%Y%m%dT%H%M%SZ).log}
diff --git a/src/maxtext/experimental/rl_trainer_bench/render_e2e.py b/src/maxtext/experimental/rl_trainer_bench/render_e2e.py
index 56a62f4988..1646c2cc55 100644
--- a/src/maxtext/experimental/rl_trainer_bench/render_e2e.py
+++ b/src/maxtext/experimental/rl_trainer_bench/render_e2e.py
@@ -45,6 +45,8 @@
 trainer's MaxText and XLA flags and mesh for aot_e2e.sh instead of rendering.
 """
 
+from __future__ import annotations
+
 import argparse
 import copy
 import os

From 70d64441eba957fe5fe934cc4346a26df4706719 Mon Sep 17 00:00:00 2001
From: Vaibhav Singh <sivaibhav@google.com>
Date: Tue, 13 Oct 2026 23:02:07 +0000
Subject: [PATCH 3/4] rl_trainer_bench: rename e2e identifiers to bench/launch
 names

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.
---
 .../experimental/rl_trainer_bench/README.md   | 38 +++++++++--------
 .../{aot_e2e.sh => aot_bench.sh}              | 16 +++----
 .../rl_trainer_bench/aot_summarize.py         |  2 +-
 .../rl_trainer_bench/aot_trainer.py           |  4 +-
 .../{e2e_driver.py => bench_driver.py}        |  6 +--
 .../rl_trainer_bench/check_numerics_cpu.py    | 14 +++----
 .../rl_trainer_bench/check_payload_routing.py |  4 +-
 .../{render_e2e.py => render_bench.py}        | 42 +++++++++----------
 .../{render_e2e.sh => render_bench.sh}        | 12 +++---
 .../rl_trainer_bench/structdiff_render.py     | 14 +++----
 .../rl_trainer_bench/tests/payload_test.py    |  4 +-
 ...ender_e2e_test.py => render_bench_test.py} | 32 +++++++-------
 ...ixture.yaml => launch_dryrun_fixture.yaml} |  0
 13 files changed, 95 insertions(+), 93 deletions(-)
 rename src/maxtext/experimental/rl_trainer_bench/{aot_e2e.sh => aot_bench.sh} (76%)
 rename src/maxtext/experimental/rl_trainer_bench/{e2e_driver.py => bench_driver.py} (99%)
 rename src/maxtext/experimental/rl_trainer_bench/{render_e2e.py => render_bench.py} (93%)
 rename src/maxtext/experimental/rl_trainer_bench/{render_e2e.sh => render_bench.sh} (87%)
 rename src/maxtext/experimental/rl_trainer_bench/tests/{render_e2e_test.py => render_bench_test.py} (89%)
 rename src/maxtext/experimental/rl_trainer_bench/tests/testdata/{e2e_dryrun_fixture.yaml => launch_dryrun_fixture.yaml} (100%)

diff --git a/src/maxtext/experimental/rl_trainer_bench/README.md b/src/maxtext/experimental/rl_trainer_bench/README.md
index a13b8412fe..f013a3f0d4 100644
--- a/src/maxtext/experimental/rl_trainer_bench/README.md
+++ b/src/maxtext/experimental/rl_trainer_bench/README.md
@@ -21,7 +21,9 @@ program the e2e run compiles (the same MaxText flags, libtpu flags, mesh, GRPO l
 the same compile key) on one trainer slice, under the same Pathways setup, with a synthetic micro-batch whose MoE
 routing load is calibrated to a production run. There is no orchestrator, no rollout engines and no weight sync. A
 variant (a MaxText flag, a libtpu flag, a slice size) can therefore be A/B tested in one short job on one slice, and
-its profile read, without the rest of the RL system.
+its profile read, without the rest of the RL system. The bench is standalone and launches no e2e job: its input,
+`LAUNCH_RENDER`, is the DRY_RUN output of the full launcher, from which the renderer takes only the trainer JobSet and
+the trainer and orchestrator argv.
 
 ## Why a separate harness
 
@@ -31,7 +33,7 @@ loss and its vocabulary-chunked `compute_per_token_logps` pass are missing, the
 optimizer step is driven by MaxText rather than by Tunix's `TrainerWorker`. Timings from it do not transfer to the RL
 step, and a variant that compiles there may not compile in the RL program (or the other way round).
 
-The bench therefore drives the trainer through Tunix's own code paths. `e2e_driver.py` builds the trainer worker with
+The bench therefore drives the trainer through Tunix's own code paths. `bench_driver.py` builds the trainer worker with
 `run_trainer_node._parse_args` and `_create_maxtext_trainer_factory` from the e2e trainer's argv, builds the
 `GRPOAdapter`, `BatchConfig` and batch assembler with `run_deepswe_dist._parse_args` and `_build_algo` from the e2e
 orchestrator's argv, applies `configure_worker`'s loss and model-input functions, compiles against Tunix's
@@ -51,11 +53,11 @@ to the team.
 
 | file | what it does |
 |---|---|
-| `e2e_driver.py` | the bench itself; runs as the Pathways client of the trainer slice |
-| `render_e2e.py`, `render_e2e.sh` | turn an e2e DRY_RUN render into the bench JobSet, with variant knobs |
+| `bench_driver.py` | the bench itself; runs as the Pathways client of the trainer slice |
+| `render_bench.py`, `render_bench.sh` | turn an e2e DRY_RUN render into the bench JobSet, with variant knobs |
 | `structdiff_render.py` | structured diff of two bench renders (YAML leaves, MaxText flags, libtpu flags, argv) |
 | `diff_config.py` | diff of the resolved MaxText `Config param` lines of two logs (e2e trainer vs bench) |
-| `aot_trainer.py`, `aot_e2e.sh`, `aot_summarize.py` | ahead-of-time compile and memory report on a CPU host |
+| `aot_trainer.py`, `aot_bench.sh`, `aot_summarize.py` | ahead-of-time compile and memory report on a CPU host |
 | `check_payload_routing.py` | CPU check of the realistic payload at the 397B production shapes (needs Tunix) |
 | `check_numerics_cpu.py` | CPU check of the numerics knobs with Tunix's GRPO loss (needs Tunix and MaxText) |
 | `hlo_collective_schedule.py` | static overlap summary of the collectives in a scheduled HLO dump |
@@ -69,7 +71,7 @@ Start from the launch you want to reproduce, with the environment you would laun
 for a dry run. For the MLPerf recipe in Tunix this is
 
 ```bash
-DRY_RUN=true bash tunix/experimental/examples/recipes/mlperf_gbs1024_submission.sh start > e2e-dryrun.yaml
+DRY_RUN=true bash tunix/experimental/examples/recipes/mlperf_gbs1024_submission.sh start > launch-dryrun.yaml
 ```
 
 The output is a stream of JobSets (orchestrator, trainer, one per rollout engine), possibly interleaved with
@@ -82,21 +84,21 @@ executing) the delete commands of a stop, so check your context afterwards.
 ### 2. Render the bench JobSet
 
 ```bash
-E2E_RENDER=e2e-dryrun.yaml NAME=tb-base K8S_NAMESPACE=<namespace> OUTPUT_PREFIX=<output-prefix> \
-  bash render_e2e.sh                     # writes rendered/tb-base.yaml and nothing else
+LAUNCH_RENDER=launch-dryrun.yaml NAME=tb-base K8S_NAMESPACE=<namespace> OUTPUT_PREFIX=<output-prefix> \
+  bash render_bench.sh                   # writes rendered/tb-base.yaml and nothing else
 ```
 
 The render keeps every scheduling field of the e2e trainer JobSet (Kueue labels and annotations, node selectors,
 tolerations, reservations, priority class, service account, the Pathways resource manager, proxy, worker and sidecar
 containers with their images, args and env). It changes the names, sets `failurePolicy.maxRestarts` to 0 (so a
 failure does not rerun the compile), and replaces the client command: the same inline env prefix as the e2e trainer
-(MAXTEXT_EXTRA_FLAGS, LIBTPU_INIT_ARGS and the rest) runs `e2e_driver.py`, written into the pod by the command itself,
-with the e2e trainer argv and orchestrator argv as `E2E_TRAINER_ARGV` and `E2E_ORCH_ARGV`. In the trainer argv,
+(MAXTEXT_EXTRA_FLAGS, LIBTPU_INIT_ARGS and the rest) runs `bench_driver.py`, written into the pod by the command itself,
+with the e2e trainer argv and orchestrator argv as `TRAINER_ARGV` and `ORCH_ARGV`. In the trainer argv,
 `--maxtext_ckpt_path` is cleared (random initialization is the one intended difference), `--maxtext_output_directory`
 becomes `<output-prefix>/<NAME>`, `--worker_id` becomes `NAME`, and the gRPC and discovery flags are dropped. The
 Tunix fail-fast wrapper is dropped as well, since it rewrites a clean exit into a restart.
 
-The knobs, all environment variables of `render_e2e.sh` (the Python script takes the same as flags):
+The knobs, all environment variables of `render_bench.sh` (the Python script takes the same as flags):
 
 `IMAGE` replaces the client image; the Pathways containers keep the e2e render's. `STEPS` (default 3) is the number of
 optimizer steps and `MICRO_STEPS` the `fwd_bwd` calls per step. Set `MICRO_STEPS` to what your e2e run does per
@@ -187,7 +189,7 @@ routing). Real tokens get replayed top-k routing with distinct experts, so MaxTe
 them. With `ROUTING=probe` (the default) each layer draws from a pool whose max expert-shard load ratio (under the
 EP=8 cp-as-ep layout) is calibrated by bisection to a per-layer target, with a seeded hot shard per layer.
 
-The default targets, `PROBE_REGULAR_MEAN_RATIO` in `e2e_driver.py`, were calibrated from a Qwen3.5-397B-A17B GRPO
+The default targets, `PROBE_REGULAR_MEAN_RATIO` in `bench_driver.py`, were calibrated from a Qwen3.5-397B-A17B GRPO
 production run (the per-layer mean over its 65 regular micro-batches). For another model or another phase of
 training, pass your own table with `ROUTING_TARGETS="r0,r1,..."` (one ratio per MoE layer, in the same probe order),
 or stretch the default with `ROUTING_SCALE` (targets become `1 + scale * (target - 1)`). `ROUTING=uniform` draws
@@ -212,19 +214,19 @@ the loss, and `ADVANTAGES=random` gives each segment a seeded N(0, 1) advantage
 and routing byte-identical. With these two, the gradient norm of two variants (e.g. a MoE padding flag on and off) is
 a meaningful comparison; the loss value itself is a function of the payload alone when the ratio is 1.
 
-`check_numerics_cpu.py --e2e-render e2e-dryrun.yaml` reproduces this on CPU with Tunix's GRPO loss and a
+`check_numerics_cpu.py --launch-render launch-dryrun.yaml` reproduces this on CPU with Tunix's GRPO loss and a
 one-parameter stand-in for the model. The stand-in replaces `compute_per_token_logps`, so it is tied to that
 function's signature in the Tunix version it runs against.
 
 ## Ahead-of-time memory check
 
 ```bash
-E2E_RENDER=e2e-dryrun.yaml AOT_IMAGE=<image> [TRAINER_SLICE=4x8x8] bash aot_e2e.sh
-python3 aot_summarize.py logs/aot-e2e-<time>.log
+LAUNCH_RENDER=launch-dryrun.yaml AOT_IMAGE=<image> [TRAINER_SLICE=4x8x8] bash aot_bench.sh
+python3 aot_summarize.py logs/aot-bench-<time>.log
 ```
 
-`aot_e2e.sh` takes the trainer's MaxText flags, `--xla_*` libtpu flags and mesh from the e2e render
-(`render_e2e.py --print-aot-env`, with the same variant knobs as `render_e2e.sh`) and runs `aot_trainer.py` in the
+`aot_bench.sh` takes the trainer's MaxText flags, `--xla_*` libtpu flags and mesh from the e2e render
+(`render_bench.py --print-aot-env`, with the same variant knobs as `render_bench.sh`) and runs `aot_trainer.py` in the
 image on a CPU host against a virtual topology. `aot_trainer.py` builds the MaxText argv as the Tunix trainer node does
 and calls `training_engine/maxtext_engine_compile.py` with the GRPO loss, packed micro-batches and router replay; its
 recipe constants (sequence lengths, optimizer, GRPO config) are the 397B recipe's and need editing for another
@@ -239,7 +241,7 @@ python3 -m pytest src/maxtext/experimental/rl_trainer_bench/tests
 ```
 
 The tests run on CPU with PyYAML and NumPy only: the renderer on a hand-written DRY_RUN fixture
-(`tests/testdata/e2e_dryrun_fixture.yaml`), and the payload builder and metric helpers on a small stand-in payload.
+(`tests/testdata/launch_dryrun_fixture.yaml`), and the payload builder and metric helpers on a small stand-in payload.
 They live next to the code, so the repository's default `testpaths = tests` does not collect them; run them by path.
 `check_payload_routing.py` and `check_numerics_cpu.py` are the heavier checks at the production shapes and need a
 matching Tunix.
diff --git a/src/maxtext/experimental/rl_trainer_bench/aot_e2e.sh b/src/maxtext/experimental/rl_trainer_bench/aot_bench.sh
similarity index 76%
rename from src/maxtext/experimental/rl_trainer_bench/aot_e2e.sh
rename to src/maxtext/experimental/rl_trainer_bench/aot_bench.sh
index fa3f596224..7e7fdb3e1e 100755
--- a/src/maxtext/experimental/rl_trainer_bench/aot_e2e.sh
+++ b/src/maxtext/experimental/rl_trainer_bench/aot_bench.sh
@@ -16,28 +16,28 @@
 # AOT memory check (CPU, no TPU) of the e2e trainer program (Tunix GRPO loss, packed micro-batch with router replay,
 # fwd_bwd / accumulate / update) with aot_trainer.py, inside an image whose MaxText and Tunix match the e2e run.
 #
-#   E2E_RENDER=<file> AOT_IMAGE=<image> [TRAINER_SLICE=4x8x8] [AOT_COMPILE_TOPOLOGY=tpu7x-512]
+#   LAUNCH_RENDER=<file> AOT_IMAGE=<image> [TRAINER_SLICE=4x8x8] [AOT_COMPILE_TOPOLOGY=tpu7x-512]
 #   [EXTRA_MAXTEXT_FLAGS=... REPLACE_MAXTEXT_FLAGS=... LIBTPU_EXTRA=... LIBTPU_DROP=... COMPUTE_LOGPS_CHUNK_SIZE=...]
-#   [LOG=<file>] bash aot_e2e.sh [extra MaxText key=value ...]
+#   [LOG=<file>] bash aot_bench.sh [extra MaxText key=value ...]
 #
-# The MaxText flags, --xla_* libtpu flags and mesh come from the e2e render (render_e2e.py --print-aot-env, with the
-# same variant knobs as render_e2e.sh), so the compile matches what the bench would run. AOT_COMPILE_TOPOLOGY must
+# The MaxText flags, --xla_* libtpu flags and mesh come from the e2e render (render_bench.py --print-aot-env, with the
+# same variant knobs as render_bench.sh), so the compile matches what the bench would run. AOT_COMPILE_TOPOLOGY must
 # name the slice: tpu7x-256 for 4x4x8 (the default), tpu7x-512 for 4x8x8. Summarize with aot_summarize.py <log>.
 set -euo pipefail
 HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
-: "${E2E_RENDER:?set E2E_RENDER to the DRY_RUN output of your e2e launch}"
+: "${LAUNCH_RENDER:?set LAUNCH_RENDER to the DRY_RUN output of your e2e launch}"
 : "${AOT_IMAGE:?set AOT_IMAGE to an image with the MaxText and Tunix of the e2e run}"
-_aot_env=$(python3 "$HERE/render_e2e.py" --print-aot-env --e2e-render "$E2E_RENDER" \
+_aot_env=$(python3 "$HERE/render_bench.py" --print-aot-env --launch-render "$LAUNCH_RENDER" \
   --extra-maxtext-flags "${EXTRA_MAXTEXT_FLAGS:-}" --replace-maxtext-flags "${REPLACE_MAXTEXT_FLAGS:-}" \
   --libtpu-extra="${LIBTPU_EXTRA:-}" --libtpu-drop="${LIBTPU_DROP:-}" --trainer-slice "${TRAINER_SLICE:-}" \
   --compute-logps-chunk-size "${COMPUTE_LOGPS_CHUNK_SIZE:-}")
 eval "$_aot_env"
 case "${TRAINER_SLICE:-}" in 4x8x8) _topo=tpu7x-512 ;; *) _topo=tpu7x-256 ;; esac
 TOPO=${AOT_COMPILE_TOPOLOGY:-$_topo}
-LOG=${LOG:-logs/aot-e2e-$(date -u +%Y%m%dT%H%M%SZ).log}
+LOG=${LOG:-logs/aot-bench-$(date -u +%Y%m%dT%H%M%SZ).log}
 mkdir -p "$(dirname "$LOG")"
 echo "log: $LOG"
-{ echo "=== $(date -u -Is) aot_e2e TOPO=$TOPO FSDP=$AOT_MESH_FSDP TP=$AOT_MESH_TP EXPERT=$AOT_MESH_EXPERT CONTEXT=$AOT_MESH_CONTEXT IMAGE=$AOT_IMAGE"
+{ echo "=== $(date -u -Is) aot_bench TOPO=$TOPO FSDP=$AOT_MESH_FSDP TP=$AOT_MESH_TP EXPERT=$AOT_MESH_EXPERT CONTEXT=$AOT_MESH_CONTEXT IMAGE=$AOT_IMAGE"
   echo "MAXTEXT_EXTRA_FLAGS=$MAXTEXT_EXTRA_FLAGS"; echo "LIBTPU_INIT_ARGS=$LIBTPU_INIT_ARGS"; } | tee -a "$LOG"
 docker run --rm --entrypoint python -v "$HERE":/bench:ro \
   -e JAX_PLATFORMS=cpu,tpu -e LIBTPU_INIT_ARGS="$LIBTPU_INIT_ARGS" -e MAXTEXT_EXTRA_FLAGS="$MAXTEXT_EXTRA_FLAGS" \
diff --git a/src/maxtext/experimental/rl_trainer_bench/aot_summarize.py b/src/maxtext/experimental/rl_trainer_bench/aot_summarize.py
index 92b5f0b5a4..d4b65072d8 100644
--- a/src/maxtext/experimental/rl_trainer_bench/aot_summarize.py
+++ b/src/maxtext/experimental/rl_trainer_bench/aot_summarize.py
@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Summarizes aot_e2e.sh logs: per kernel (fwd_bwd, accumulate, update) the CompiledMemoryStats and cost analysis,
+"""Summarizes aot_bench.sh logs: per kernel (fwd_bwd, accumulate, update) the CompiledMemoryStats and cost analysis,
 and the tool's TRAIN-KERNEL DEVICE PEAK against an HBM cap.
 
   python3 aot_summarize.py <log> [<log> ...] [--cap-mib 84992]
diff --git a/src/maxtext/experimental/rl_trainer_bench/aot_trainer.py b/src/maxtext/experimental/rl_trainer_bench/aot_trainer.py
index ef7d98a05f..dc4b096037 100644
--- a/src/maxtext/experimental/rl_trainer_bench/aot_trainer.py
+++ b/src/maxtext/experimental/rl_trainer_bench/aot_trainer.py
@@ -22,8 +22,8 @@
 
 The recipe constants below are those of the Qwen3.5-397B-A17B GRPO recipe the bench was built for (prompt and
 response lengths, optimizer, GRPO config). The mesh, the logps chunk and the topology come from the environment
-(`render_e2e.py --print-aot-env` derives the first two from an e2e DRY_RUN render); edit the rest for another
-recipe so that the compile matches your trainer. Run with `aot_e2e.sh`, inside an image whose MaxText and Tunix
+(`render_bench.py --print-aot-env` derives the first two from an e2e DRY_RUN render); edit the rest for another
+recipe so that the compile matches your trainer. Run with `aot_bench.sh`, inside an image whose MaxText and Tunix
 match the e2e run.
 
 Environment: MAXTEXT_EXTRA_FLAGS (required), LIBTPU_INIT_ARGS (--xla_* flags), AOT_MESH_FSDP (32), AOT_MESH_TP (1),
diff --git a/src/maxtext/experimental/rl_trainer_bench/e2e_driver.py b/src/maxtext/experimental/rl_trainer_bench/bench_driver.py
similarity index 99%
rename from src/maxtext/experimental/rl_trainer_bench/e2e_driver.py
rename to src/maxtext/experimental/rl_trainer_bench/bench_driver.py
index cb5d828954..0d05ba1d06 100644
--- a/src/maxtext/experimental/rl_trainer_bench/e2e_driver.py
+++ b/src/maxtext/experimental/rl_trainer_bench/bench_driver.py
@@ -57,7 +57,7 @@
 import sys
 import time
 
-logging.basicConfig(level=logging.INFO, format="%(asctime)s - [E2EBench] %(message)s", force=True)
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - [TrainerBench] %(message)s", force=True)
 
 import dataclasses
 
@@ -323,8 +323,8 @@ def metric_log_lines(step, summary):
 
 
 def main() -> int:
-  trainer_argv = shlex.split(os.environ["E2E_TRAINER_ARGV"])
-  orch_argv = shlex.split(os.environ["E2E_ORCH_ARGV"])
+  trainer_argv = shlex.split(os.environ["TRAINER_ARGV"])
+  orch_argv = shlex.split(os.environ["ORCH_ARGV"])
   # LOGP_GATES=off (numerics A/B): the payload's rollout log-probs are synthetic (-0.7 per completion token), so
   # against a random-init trainer the seq_logprob_error_threshold and seq-mask-tis gates drop every sequence, the
   # loss mask is empty and the gradient is exactly 0. Stripping the gate flags keeps every sequence in the loss.
diff --git a/src/maxtext/experimental/rl_trainer_bench/check_numerics_cpu.py b/src/maxtext/experimental/rl_trainer_bench/check_numerics_cpu.py
index d4b56d0001..6b98b59a9c 100644
--- a/src/maxtext/experimental/rl_trainer_bench/check_numerics_cpu.py
+++ b/src/maxtext/experimental/rl_trainer_bench/check_numerics_cpu.py
@@ -13,10 +13,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""CPU check of the numerics knobs of e2e_driver.py (no TPU, no cluster): why gradient_norm was 0, and the fix.
+"""CPU check of the numerics knobs of bench_driver.py (no TPU, no cluster): why gradient_norm was 0, and the fix.
 
 Builds the TRAINER_WARMUP_COMPILE payload with Tunix's own builder from the e2e orchestrator argv of a DRY_RUN
-render (--e2e-render; exactly as the driver does: `run_deepswe_dist._parse_args`, `_build_algo`,
+render (--launch-render; exactly as the driver does: `run_deepswe_dist._parse_args`, `_build_algo`,
 `BatchConfig`, routed_experts left out to keep it light), then the realistic payload, then runs Tunix's GRPO
 loss (`algo.loss_fn()`, the function `configure_worker` hands the trainer) on it with a one-parameter stand-in for
 the model: `compute_per_token_logps` is replaced by `theta + base_logp` per token, so everything after the forward
@@ -34,7 +34,7 @@
 
 Run with a Tunix and MaxText that match the e2e run, e.g. inside the trainer image:
   docker run --rm --entrypoint python -v $PWD:/tb:ro -v <dir of render>:/r:ro -e JAX_PLATFORMS=cpu <image> \
-    /tb/check_numerics_cpu.py --e2e-render /r/<e2e-dryrun>.yaml
+    /tb/check_numerics_cpu.py --launch-render /r/<launch-dryrun>.yaml
 """
 
 import argparse
@@ -48,11 +48,11 @@
 
 HERE = os.path.dirname(os.path.abspath(__file__))
 sys.path.insert(0, HERE)
-import e2e_driver as d  # pylint: disable=g-import-not-at-top
+import bench_driver as d  # pylint: disable=g-import-not-at-top
 
 
 def orch_argv_from_render(path, fsdp=None):
-  from render_e2e import _proc_main, load_documents, select_jobsets, split_command  # pylint: disable=g-import-not-at-top
+  from render_bench import _proc_main, load_documents, select_jobsets, split_command  # pylint: disable=g-import-not-at-top
   _, orch = select_jobsets(load_documents(path))
   _, flags = split_command(_proc_main(orch)[1]["command"][-1])
   if fsdp is not None:
@@ -141,14 +141,14 @@ def check_summary():
 
 def main():
   ap = argparse.ArgumentParser()
-  ap.add_argument("--e2e-render", required=True, help="DRY_RUN output of the e2e launcher (see README)")
+  ap.add_argument("--launch-render", required=True, help="DRY_RUN output of the e2e launcher (see README)")
   ap.add_argument("--fsdp", type=int, default=None, help="override --trainer_fsdp (64 for the 256-chip variant)")
   ap.add_argument("--rows", type=int, default=4, help="rows of the micro-batch fed to the loss (CPU time)")
   ap.add_argument("--far-logp", type=float, default=-12.0,
                   help="stand-in trainer logp per token; ~ -ln(vocab) for a near-uniform random-init model")
   a = ap.parse_args()
 
-  argv = orch_argv_from_render(a.e2e_render, a.fsdp)
+  argv = orch_argv_from_render(a.launch_render, a.fsdp)
   gates = [t for t in argv if t.split("=", 1)[0] in d.LOGP_GATE_FLAGS]
   print("orch gate flags:", gates)
   o, algo_on, warm = build(argv)
diff --git a/src/maxtext/experimental/rl_trainer_bench/check_payload_routing.py b/src/maxtext/experimental/rl_trainer_bench/check_payload_routing.py
index ed2f07c4d6..8d5ce39676 100644
--- a/src/maxtext/experimental/rl_trainer_bench/check_payload_routing.py
+++ b/src/maxtext/experimental/rl_trainer_bench/check_payload_routing.py
@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""CPU check of e2e_driver.realistic_payload at the 397B production shapes (no TPU, no cluster; needs Tunix).
+"""CPU check of bench_driver.realistic_payload at the 397B production shapes (no TPU, no cluster; needs Tunix).
 
 Builds the TRAINER_WARMUP_COMPILE payload with Tunix's own builder at the 397B e2e run's BatchConfig (64 rows x 65,536
 tokens, 16 segments per row at most, trainer_fsdp=32, trainer_expert=2, routed_experts [60 layers, top-10]), then
@@ -42,7 +42,7 @@
 import numpy as np
 
 sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-import e2e_driver as d  # pylint: disable=g-import-not-at-top
+import bench_driver as d  # pylint: disable=g-import-not-at-top
 
 from tunix.experimental.common import datatypes  # pylint: disable=g-import-not-at-top
 from tunix.experimental.examples.deepswe_dist import trainer_warmup  # pylint: disable=g-import-not-at-top
diff --git a/src/maxtext/experimental/rl_trainer_bench/render_e2e.py b/src/maxtext/experimental/rl_trainer_bench/render_bench.py
similarity index 93%
rename from src/maxtext/experimental/rl_trainer_bench/render_e2e.py
rename to src/maxtext/experimental/rl_trainer_bench/render_bench.py
index 1646c2cc55..ff42e870ba 100644
--- a/src/maxtext/experimental/rl_trainer_bench/render_e2e.py
+++ b/src/maxtext/experimental/rl_trainer_bench/render_bench.py
@@ -16,7 +16,7 @@
 """Renders the RL trainer benchmark JobSet from a DRY_RUN render of a Tunix GRPO e2e run.
 
 Input: the output of the e2e launcher with DRY_RUN=true (e.g. `DRY_RUN=true mlperf_gbs1024_submission.sh start`),
-saved to a file (`--e2e-render`). That output is a stream of YAML documents (one JobSet per orchestrator, trainer
+saved to a file (`--launch-render`). That output is a stream of YAML documents (one JobSet per orchestrator, trainer
 and rollout engine) that may be interleaved with launcher log lines; the trainer JobSet and the orchestrator JobSet
 are picked from it by the `--process_main=` of their `main` container. Alternatively pass the two JobSets as
 separate files (`--train-yaml`, `--orch-yaml`).
@@ -27,8 +27,8 @@
   * names: JobSet name, namespace, gang-id label; failurePolicy.maxRestarts=0 (a failed bench should not rerun
     the compile);
   * the client command: the same inline env prefix (MAXTEXT_EXTRA_FLAGS, LIBTPU_INIT_ARGS, ...) in front of
-    e2e_driver.py instead of the Tunix runtime main, with the e2e trainer's flags (E2E_TRAINER_ARGV) and the e2e
-    orchestrator's flags (E2E_ORCH_ARGV); the Tunix fail-fast wrapper and the gRPC port are dropped;
+    bench_driver.py instead of the Tunix runtime main, with the e2e trainer's flags (TRAINER_ARGV) and the e2e
+    orchestrator's flags (ORCH_ARGV); the Tunix fail-fast wrapper and the gRPC port are dropped;
   * trainer flags: --maxtext_ckpt_path is cleared (random init, the one intended difference),
     --maxtext_output_directory points at `--output-prefix/<name>` and --worker_id is the bench name;
   * optional knobs for A/B variants: --libtpu-extra / --libtpu-drop (written to every place the e2e render carries
@@ -42,7 +42,7 @@
     the orchestrator's --trainer_fsdp (which sets the assembler's pack size, fsdp x expert rows).
 
 Writes a file and nothing else; it never talks to a cluster. With --print-aot-env it prints shell exports of the
-trainer's MaxText and XLA flags and mesh for aot_e2e.sh instead of rendering.
+trainer's MaxText and XLA flags and mesh for aot_bench.sh instead of rendering.
 """
 
 from __future__ import annotations
@@ -302,24 +302,24 @@ def client_script(
 ) -> str:
   """The client container's bash script: writes the driver to /tmp and runs it under the e2e env prefix."""
   bench_env = (
-      f"E2E_TRAINER_ARGV={shlex.quote(shlex.join(trainer_argv))} "
-      f"E2E_ORCH_ARGV={shlex.quote(shlex.join(orch_argv))} STEPS={steps} MICRO_STEPS={micro_steps}"
+      f"TRAINER_ARGV={shlex.quote(shlex.join(trainer_argv))} "
+      f"ORCH_ARGV={shlex.quote(shlex.join(orch_argv))} STEPS={steps} MICRO_STEPS={micro_steps}"
   )
   for kv in shlex.split(driver_env):
     k, _, v = kv.partition("=")
     if not re.fullmatch(r"[A-Z_][A-Z0-9_]*", k):
       raise ValueError(f"bad --driver-env entry {kv!r}")
     bench_env += f" {k}={shlex.quote(v)}"
-  if "E2E_DRIVER_EOF" in driver_src:
+  if "BENCH_DRIVER_EOF" in driver_src:
     raise ValueError("driver source contains the heredoc terminator")
   return f"""# PID 1 wrapper: forwards SIGTERM to the benchmark so the grace period shuts it down cleanly.
 echo "Program Start: $(date)"
-cat > /tmp/e2e_driver.py <<'E2E_DRIVER_EOF'
-{driver_src}E2E_DRIVER_EOF
+cat > /tmp/bench_driver.py <<'BENCH_DRIVER_EOF'
+{driver_src}BENCH_DRIVER_EOF
 _sigterm() {{ kill -SIGTERM "$!" 2>/dev/null; }}
 trap _sigterm SIGTERM
 (
-{env_prefix} {bench_env} exec python /tmp/e2e_driver.py
+{env_prefix} {bench_env} exec python /tmp/bench_driver.py
 ) & PID=$!
 while kill -0 $PID 2>/dev/null; do
   sleep 5
@@ -333,7 +333,7 @@ def client_script(
 
 
 def aot_env(env_prefix: str, libtpu: list[str], trainer_flags: list[str]) -> str:
-  """Shell exports for aot_e2e.sh: the trainer's MaxText flags, its --xla_* libtpu flags and its mesh."""
+  """Shell exports for aot_bench.sh: the trainer's MaxText flags, its --xla_* libtpu flags and its mesh."""
   lines = [
       f"export MAXTEXT_EXTRA_FLAGS={shlex.quote(maxtext_extra_flags(env_prefix))}",
       f"export LIBTPU_INIT_ARGS={shlex.quote(' '.join(f for f in libtpu if f.startswith('--xla_')))}",
@@ -348,12 +348,12 @@ def aot_env(env_prefix: str, libtpu: list[str], trainer_flags: list[str]) -> str
 
 def build_parser() -> argparse.ArgumentParser:
   p = argparse.ArgumentParser(description=__doc__.split("\n", 1)[0])
-  p.add_argument("--e2e-render", default="", help="DRY_RUN output of the e2e launcher (multi-document YAML)")
-  p.add_argument("--train-yaml", default="", help="the e2e trainer JobSet alone (instead of --e2e-render)")
-  p.add_argument("--orch-yaml", default="", help="the e2e orchestrator JobSet alone (instead of --e2e-render)")
+  p.add_argument("--launch-render", default="", help="DRY_RUN output of the e2e launcher (multi-document YAML)")
+  p.add_argument("--train-yaml", default="", help="the e2e trainer JobSet alone (instead of --launch-render)")
+  p.add_argument("--orch-yaml", default="", help="the e2e orchestrator JobSet alone (instead of --launch-render)")
   p.add_argument("--trainer-main", default=TRAINER_MAIN, help="--process_main of the trainer JobSet")
   p.add_argument("--orch-main", default=ORCH_MAIN, help="--process_main of the orchestrator JobSet")
-  p.add_argument("--driver", default=os.path.join(os.path.dirname(os.path.abspath(__file__)), "e2e_driver.py"))
+  p.add_argument("--driver", default=os.path.join(os.path.dirname(os.path.abspath(__file__)), "bench_driver.py"))
   p.add_argument("--out", default="")
   p.add_argument("--name", default="")
   p.add_argument("--namespace", default="")
@@ -367,7 +367,7 @@ def build_parser() -> argparse.ArgumentParser:
   p.add_argument("--set-trainer-flags", default="", help='"k=v ...": replace existing trainer flags in place')
   p.add_argument("--set-orch-flags", default="", help='"k=v ...": replace existing orchestrator flags in place')
   p.add_argument("--image", default="", help="client (main container) image; empty keeps the e2e render's")
-  # Driver knobs (e2e_driver.py), as "K=V K=V": PAYLOAD ROUTING ROUTING_SCALE ROUTING_TARGETS PAD_FRACTION
+  # Driver knobs (bench_driver.py), as "K=V K=V": PAYLOAD ROUTING ROUTING_SCALE ROUTING_TARGETS PAD_FRACTION
   # FINAL_PAD_FRACTION PAYLOAD_SEED EP_SIZE ADVANTAGES LOGP_GATES.
   p.add_argument("--driver-env", default="")
   p.add_argument("--trainer-slice", default="", help="e.g. 4x8x8: FSDP fills it, expert/context/tp kept")
@@ -375,23 +375,23 @@ def build_parser() -> argparse.ArgumentParser:
   p.add_argument("--compute-logps-chunk-size", default="", help="trainer --compute_logps_chunk_size; empty keeps")
   p.add_argument("--profiler-skip", default="", help="trainer --skip_first_n_profiler_steps (counts fwd_bwd calls)")
   p.add_argument("--profiler-steps", default="", help="trainer --profiler_steps (fwd_bwd calls in the window)")
-  p.add_argument("--print-aot-env", action="store_true", help="print exports for aot_e2e.sh instead of rendering")
+  p.add_argument("--print-aot-env", action="store_true", help="print exports for aot_bench.sh instead of rendering")
   return p
 
 
 def main(argv=None):
   a = build_parser().parse_args(argv)
-  if a.e2e_render:
+  if a.launch_render:
     if a.train_yaml or a.orch_yaml:
-      sys.exit("pass either --e2e-render or --train-yaml/--orch-yaml")
-    train, orch = select_jobsets(load_documents(a.e2e_render), a.trainer_main, a.orch_main)
+      sys.exit("pass either --launch-render or --train-yaml/--orch-yaml")
+    train, orch = select_jobsets(load_documents(a.launch_render), a.trainer_main, a.orch_main)
   elif a.train_yaml and a.orch_yaml:
     with open(a.train_yaml) as f:
       train = yaml.safe_load(f)
     with open(a.orch_yaml) as f:
       orch = yaml.safe_load(f)
   else:
-    sys.exit("pass --e2e-render, or both --train-yaml and --orch-yaml")
+    sys.exit("pass --launch-render, or both --train-yaml and --orch-yaml")
   if not a.print_aot_env:
     missing = [k for k in ("out", "name", "namespace", "output_prefix") if not getattr(a, k)]
     if missing:
diff --git a/src/maxtext/experimental/rl_trainer_bench/render_e2e.sh b/src/maxtext/experimental/rl_trainer_bench/render_bench.sh
similarity index 87%
rename from src/maxtext/experimental/rl_trainer_bench/render_e2e.sh
rename to src/maxtext/experimental/rl_trainer_bench/render_bench.sh
index d3ae4e0409..99adfa10d0 100755
--- a/src/maxtext/experimental/rl_trainer_bench/render_e2e.sh
+++ b/src/maxtext/experimental/rl_trainer_bench/render_bench.sh
@@ -13,11 +13,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# Renders the RL trainer benchmark JobSet (e2e_driver.py on the trainer JobSet of an e2e DRY_RUN render).
+# Renders the RL trainer benchmark JobSet (bench_driver.py on the trainer JobSet of an e2e DRY_RUN render).
 # Writes one YAML file and nothing else; it never talks to a cluster.
 #
 # Required:
-#   E2E_RENDER=<file>        DRY_RUN output of the e2e launcher (see README)
+#   LAUNCH_RENDER=<file>     DRY_RUN output of the e2e launcher (see README)
 #   NAME=<name>              JobSet name, worker id and output folder (lowercase, at most 30 characters)
 #   K8S_NAMESPACE=<namespace>
 #   OUTPUT_PREFIX=<prefix>   the bench writes to <prefix>/<NAME> (e.g. a GCS folder you own)
@@ -35,11 +35,11 @@
 #   COMPUTE_LOGPS_CHUNK_SIZE=<n>       trainer --compute_logps_chunk_size in place; empty keeps the render's
 #   TRAINER_SLICE=<topology>           e.g. 4x8x8: re-target the trainer slice, FSDP fills it
 #   PROFILER_SKIP=<n> PROFILER_STEPS=<m>  trace fwd_bwd calls n .. n+m-1 of the run (see README)
-#   DRIVER_ENV="K=V ..."               e2e_driver.py knobs: PAYLOAD ROUTING ROUTING_SCALE ROUTING_TARGETS
+#   DRIVER_ENV="K=V ..."               bench_driver.py knobs: PAYLOAD ROUTING ROUTING_SCALE ROUTING_TARGETS
 #                                      PAD_FRACTION FINAL_PAD_FRACTION PAYLOAD_SEED EP_SIZE ADVANTAGES LOGP_GATES
 set -euo pipefail
 HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
-: "${E2E_RENDER:?set E2E_RENDER to the DRY_RUN output of your e2e launch}"
+: "${LAUNCH_RENDER:?set LAUNCH_RENDER to the DRY_RUN output of your e2e launch}"
 : "${NAME:?set NAME (lowercase, <=30 chars)}"
 : "${K8S_NAMESPACE:?set K8S_NAMESPACE}"
 : "${OUTPUT_PREFIX:?set OUTPUT_PREFIX (the bench writes to OUTPUT_PREFIX/NAME)}"
@@ -49,8 +49,8 @@ mkdir -p "$(dirname "$OUT")"
 # trajectories packed into 9 micro-batches of 64 rows on a 4x4x8 trainer, and the same token mass into 5 micro-batches
 # of 128 rows on 4x8x8. Set MICRO_STEPS to the fwd_bwd count per step of your own e2e run.
 case "${TRAINER_SLICE:-}" in 4x8x8) _ms=5 ;; *) _ms=9 ;; esac
-python3 "$HERE/render_e2e.py" \
-  --e2e-render "$E2E_RENDER" --driver "$HERE/e2e_driver.py" --out "$OUT" \
+python3 "$HERE/render_bench.py" \
+  --launch-render "$LAUNCH_RENDER" --driver "$HERE/bench_driver.py" --out "$OUT" \
   --name "$NAME" --namespace "$K8S_NAMESPACE" --output-prefix "$OUTPUT_PREFIX" \
   --steps "${STEPS:-3}" --micro-steps "${MICRO_STEPS:-$_ms}" \
   --extra-maxtext-flags "${EXTRA_MAXTEXT_FLAGS:-}" --replace-maxtext-flags "${REPLACE_MAXTEXT_FLAGS:-}" \
diff --git a/src/maxtext/experimental/rl_trainer_bench/structdiff_render.py b/src/maxtext/experimental/rl_trainer_bench/structdiff_render.py
index 9ce85d4590..615acedca6 100644
--- a/src/maxtext/experimental/rl_trainer_bench/structdiff_render.py
+++ b/src/maxtext/experimental/rl_trainer_bench/structdiff_render.py
@@ -13,9 +13,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Structured diff of two render_e2e.py JobSets: every YAML leaf except the client command, plus the client command
+"""Structured diff of two render_bench.py JobSets: every YAML leaf except the client command, plus the client command
 split into the env prefix (MAXTEXT_EXTRA_FLAGS tokens, LIBTPU_INIT_ARGS tokens, other env), the trainer argv
-(E2E_TRAINER_ARGV), the orchestrator argv (E2E_ORCH_ARGV) and the driver env. Prints only differences.
+(TRAINER_ARGV), the orchestrator argv (ORCH_ARGV) and the driver env. Prints only differences.
 
   python3 structdiff_render.py a.yaml b.yaml
 """
@@ -25,7 +25,7 @@
 import yaml
 
 sys.path.insert(0, __import__("os").path.dirname(__import__("os").path.abspath(__file__)))
-from render_e2e import _proc_main  # pylint: disable=g-import-not-at-top
+from render_bench import _proc_main  # pylint: disable=g-import-not-at-top
 
 
 def leaves(x, path=""):
@@ -54,11 +54,11 @@ def parse(path):
   _, main_c = _proc_main(js)
   script = main_c["command"][-1]
   main_c["command"] = ["<client script>"]
-  line = next(l for l in script.splitlines() if "exec python /tmp/e2e_driver.py" in l)
+  line = next(l for l in script.splitlines() if "exec python /tmp/bench_driver.py" in l)
   toks = shlex.split(line)
   env = {}
   for t in toks:
-    if t in ("exec", "python", "/tmp/e2e_driver.py"):
+    if t in ("exec", "python", "/tmp/bench_driver.py"):
       continue
     k, _, v = t.partition("=")
     env[k] = v
@@ -66,8 +66,8 @@ def parse(path):
       "leaf": dict(leaves(js)),
       "MAXTEXT_EXTRA_FLAGS": flagmap(env.pop("MAXTEXT_EXTRA_FLAGS", "").split()),
       "LIBTPU_INIT_ARGS": flagmap(env.pop("LIBTPU_INIT_ARGS", "").split()),
-      "trainer_argv": flagmap(shlex.split(env.pop("E2E_TRAINER_ARGV", ""))),
-      "orch_argv": flagmap(shlex.split(env.pop("E2E_ORCH_ARGV", ""))),
+      "trainer_argv": flagmap(shlex.split(env.pop("TRAINER_ARGV", ""))),
+      "orch_argv": flagmap(shlex.split(env.pop("ORCH_ARGV", ""))),
       "env": env,
   }
   return sec
diff --git a/src/maxtext/experimental/rl_trainer_bench/tests/payload_test.py b/src/maxtext/experimental/rl_trainer_bench/tests/payload_test.py
index c164581827..d140a9d41a 100644
--- a/src/maxtext/experimental/rl_trainer_bench/tests/payload_test.py
+++ b/src/maxtext/experimental/rl_trainer_bench/tests/payload_test.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""CPU-only tests of e2e_driver.py's payload builder and metric helpers (numpy only; no Tunix, no JAX).
+"""CPU-only tests of bench_driver.py's payload builder and metric helpers (numpy only; no Tunix, no JAX).
 
 The payload is a small stand-in for Tunix's packed RLTrainerPayload with the fields `realistic_payload` reads; it
 checks the same invariants as check_payload_routing.py does at the production shapes with Tunix's own builder.
@@ -29,7 +29,7 @@
 
 HERE = os.path.dirname(os.path.abspath(__file__))
 sys.path.insert(0, os.path.dirname(HERE))
-import e2e_driver as d  # pylint: disable=g-import-not-at-top,wrong-import-position
+import bench_driver as d  # pylint: disable=g-import-not-at-top,wrong-import-position
 
 pytestmark = pytest.mark.cpu_only
 
diff --git a/src/maxtext/experimental/rl_trainer_bench/tests/render_e2e_test.py b/src/maxtext/experimental/rl_trainer_bench/tests/render_bench_test.py
similarity index 89%
rename from src/maxtext/experimental/rl_trainer_bench/tests/render_e2e_test.py
rename to src/maxtext/experimental/rl_trainer_bench/tests/render_bench_test.py
index 24e2153ce1..f793346664 100644
--- a/src/maxtext/experimental/rl_trainer_bench/tests/render_e2e_test.py
+++ b/src/maxtext/experimental/rl_trainer_bench/tests/render_bench_test.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""CPU-only tests of render_e2e.py on a hand-written DRY_RUN fixture (no cluster, no Tunix, no JAX)."""
+"""CPU-only tests of render_bench.py on a hand-written DRY_RUN fixture (no cluster, no Tunix, no JAX)."""
 
 import io
 import os
@@ -29,19 +29,19 @@
 HERE = os.path.dirname(os.path.abspath(__file__))
 PKG = os.path.dirname(HERE)
 sys.path.insert(0, PKG)
-import render_e2e as R  # pylint: disable=g-import-not-at-top,wrong-import-position
+import render_bench as R  # pylint: disable=g-import-not-at-top,wrong-import-position
 import structdiff_render as SD  # pylint: disable=g-import-not-at-top,wrong-import-position
 
-FIXTURE = os.path.join(HERE, "testdata", "e2e_dryrun_fixture.yaml")
-DRIVER = os.path.join(PKG, "e2e_driver.py")
+FIXTURE = os.path.join(HERE, "testdata", "launch_dryrun_fixture.yaml")
+DRIVER = os.path.join(PKG, "bench_driver.py")
 
 pytestmark = pytest.mark.cpu_only
 
 
 def _client_env(js):
-  """The driver's inline env (E2E_TRAINER_ARGV, E2E_ORCH_ARGV, MAXTEXT_EXTRA_FLAGS, ...) of a rendered JobSet."""
+  """The driver's inline env (TRAINER_ARGV, ORCH_ARGV, MAXTEXT_EXTRA_FLAGS, ...) of a rendered JobSet."""
   _, main_c = R._proc_main(js)  # pylint: disable=protected-access
-  line = next(l for l in main_c["command"][-1].splitlines() if "exec python /tmp/e2e_driver.py" in l)
+  line = next(l for l in main_c["command"][-1].splitlines() if "exec python /tmp/bench_driver.py" in l)
   env = {}
   for t in shlex.split(line):
     k, eq, v = t.partition("=")
@@ -50,7 +50,7 @@ def _client_env(js):
   return env
 
 
-class RenderE2ETest(unittest.TestCase):
+class RenderBenchTest(unittest.TestCase):
 
   def setUp(self):
     self.tmp = tempfile.TemporaryDirectory()  # pylint: disable=consider-using-with
@@ -60,7 +60,7 @@ def tearDown(self):
 
   def render(self, *extra, name="tb-test"):
     out = os.path.join(self.tmp.name, f"{name}.yaml")
-    argv = ["--e2e-render", FIXTURE, "--driver", DRIVER, "--out", out, "--name", name, "--namespace", "bench-ns",
+    argv = ["--launch-render", FIXTURE, "--driver", DRIVER, "--out", out, "--name", name, "--namespace", "bench-ns",
             "--output-prefix", "/bench-out", *extra]
     with redirect_stdout(io.StringIO()):
       self.assertEqual(R.main(argv), 0)
@@ -87,19 +87,19 @@ def test_base_render(self):
     self.assertNotIn("ports", main_c)
     self.assertEqual(main_c["image"], "example.invalid/tunix:tag")
     script = main_c["command"][-1]
-    self.assertIn("cat > /tmp/e2e_driver.py <<'E2E_DRIVER_EOF'", script)
+    self.assertIn("cat > /tmp/bench_driver.py <<'BENCH_DRIVER_EOF'", script)
     self.assertNotIn("FT_REGISTERED_MARKER", script)  # Tunix fail-fast wrapper dropped
     with open(DRIVER) as f:
       self.assertIn(f.read(), script)
 
     env = _client_env(js)
-    targv = shlex.split(env["E2E_TRAINER_ARGV"])
+    targv = shlex.split(env["TRAINER_ARGV"])
     self.assertIn("--maxtext_ckpt_path=", targv)
     self.assertIn("--maxtext_output_directory=/bench-out/tb-test", targv)
     self.assertIn("--worker_id=tb-test", targv)
     self.assertFalse([t for t in targv if t.startswith(("--port=", "--discovery_addrs=", "--process_main="))])
     self.assertEqual([t for t in targv if t.startswith("--mesh_fsdp=")], ["--mesh_fsdp=32"])
-    oargv = shlex.split(env["E2E_ORCH_ARGV"])
+    oargv = shlex.split(env["ORCH_ARGV"])
     self.assertIn("--trainer_fsdp=32", oargv)
     self.assertIn("--seq_logprob_error_threshold=2.0", oargv)
     self.assertEqual((env["STEPS"], env["MICRO_STEPS"]), ("3", "4"))
@@ -134,12 +134,12 @@ def test_trainer_and_orch_flag_knobs(self):
                         "train_micro_batch_size=128", "--compute-logps-chunk-size", "1024",
                         "--profiler-skip", "16", "--profiler-steps", "3")
     env = _client_env(js)
-    targv = shlex.split(env["E2E_TRAINER_ARGV"])
+    targv = shlex.split(env["TRAINER_ARGV"])
     for flag in ("--train_micro_batch_size=128", "--compute_logps_chunk_size=1024", "--profiler_steps=3",
                  "--skip_first_n_profiler_steps=16"):
       self.assertEqual(sum(t.split("=", 1)[0] == flag.split("=", 1)[0] for t in targv), 1, flag)
       self.assertIn(flag, targv)
-    self.assertIn("--train_micro_batch_size=128", shlex.split(env["E2E_ORCH_ARGV"]))
+    self.assertIn("--train_micro_batch_size=128", shlex.split(env["ORCH_ARGV"]))
     with self.assertRaises(ValueError):
       self.render("--set-trainer-flags", "no_such_flag=1")
     with self.assertRaises(ValueError):
@@ -155,8 +155,8 @@ def test_trainer_slice(self):
     self.assertEqual((pw["completions"], pw["parallelism"]), (64, 64))
     self.assertEqual(pw["template"]["metadata"]["annotations"]["cloud.google.com/gke-tpu-slice-topology"], "4x8x8")
     env = _client_env(js)
-    self.assertIn("--mesh_fsdp=64", shlex.split(env["E2E_TRAINER_ARGV"]))
-    self.assertIn("--trainer_fsdp=64", shlex.split(env["E2E_ORCH_ARGV"]))
+    self.assertIn("--mesh_fsdp=64", shlex.split(env["TRAINER_ARGV"]))
+    self.assertIn("--trainer_fsdp=64", shlex.split(env["ORCH_ARGV"]))
     # Only the documented fields (plus the names) differ from the 4x4x8 render.
     buf = io.StringIO()
     with redirect_stdout(buf):
@@ -188,7 +188,7 @@ def test_driver_env_and_image(self):
   def test_print_aot_env(self):
     buf = io.StringIO()
     with redirect_stdout(buf):
-      R.main(["--print-aot-env", "--e2e-render", FIXTURE, "--trainer-slice", "4x8x8", "--libtpu-extra=--xla_new=1"])
+      R.main(["--print-aot-env", "--launch-render", FIXTURE, "--trainer-slice", "4x8x8", "--libtpu-extra=--xla_new=1"])
     out = buf.getvalue()
     self.assertIn("export MAXTEXT_EXTRA_FLAGS='context=offload num_moe_token_chunks=4 packing=True'", out)
     self.assertIn("export LIBTPU_INIT_ARGS='--xla_tpu_flag_a=true --xla_tpu_flag_b=4 --xla_new=1'", out)
diff --git a/src/maxtext/experimental/rl_trainer_bench/tests/testdata/e2e_dryrun_fixture.yaml b/src/maxtext/experimental/rl_trainer_bench/tests/testdata/launch_dryrun_fixture.yaml
similarity index 100%
rename from src/maxtext/experimental/rl_trainer_bench/tests/testdata/e2e_dryrun_fixture.yaml
rename to src/maxtext/experimental/rl_trainer_bench/tests/testdata/launch_dryrun_fixture.yaml

From a03bcd4f362b56717ff92f146151823773b8df07 Mon Sep 17 00:00:00 2001
From: Param Bole <parambole@google.com>
Date: Wed, 14 Oct 2026 08:03:56 -0700
Subject: [PATCH 4/4] rl_trainer_bench: STAGE_BEFORE_UPDATE switch for the step
 boundary

STAGE_BEFORE_UPDATE=on (default off) times the step boundary of a trainer
that stages the next step's first micro-batch before the previous step's
update, as Tunix's TrainerWorker.stage_fwd_bwd does (google/tunix#2877).
After step k's last fwd_bwd and before its update, step k+1's first
micro-batch goes through the engine's prepare_fwd_bwd_input, so its
host-to-device copy runs while the device still computes step k's last
micro-batch. The driver then waits for the device (wait_for_device), runs
the update, the drain and the metrics read as before, and step k+1's
first fwd_bwd takes the prepared input.

Off, the loop is the original one: bench_driver.py only gains lines (no
existing line changes), and each added line is either behind the switch
or leaves `staged` None. On, the engine must have prepare_fwd_bwd_input
and wait_for_device (AI-Hypercomputer/maxtext#5718). Without them the
driver raises right after the trainer worker initializes, before the
compile. A trainer worker that runs the update inside the step's last
fwd_bwd (Tunix's fused update, google/tunix#2719) leaves no update to
stage before; the driver sees the engine's train_step move and raises at
the first boundary. Any value other than on/off raises before Pathways
starts.

Same-pod A/B on the 397B recipe (128 chips, MICRO_STEPS=9, steps 1-4):
107.727 -> 100.007 s/step. The step's second fwd_bwd call loses its
7.5-7.8 s extra; regular micro-batches and the step tail are unchanged.

tests/stage_before_update_test.py runs main() against stand-ins for
Pathways, JAX, Tunix and the engine. Off reproduces the original loop's
engine calls, on trains the same micro-batches in the same order with one
update per step, and an engine without the two methods or a fused-update
worker is refused.
---
 .../experimental/rl_trainer_bench/README.md   |  34 +-
 .../rl_trainer_bench/bench_driver.py          |  60 ++
 .../rl_trainer_bench/render_bench.py          |   2 +-
 .../rl_trainer_bench/render_bench.sh          |   1 +
 .../tests/stage_before_update_test.py         | 652 ++++++++++++++++++
 5 files changed, 747 insertions(+), 2 deletions(-)
 create mode 100644 src/maxtext/experimental/rl_trainer_bench/tests/stage_before_update_test.py

diff --git a/src/maxtext/experimental/rl_trainer_bench/README.md b/src/maxtext/experimental/rl_trainer_bench/README.md
index f013a3f0d4..86cdf794c1 100644
--- a/src/maxtext/experimental/rl_trainer_bench/README.md
+++ b/src/maxtext/experimental/rl_trainer_bench/README.md
@@ -61,7 +61,7 @@ to the team.
 | `check_payload_routing.py` | CPU check of the realistic payload at the 397B production shapes (needs Tunix) |
 | `check_numerics_cpu.py` | CPU check of the numerics knobs with Tunix's GRPO loss (needs Tunix and MaxText) |
 | `hlo_collective_schedule.py` | static overlap summary of the collectives in a scheduled HLO dump |
-| `tests/` | CPU unit tests of the renderer and the payload builder |
+| `tests/` | CPU unit tests of the renderer, the payload builder and the driver's step loop |
 
 ## How to use it
 
@@ -218,6 +218,35 @@ a meaningful comparison; the loss value itself is a function of the payload alon
 one-parameter stand-in for the model. The stand-in replaces `compute_per_token_logps`, so it is tied to that
 function's signature in the Tunix version it runs against.
 
+## Step boundary: STAGE_BEFORE_UPDATE
+
+`DRIVER_ENV="STAGE_BEFORE_UPDATE=on"` (default `off`; any other value raises) times the step boundary of a trainer that
+stages the next step's first micro-batch before the previous step's update, as Tunix's `TrainerWorker.stage_fwd_bwd`
+does before `apply_deferred_update` (google/tunix#2877). After step k's last `fwd_bwd` and before its `update`, the
+driver hands step k+1's first micro-batch to the engine's `prepare_fwd_bwd_input`, which maps it to loss inputs and
+places it on `fwd_bwd`'s input shardings while the device still computes step k's last micro-batch. The driver then
+calls `wait_for_device` (it returns how long that micro-batch still ran) and runs the update, the drain and the
+metrics read as with the switch off; step k+1's first `fwd_bwd` takes the prepared input. The trainer sees the same
+micro-batches, in the same order, with the same `apply_optimizer` flags and one `update` per step. Each staged
+boundary logs, after step k's `completed step` line:
+
+```
+STAGE_BEFORE_UPDATE step=<k> staged=bench_<k+1>_0 stage_s=<host seconds> device_wait_s=<seconds>
+```
+
+A `device_wait_s` well above `stage_s` means the copy was issued while the device was still busy.
+`prepare_fwd_bwd_input` and `wait_for_device` come with AI-Hypercomputer/maxtext#5718. With an engine that lacks them
+the driver raises a `RuntimeError` naming them right after the trainer worker initializes, before the compile; the
+switch off needs nothing new. A trainer worker that runs the update inside the step's last `fwd_bwd` (Tunix's fused
+update, google/tunix#2719) leaves no update to stage before, so with it the driver raises at the first boundary.
+
+Why it shows on this bench: with replayed routing a micro-batch's input is several GB (`routed_experts` is
+`[rows, seq, layers, top_k]`), and with the switch off step k+1's first micro-batch is copied to the device after the
+drain, while the device is idle. In one same-pod A/B on the 397B recipe (128 TPU v7x chips, `MICRO_STEPS=9`, steps
+1-4), the step's second `fwd_bwd` call took 7.5-7.8 s longer than a regular one with the switch off and about the same
+as a regular one with it on, and the step went from 107.7 s to 100.0 s; the regular micro-batches and the step's tail
+(last micro-batch, update and drain) did not change.
+
 ## Ahead-of-time memory check
 
 ```bash
@@ -242,6 +271,9 @@ python3 -m pytest src/maxtext/experimental/rl_trainer_bench/tests
 
 The tests run on CPU with PyYAML and NumPy only: the renderer on a hand-written DRY_RUN fixture
 (`tests/testdata/launch_dryrun_fixture.yaml`), and the payload builder and metric helpers on a small stand-in payload.
+`tests/stage_before_update_test.py` runs the driver's `main()` against stand-ins for Pathways, JAX, Tunix and the
+training engine: with `STAGE_BEFORE_UPDATE` off the engine sees the original loop's calls, with it on the same
+micro-batches in the same order with one update per step, and an engine without the two methods is refused.
 They live next to the code, so the repository's default `testpaths = tests` does not collect them; run them by path.
 `check_payload_routing.py` and `check_numerics_cpu.py` are the heavier checks at the production shapes and need a
 matching Tunix.
diff --git a/src/maxtext/experimental/rl_trainer_bench/bench_driver.py b/src/maxtext/experimental/rl_trainer_bench/bench_driver.py
index 0d05ba1d06..68a96d98ad 100644
--- a/src/maxtext/experimental/rl_trainer_bench/bench_driver.py
+++ b/src/maxtext/experimental/rl_trainer_bench/bench_driver.py
@@ -49,6 +49,20 @@
 
 The intended difference from e2e: no checkpoint is loaded (random init; the trainer argv's --maxtext_ckpt_path is
 cleared by the launcher), so weights and routing values differ while shapes, shardings and the program match.
+
+STAGE_BEFORE_UPDATE=on (default off; anything else raises): the step boundary of a trainer that stages the next step's
+first micro-batch before the previous step's update (Tunix's `TrainerWorker.stage_fwd_bwd` before
+`apply_deferred_update`, google/tunix#2877). After step k's last `fwd_bwd`, and before its `update`, step k+1's first
+micro-batch is mapped to loss inputs and placed on the mesh (`prepare_fwd_bwd_input`, i.e. its host-to-device copy runs
+while the device still computes step k's last micro-batch), the driver waits for that micro-batch (`wait_for_device`,
+which reports how long it still ran), then runs the update, the drain and the metrics read exactly as with the switch
+off, and step k+1's first `fwd_bwd` takes the prepared input. The trainer sees the same micro-batches, in the same
+order, with the same `apply_optimizer` flags and one `update` per step. It needs a training engine with
+`prepare_fwd_bwd_input` and `wait_for_device` (AI-Hypercomputer/maxtext#5718); without them the driver raises right
+after the trainer worker initializes, before the compile. A trainer worker that runs the update inside the step's last
+`fwd_bwd` (Tunix's fused update, google/tunix#2719) leaves no update to stage before, and the driver raises at the
+first boundary. Off, the loop runs as before: the added lines are behind the switch or leave `staged` None. Each
+staged boundary logs `STAGE_BEFORE_UPDATE step=k ...` after the step's `completed step` line.
 """
 
 import logging
@@ -335,6 +349,10 @@ def main() -> int:
     raise ValueError(f"LOGP_GATES must be on or off, got {logp_gates!r}")
   steps = int(os.environ.get("STEPS", "3"))
   micro_steps = int(os.environ.get("MICRO_STEPS", "4"))
+  stage_before_update = os.environ.get("STAGE_BEFORE_UPDATE", "off")
+  if stage_before_update not in ("on", "off"):
+    raise ValueError(f"STAGE_BEFORE_UPDATE must be on or off, got {stage_before_update!r}")
+  stage_before_update = stage_before_update == "on"
 
   # Same JAX init as K8sJaxContext.initialize() under JAX_PLATFORMS=proxy.
   import pathwaysutils  # pylint: disable=g-import-not-at-top
@@ -366,6 +384,20 @@ def main() -> int:
   t0 = time.monotonic()
   worker.initialize()
   logging.info("Trainer worker initialized in %.1fs on %d devices.", time.monotonic() - t0, jax.device_count())
+  if stage_before_update:
+    engine = worker._trainer  # pylint: disable=protected-access
+    missing = [n for n in ("prepare_fwd_bwd_input", "wait_for_device") if not callable(getattr(engine, n, None))]
+    if missing:
+      raise RuntimeError(
+          f"STAGE_BEFORE_UPDATE=on needs the training engine's prepare_fwd_bwd_input and "
+          f"wait_for_device (AI-Hypercomputer/maxtext#5718); {type(engine).__name__} lacks {missing}. "
+          "Use a MaxText that has them, or STAGE_BEFORE_UPDATE=off."
+      )
+    logging.info(
+        "STAGE_BEFORE_UPDATE on: step k+1's first micro-batch is prepared on the mesh before step k's update "
+        "(%s.prepare_fwd_bwd_input, then wait_for_device).",
+        type(engine).__name__,
+    )
 
   # Tokenizer ids, as run_deepswe_dist.main derives them.
   tokenizer_path = o_args.tokenizer_path or os.getenv("MODEL_DIR") or o_args.model_id
@@ -465,17 +497,37 @@ def main() -> int:
   elif mode != "warmup":
     raise ValueError(f"PAYLOAD must be realistic or warmup, got {mode!r}")
 
+  staged = None  # STAGE_BEFORE_UPDATE: step k+1's first micro-batch as prepare_fwd_bwd_input returned it
   for step in range(steps):
     t0 = time.monotonic()
     times = []
+    if stage_before_update:
+      train_step_at_start = engine.train_step  # Moves only in update(): see the check before the stage below.
     for m in range(micro_steps):
       tm = time.monotonic()
       last = m == micro_steps - 1
       req = datatypes.TrainRequest(request_id=f"bench_{step}_{m}", payload=final if last else regular, metadata={})
+      if m == 0 and staged is not None:
+        # The same micro-batch (step k+1's first, `regular`, or `final` when MICRO_STEPS is 1), already prepared.
+        req = datatypes.TrainRequest(request_id=req.request_id, payload=staged, metadata={})
+        staged = None
       # As DistributedRLEngine.train_step: the step's last micro-batch is the one with apply_optimizer=True.
       # Per-micro-step runtime load ratios: MOE_DROP_AUDIT lines with the MaxText flag log_moe_drops=true.
       worker.fwd_bwd(request=req, apply_optimizer=last)
       times.append(time.monotonic() - tm)
+    if stage_before_update and step + 1 < steps:
+      if engine.train_step != train_step_at_start:  # pylint: disable=possibly-used-before-assignment
+        raise RuntimeError(
+            "STAGE_BEFORE_UPDATE=on, but the trainer worker already ran this step's update inside its last fwd_bwd "
+            "(Tunix's fused update, google/tunix#2719), so nothing can be staged before it. Use "
+            "STAGE_BEFORE_UPDATE=off with this Tunix."
+        )
+      # Step k+1's first micro-batch (`regular`, or `final` when MICRO_STEPS is 1), mapped and placed on the mesh
+      # while the device still computes this step's last micro-batch; then how long that micro-batch still ran.
+      ts = time.monotonic()
+      staged = engine.prepare_fwd_bwd_input(final if micro_steps == 1 else regular)
+      stage_s = time.monotonic() - ts
+      device_wait_s = engine.wait_for_device()
     worker.update()
     # Drain the step's device work before reading the clock, as maxtext_engine/train.py does.
     throttler = getattr(worker._trainer, "_throttler", None)  # pylint: disable=protected-access
@@ -488,6 +540,14 @@ def main() -> int:
         "completed step: %d, seconds: %.3f, micro_steps: %d, seconds/micro_step: %.3f, fwd_bwd dispatch s: %s",
         step, dt, micro_steps, dt / micro_steps, [round(x, 3) for x in times],
     )
+    if stage_before_update and step + 1 < steps:
+      logging.info(
+          "STAGE_BEFORE_UPDATE step=%d staged=bench_%d_0 stage_s=%.3f device_wait_s=%.3f",
+          step,
+          step + 1,
+          stage_s,
+          device_wait_s,
+      )
     # Numerics comparison (e.g. moe_exclude_padding_from_routing on vs off): the step's loss (per micro-batch and
     # pooled), gradient norm and every other trainer metric, read after the step's timing so they do not affect it.
     try:
diff --git a/src/maxtext/experimental/rl_trainer_bench/render_bench.py b/src/maxtext/experimental/rl_trainer_bench/render_bench.py
index ff42e870ba..7faf5868ad 100644
--- a/src/maxtext/experimental/rl_trainer_bench/render_bench.py
+++ b/src/maxtext/experimental/rl_trainer_bench/render_bench.py
@@ -368,7 +368,7 @@ def build_parser() -> argparse.ArgumentParser:
   p.add_argument("--set-orch-flags", default="", help='"k=v ...": replace existing orchestrator flags in place')
   p.add_argument("--image", default="", help="client (main container) image; empty keeps the e2e render's")
   # Driver knobs (bench_driver.py), as "K=V K=V": PAYLOAD ROUTING ROUTING_SCALE ROUTING_TARGETS PAD_FRACTION
-  # FINAL_PAD_FRACTION PAYLOAD_SEED EP_SIZE ADVANTAGES LOGP_GATES.
+  # FINAL_PAD_FRACTION PAYLOAD_SEED EP_SIZE ADVANTAGES LOGP_GATES STAGE_BEFORE_UPDATE.
   p.add_argument("--driver-env", default="")
   p.add_argument("--trainer-slice", default="", help="e.g. 4x8x8: FSDP fills it, expert/context/tp kept")
   p.add_argument("--devices-per-chip", type=int, default=2, help="JAX devices per chip (tpu7x: 2)")
diff --git a/src/maxtext/experimental/rl_trainer_bench/render_bench.sh b/src/maxtext/experimental/rl_trainer_bench/render_bench.sh
index 99adfa10d0..40993ac922 100755
--- a/src/maxtext/experimental/rl_trainer_bench/render_bench.sh
+++ b/src/maxtext/experimental/rl_trainer_bench/render_bench.sh
@@ -37,6 +37,7 @@
 #   PROFILER_SKIP=<n> PROFILER_STEPS=<m>  trace fwd_bwd calls n .. n+m-1 of the run (see README)
 #   DRIVER_ENV="K=V ..."               bench_driver.py knobs: PAYLOAD ROUTING ROUTING_SCALE ROUTING_TARGETS
 #                                      PAD_FRACTION FINAL_PAD_FRACTION PAYLOAD_SEED EP_SIZE ADVANTAGES LOGP_GATES
+#                                      STAGE_BEFORE_UPDATE
 set -euo pipefail
 HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
 : "${LAUNCH_RENDER:?set LAUNCH_RENDER to the DRY_RUN output of your e2e launch}"
diff --git a/src/maxtext/experimental/rl_trainer_bench/tests/stage_before_update_test.py b/src/maxtext/experimental/rl_trainer_bench/tests/stage_before_update_test.py
new file mode 100644
index 0000000000..6c72ddc0b3
--- /dev/null
+++ b/src/maxtext/experimental/rl_trainer_bench/tests/stage_before_update_test.py
@@ -0,0 +1,652 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""CPU-only tests of bench_driver.py's STAGE_BEFORE_UPDATE switch (NumPy only; no Tunix, no JAX, no Pathways).
+
+`main()` runs end to end against stand-ins for what it imports: Pathways, JAX, the Tunix example modules (argv
+parsing, the algorithm, the warm-up payload), the batch assembler, the tokenizer, Tunix's `TrainerWorker` (which, like
+the real one, hands `request.payload` to the engine's `fwd_bwd` and calls the engine's `update`) and the MaxText
+training engine, which records every call with a digest of the micro-batch it trains. The micro-batches are the
+driver's own `realistic_payload` of a small warm-up-style payload, so the regular and the final one differ.
+
+  * Off (unset or `off`): the engine sees exactly the calls of the original loop, per step MICRO_STEPS `fwd_bwd` (the
+    regular micro-batch, then the final one, which has fully padded rows, with apply_optimizer=True), one `update`,
+    the drain and the metrics read. Neither hook is called, nothing STAGE_BEFORE_UPDATE is logged, and an engine
+    without the hooks works.
+  * On: the same `fwd_bwd` sequence (request ids, micro-batch digests, apply_optimizer flags) and the same STEP_LOSS /
+    STEP_METRICS lines as off. Step k+1's first micro-batch arrives prepared, made by one `prepare_fwd_bwd_input`
+    between step k's last `fwd_bwd` and its `update`, with one `wait_for_device` in between; one `update` per step;
+    one STAGE_BEFORE_UPDATE line per staged boundary, carrying the wait the engine returned.
+  * Refusal: an engine without either hook stops the driver right after the worker initializes, before the compile;
+    a worker that runs the update inside the last `fwd_bwd` (Tunix's fused update) stops it at the first boundary,
+    before anything is staged; and a value other than on/off stops it before Pathways is initialized.
+
+The checks are functions that return a list of failures, and `CheckerTest` shows each one failing on a wrong run.
+`stand_in_modules` mirrors `main()`'s imports, so a new import in `main()` needs a stand-in there.
+"""
+
+import dataclasses
+import hashlib
+import logging
+import os
+import re
+import sys
+import types
+import unittest
+from unittest import mock
+
+import numpy as np
+import pytest
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, os.path.dirname(HERE))
+import bench_driver  # pylint: disable=g-import-not-at-top,wrong-import-position
+
+pytestmark = pytest.mark.cpu_only
+
+ROWS, BUDGET, LAYERS, TOP_K, EXPERTS = 4, 256, 4, 2, 16
+# Every environment variable the driver reads, cleared for each run so that the caller's environment cannot leak in.
+DRIVER_ENV = (
+    "TRAINER_ARGV",
+    "ORCH_ARGV",
+    "STEPS",
+    "MICRO_STEPS",
+    "STAGE_BEFORE_UPDATE",
+    "LOGP_GATES",
+    "PAYLOAD",
+    "ROUTING",
+    "ROUTING_SCALE",
+    "ROUTING_TARGETS",
+    "PAD_FRACTION",
+    "FINAL_PAD_FRACTION",
+    "PAYLOAD_SEED",
+    "EP_SIZE",
+    "ADVANTAGES",
+    "MODEL_DIR",
+)
+
+
+@dataclasses.dataclass(frozen=True)
+class StandInPayload:
+  """The fields of Tunix's packed RLTrainerPayload that `realistic_payload` reads and writes."""
+
+  completion_ids: np.ndarray
+  segment_ids: np.ndarray
+  segment_positions: np.ndarray
+  completion_mask: np.ndarray
+  advantages: np.ndarray
+  old_per_token_logps: np.ndarray | None
+  rollout_per_token_logps: np.ndarray | None
+  routed_experts: np.ndarray | None
+
+
+def warmup_payload():
+  """A warm-up-style micro-batch: two short segments at the start of each row, the rest padding, routing UNSET."""
+  seg = np.zeros((ROWS, BUDGET), np.int32)
+  cm = np.zeros((ROWS, BUDGET), np.int32)
+  adv = np.zeros((ROWS, BUDGET), np.float32)
+  for s in range(2):
+    seg[:, s * 40 : (s + 1) * 40] = s + 1
+    cm[:, s * 40 + 16 : (s + 1) * 40] = 1
+    adv[:, s * 40 + 16 : (s + 1) * 40] = 1.0 if s else -1.0
+  return StandInPayload(
+      completion_ids=np.where(seg > 0, 7, 0).astype(np.int32),
+      segment_ids=seg,
+      segment_positions=np.tile(np.arange(BUDGET, dtype=np.int32), (ROWS, 1)),
+      completion_mask=cm,
+      advantages=adv,
+      old_per_token_logps=None,
+      rollout_per_token_logps=np.zeros((ROWS, BUDGET), np.float32),
+      routed_experts=np.full((ROWS, BUDGET, LAYERS, TOP_K), bench_driver.UNSET_ROUTED_EXPERT, np.int16),
+  )
+
+
+def digest(payload):
+  """Content digest of a micro-batch: every field's name, shape, dtype and bytes."""
+  h = hashlib.md5()
+  for f in dataclasses.fields(payload):
+    v = getattr(payload, f.name)
+    h.update(f.name.encode())
+    if v is not None:
+      a = np.ascontiguousarray(np.asarray(v))
+      h.update(f"{a.shape}{a.dtype.str}".encode())
+      h.update(a.tobytes())
+  return h.hexdigest()[:12]
+
+
+class Prepared:
+  """What the stand-in engine's `prepare_fwd_bwd_input` returns (MaxText's `PreparedFwdBwdInput`)."""
+
+  def __init__(self, payload_digest):
+    self.digest = payload_digest
+    self.used = False
+
+
+@dataclasses.dataclass
+class WeightedMetric:
+  unreduced_sum: np.ndarray
+  denominator: np.ndarray
+  eps: float | None = None
+  min_denom: float | None = None
+
+
+def fully_padded_rows(payload):
+  """Rows with no real token: the final micro-batch has FINAL_PAD_FRACTION of them, a regular one has none."""
+  return int((np.asarray(payload.segment_ids) == 0).all(axis=1).sum())
+
+
+class EngineWithoutHooks:
+  """Stand-in for MaxText's training engine without the staging methods (as on #5701's base branch).
+
+  Records each call in `calls`: ("fwd_bwd", step, micro index, micro-batch digest, prepared, apply_optimizer),
+  ("update", step, micro-batches accumulated), ("drain", step), ("get_metrics", step), and each micro-batch's fully
+  padded row count in `pad_rows` (by digest). Its metrics are a function of the micro-batches the step trained, so a
+  different micro-batch changes the driver's STEP_LOSS line.
+  """
+
+  def __init__(self, calls):
+    self.calls = calls
+    self._config = types.SimpleNamespace(num_experts=EXPERTS, inhomogeneous_layer_cycle_interval=4)
+    self._throttler = types.SimpleNamespace(wait_for_all=lambda: calls.append(("drain", self.step)))
+    self.step = 0
+    self.micro = []
+    self.metrics = None
+    self.pad_rows = {}
+
+  @property
+  def train_step(self):
+    """As MaxText's engine: the number of updates applied."""
+    return self.step
+
+  def fwd_bwd(self, payload, **kwargs):
+    """Trains one micro-batch: a payload, or what `prepare_fwd_bwd_input` made of one."""
+    if isinstance(payload, Prepared):
+      if payload.used:
+        raise AssertionError("a prepared input was trained twice")
+      payload.used = True
+      d, prepared = payload.digest, True
+    else:
+      d, prepared = digest(payload), False
+      self.pad_rows[d] = fully_padded_rows(payload)
+    self.calls.append(("fwd_bwd", self.step, len(self.micro), d, prepared, kwargs.get("apply_optimizer")))
+    self.micro.append(d)
+
+  def update(self, **kwargs):
+    """Applies the step: its metrics are a function of the micro-batches it trained."""
+    del kwargs
+    self.calls.append(("update", self.step, len(self.micro)))
+    if not self.micro:  # As MaxText's engine: nothing accumulated, nothing applied.
+      return self.step
+    sums = np.array([int(d[:6], 16) % 1000 for d in self.micro], np.float64)
+    dens = np.array([1 + int(d[6:8], 16) % 3 for d in self.micro], np.float64)
+    norm = float(int(hashlib.md5("".join(self.micro).encode()).hexdigest()[:6], 16))
+    self.metrics = types.SimpleNamespace(
+        weighted_metrics={"loss": WeightedMetric(sums, dens)}, scalar_metrics={"gradient_norm": np.float64(norm)}
+    )
+    self.micro = []
+    self.step += 1
+    return self.step
+
+  def get_metrics(self):
+    self.calls.append(("get_metrics", self.step))
+    return self.metrics
+
+
+class Engine(EngineWithoutHooks):
+  """The engine with the two methods of AI-Hypercomputer/maxtext#5718."""
+
+  def prepare_fwd_bwd_input(self, payload):
+    d = digest(payload)
+    self.pad_rows[d] = fully_padded_rows(payload)
+    self.calls.append(("prepare", self.step, len(self.micro), d))
+    return Prepared(d)
+
+  def wait_for_device(self):
+    self.calls.append(("wait_for_device", self.step, len(self.micro)))
+    return 0.25
+
+
+class EngineWithPrepareOnly(EngineWithoutHooks):
+  prepare_fwd_bwd_input = Engine.prepare_fwd_bwd_input
+
+
+class TrainerWorker:
+  """Stand-in for Tunix's TrainerWorker without a fused update: `fwd_bwd` hands `request.payload` to the engine as
+  is, `update` and `get_metrics` call the engine's."""
+
+  def __init__(self, run, trainer_factory, **kwargs):
+    del kwargs  # worker_id, logps_chunk_size, logps_micro_batch_size, execution_context
+    self._trainer_factory = trainer_factory
+    self._trainer = None
+    self.request_ids = []
+    run.worker = self
+
+  def initialize(self):
+    self._trainer = self._trainer_factory()
+    self._trainer.calls.append(("initialize",))
+
+  def with_loss_fn(self, loss_fn, has_aux=False):
+    del loss_fn, has_aux
+
+  def with_gen_model_input_fn(self, gen_model_input_fn):
+    del gen_model_input_fn
+
+  def compile(self, dummy_data=None):
+    del dummy_data
+    self._trainer.calls.append(("compile",))
+
+  def fwd_bwd(self, request, apply_optimizer=False, **kwargs):
+    self.request_ids.append(request.request_id)
+    self._trainer.fwd_bwd(request.payload, apply_optimizer=apply_optimizer, **kwargs)
+
+  def update(self, **kwargs):
+    return self._trainer.update(**kwargs)
+
+  def get_metrics(self):
+    return self._trainer.get_metrics()
+
+  def stop(self):
+    self._trainer.calls.append(("stop",))
+
+
+class FusedUpdateTrainerWorker(TrainerWorker):
+  """A TrainerWorker with Tunix's fused update (google/tunix#2719): `fwd_bwd(apply_optimizer=True)` also runs
+  `update()`."""
+
+  def fwd_bwd(self, request, apply_optimizer=False, **kwargs):
+    super().fwd_bwd(request, apply_optimizer=apply_optimizer, **kwargs)
+    if apply_optimizer:
+      self.update()
+
+
+@dataclasses.dataclass
+class TrainRequest:
+  request_id: str
+  payload: object
+  metadata: dict
+
+
+def _package(name, **attrs):
+  m = types.ModuleType(name)
+  m.__path__ = []  # A package, so that `from <name> import <submodule>` resolves through sys.modules.
+  m.__dict__.update(attrs)
+  return m
+
+
+def stand_in_modules(run, engine_cls, worker_cls):
+  """sys.modules entries for everything `main()` imports, recording into `run`."""
+  o_args = types.SimpleNamespace(
+      tokenizer_path="stand-in",
+      model_id="stand-in",
+      max_prompt_length=64,
+      max_response_length=BUDGET - 64,
+      max_seq_token_per_tpu=BUDGET,
+      max_segments_per_packed_row=2,
+      trainer_fsdp=4,
+      trainer_dp=1,
+      trainer_expert=1,
+      num_generations=2,
+      mini_batch_size=2,
+      trainer_warmup_routed_experts=True,
+      trainer_warmup_routed_layers=LAYERS,
+      trainer_warmup_routed_top_k=TOP_K,
+  )
+  algo = types.SimpleNamespace(
+      train_micro_batch_size=ROWS, loss_fn=lambda: None, build_gen_model_input_fn=lambda pad_id, eos_id: None
+  )
+  t_args = types.SimpleNamespace(worker_id="stand-in", compute_logps_chunk_size=0, compute_logps_micro_batch_size=None)
+  tokenizer = types.SimpleNamespace(pad_token_id=0, eos_token="<eos>", vocab_size=1000)
+  leaves = {
+      "tunix.experimental.examples.common.run_trainer_node": {
+          "_parse_args": lambda argv: t_args,
+          "_create_maxtext_trainer_factory": lambda args: (lambda: engine_cls(run.calls), None),
+      },
+      "tunix.experimental.examples.deepswe_dist.run_deepswe_dist": {
+          "_parse_args": lambda argv: o_args,
+          "_build_algo": lambda args: algo,
+      },
+      "tunix.experimental.examples.deepswe_dist.trainer_warmup": {
+          "routed_experts_shape_from_config": lambda path: (LAYERS, TOP_K),
+          "build_warmup_payload": lambda **kwargs: warmup_payload(),
+      },
+      "tunix.experimental.orchestrator.batch_assembly": {
+          "BatchConfig": types.SimpleNamespace,
+          "create_batch_assembler": lambda **kwargs: types.SimpleNamespace(pad_id=0, eos_id=0),
+          "describe_payload": lambda payload: "stand-in micro-batch",
+      },
+      "tunix.experimental.common.datatypes": {"TrainRequest": TrainRequest},
+      "tunix.experimental.worker.trainer_worker": {"TrainerWorker": lambda **kwargs: worker_cls(run, **kwargs)},
+  }
+  mods = {
+      "pathwaysutils": _package("pathwaysutils", initialize=lambda: run.calls.append(("pathways_initialize",))),
+      "jax": _package("jax", device_count=lambda: 4, device_get=lambda x: x, block_until_ready=lambda x: x),
+      "transformers": _package(
+          "transformers", AutoTokenizer=types.SimpleNamespace(from_pretrained=lambda *a, **k: tokenizer)
+      ),
+  }
+  for name, attrs in leaves.items():
+    mods[name] = _package(name, **attrs)
+    parts = name.split(".")
+    for i in range(len(parts) - 1, 0, -1):
+      parent = ".".join(parts[:i])
+      mods.setdefault(parent, _package(parent))
+      setattr(mods[parent], parts[i], mods[".".join(parts[: i + 1])])
+  return mods
+
+
+class _ListHandler(logging.Handler):
+
+  def __init__(self, lines):
+    super().__init__(level=logging.INFO)
+    self.lines = lines
+
+  def emit(self, record):
+    self.lines.append(record.getMessage())
+
+
+def run_driver(env, engine_cls=Engine, driver=bench_driver, worker_cls=TrainerWorker):
+  """Runs `driver.main()` against the stand-ins with STEPS=3, MICRO_STEPS=4 and `env` on top (a value of None
+  unsets the variable). Returns the run: calls (in order), worker.request_ids, log (messages), error (or None)."""
+  run = types.SimpleNamespace(calls=[], worker=None, log=[], error=None)
+  environ = {k: v for k, v in os.environ.items() if k not in DRIVER_ENV}
+  environ.update(TRAINER_ARGV="--worker_id=stand-in", ORCH_ARGV="--model_id=stand-in", STEPS="3", MICRO_STEPS="4")
+  for k, v in env.items():
+    if v is None:
+      environ.pop(k, None)
+    else:
+      environ[k] = v
+  root = logging.getLogger()
+  handler, level = _ListHandler(run.log), root.level
+  root.addHandler(handler)
+  root.setLevel(logging.INFO)
+  try:
+    with (
+        mock.patch.dict(os.environ, environ, clear=True),
+        mock.patch.dict(sys.modules, stand_in_modules(run, engine_cls, worker_cls)),
+    ):
+      try:
+        driver.main()
+      except Exception as e:  # pylint: disable=broad-except
+        run.error = e
+  finally:
+    root.removeHandler(handler)
+    root.setLevel(level)
+  return run
+
+
+def _fwd_bwd(run):
+  return [c for c in run.calls if c[0] == "fwd_bwd"]
+
+
+def _metric_lines(run):
+  return [line for line in run.log if line.startswith(("STEP_LOSS ", "STEP_METRICS "))]
+
+
+def original_loop_failures(run, steps, micro_steps):
+  """How `run` departs from the original loop's call sequence (an empty list when it does not)."""
+  if run.error is not None:
+    return [f"the driver raised {run.error!r}"]
+  fb = _fwd_bwd(run)
+  if len(fb) != steps * micro_steps:
+    return [f"{len(fb)} fwd_bwd calls, want {steps * micro_steps}"]
+  regular, final = fb[0][3], fb[micro_steps - 1][3]
+  failures = []
+  if micro_steps > 1 and regular == final:
+    failures.append("the regular and the final micro-batch are identical, so the check cannot tell them apart")
+  pad_rows = run.worker._trainer.pad_rows  # pylint: disable=protected-access
+  if pad_rows.get(final, 0) == 0 or (micro_steps > 1 and pad_rows.get(regular) != 0):
+    failures.append(
+        f"the last micro-batch must be the final payload (fully padded rows: {pad_rows.get(final)}) and the others "
+        f"the regular one (fully padded rows: {pad_rows.get(regular)})"
+    )
+  want = [("pathways_initialize",), ("initialize",), ("compile",)]
+  for s in range(steps):
+    for m in range(micro_steps):
+      last = m == micro_steps - 1
+      want.append(("fwd_bwd", s, m, final if last else regular, False, last))
+    want += [("update", s, micro_steps), ("drain", s + 1), ("get_metrics", s + 1)]
+  want.append(("stop",))
+  if run.calls != want:
+    i = next((i for i, (a, b) in enumerate(zip(run.calls, want)) if a != b), min(len(run.calls), len(want)))
+    got_i = run.calls[i] if i < len(run.calls) else None
+    want_i = want[i] if i < len(want) else None
+    failures.append(f"engine calls differ from the original loop at call {i}: got {got_i}, want {want_i}")
+  want_ids = [f"bench_{s}_{m}" for s in range(steps) for m in range(micro_steps)]
+  if run.worker.request_ids != want_ids:
+    failures.append(f"request ids {run.worker.request_ids}, want {want_ids}")
+  if len(_metric_lines(run)) != 2 * steps:
+    failures.append(f"{len(_metric_lines(run))} STEP_LOSS/STEP_METRICS lines, want {2 * steps}")
+  stage_lines = [line for line in run.log if "STAGE_BEFORE_UPDATE" in line]
+  if stage_lines:
+    failures.append(f"STAGE_BEFORE_UPDATE lines with the switch off: {stage_lines}")
+  return failures
+
+
+def _index(calls, *head):
+  """Index of the first call that starts with `head`, or -1."""
+  return next((i for i, c in enumerate(calls) if c[: len(head)] == head), -1)
+
+
+def _without_staging(run):
+  """The run's calls without prepare_fwd_bwd_input / wait_for_device, and without fwd_bwd's `prepared` field."""
+  return [c[:4] + c[5:] if c[0] == "fwd_bwd" else c for c in run.calls if c[0] not in ("prepare", "wait_for_device")]
+
+
+def staged_failures(on, off, steps, micro_steps):
+  """How the switch-on run `on` departs from the switch-off run `off` of the same configuration (empty = correct)."""
+  if on.error is not None or off.error is not None:
+    return [f"a run raised: on {on.error!r}, off {off.error!r}"]
+  failures = []
+  on_fb, off_fb = _fwd_bwd(on), _fwd_bwd(off)
+  # Per call: step, micro index, micro-batch digest, apply_optimizer.
+  if [c[1:4] + c[5:] for c in on_fb] != [c[1:4] + c[5:] for c in off_fb]:
+    failures.append("the fwd_bwd sequence (step, micro index, micro-batch, apply_optimizer) differs from off")
+  if on.worker.request_ids != off.worker.request_ids:
+    failures.append("the request ids differ from off")
+  want_prepared = [(s, 0) for s in range(1, steps)]
+  got_prepared = [(c[1], c[2]) for c in on_fb if c[4]]
+  if got_prepared != want_prepared:
+    failures.append(f"prepared inputs trained at (step, micro) {got_prepared}, want {want_prepared}")
+  updates = [(c[1], c[2]) for c in on.calls if c[0] == "update"]
+  if updates != [(s, micro_steps) for s in range(steps)]:
+    failures.append(f"updates (step, micro-batches) {updates}, want one per step after its {micro_steps}")
+  prepares = [c for c in on.calls if c[0] == "prepare"]
+  if len(prepares) != steps - 1:
+    failures.append(f"{len(prepares)} prepare_fwd_bwd_input calls, want {steps - 1}")
+  waits = [c[1] for c in on.calls if c[0] == "wait_for_device"]
+  if waits != list(range(steps - 1)):
+    failures.append(f"wait_for_device at steps {waits}, want one per staged boundary {list(range(steps - 1))}")
+  first_of = {c[1]: c[3] for c in off_fb if c[2] == 0}
+  for s in range(steps - 1):
+    order = (
+        _index(on.calls, "fwd_bwd", s, micro_steps - 1),
+        _index(on.calls, "prepare", s),
+        _index(on.calls, "wait_for_device", s),
+        _index(on.calls, "update", s),
+        _index(on.calls, "fwd_bwd", s + 1, 0),
+    )
+    if not -1 < order[0] < order[1] < order[2] < order[3] < order[4]:
+      failures.append(f"step {s}: want last fwd_bwd < prepare < wait_for_device < update < next fwd_bwd, got {order}")
+    if order[1] >= 0 and on.calls[order[1]][3] != first_of.get(s + 1):
+      failures.append(
+          f"step {s} prepared {on.calls[order[1]][3]}, off trains {first_of.get(s + 1)} first at step {s + 1}"
+      )
+  if _without_staging(on) != _without_staging(off):
+    failures.append("the calls other than prepare_fwd_bwd_input / wait_for_device differ from off")
+  if len(_metric_lines(off)) != 2 * steps or _metric_lines(on) != _metric_lines(off):
+    failures.append("the STEP_LOSS / STEP_METRICS lines differ from off (or are missing)")
+  if sum(line.startswith("STAGE_BEFORE_UPDATE on:") for line in on.log) != 1:
+    failures.append("want exactly one 'STAGE_BEFORE_UPDATE on:' line")
+  stage = [i for i, line in enumerate(on.log) if line.startswith("STAGE_BEFORE_UPDATE step=")]
+  completed = [i for i, line in enumerate(on.log) if line.startswith("completed step: ")]
+  if len(stage) != steps - 1 or len(completed) != steps:
+    failures.append(f"{len(stage)} STAGE_BEFORE_UPDATE step lines and {len(completed)} completed lines")
+  else:
+    for s, i in enumerate(stage):
+      want = rf"STAGE_BEFORE_UPDATE step={s} staged=bench_{s + 1}_0 stage_s=\d+\.\d{{3}} device_wait_s=0\.250"
+      in_place = completed[s] < i < completed[s + 1]
+      if not re.fullmatch(want, on.log[i]) or not in_place:
+        failures.append(f"STAGE_BEFORE_UPDATE line {s} is wrong or out of place: {on.log[i]!r}")
+  return failures
+
+
+class OffTest(unittest.TestCase):
+
+  def test_unset_runs_the_original_loop(self):
+    self.assertEqual(original_loop_failures(run_driver({"STAGE_BEFORE_UPDATE": None}), 3, 4), [])
+
+  def test_off_runs_the_original_loop(self):
+    self.assertEqual(original_loop_failures(run_driver({"STAGE_BEFORE_UPDATE": "off"}), 3, 4), [])
+
+  def test_off_needs_no_hooks(self):
+    run = run_driver({"STAGE_BEFORE_UPDATE": "off"}, engine_cls=EngineWithoutHooks)
+    self.assertEqual(original_loop_failures(run, 3, 4), [])
+
+
+class OnTest(unittest.TestCase):
+
+  def _check(self, steps, micro_steps):
+    env = {"STEPS": str(steps), "MICRO_STEPS": str(micro_steps)}
+    off = run_driver({**env, "STAGE_BEFORE_UPDATE": "off"})
+    on = run_driver({**env, "STAGE_BEFORE_UPDATE": "on"})
+    self.assertEqual(original_loop_failures(off, steps, micro_steps), [])
+    self.assertEqual(staged_failures(on, off, steps, micro_steps), [])
+    return on, off
+
+  def test_same_micro_batches_order_and_one_update_per_step(self):
+    on, _ = self._check(3, 4)
+    self.assertEqual([c[1] for c in on.calls if c[0] == "prepare"], [0, 1])
+
+  def test_one_micro_batch_per_step_stages_the_final_one(self):
+    on, off = self._check(3, 1)
+    final = _fwd_bwd(off)[0][3]
+    self.assertEqual([c[3] for c in on.calls if c[0] == "prepare"], [final, final])
+
+  def test_last_step_stages_nothing(self):
+    on, off = self._check(1, 3)
+    self.assertEqual([c for c in on.calls if c[0] in ("prepare", "wait_for_device")], [])
+    self.assertEqual(on.calls, off.calls)
+
+
+class RefusalTest(unittest.TestCase):
+
+  def test_engine_without_hooks_is_refused_before_the_compile(self):
+    run = run_driver({"STAGE_BEFORE_UPDATE": "on"}, engine_cls=EngineWithoutHooks)
+    self.assertIsInstance(run.error, RuntimeError)
+    for word in ("prepare_fwd_bwd_input", "wait_for_device", "AI-Hypercomputer/maxtext#5718", "EngineWithoutHooks"):
+      self.assertIn(word, str(run.error))
+    self.assertEqual(run.calls, [("pathways_initialize",), ("initialize",)])
+
+  def test_engine_with_one_hook_is_refused(self):
+    run = run_driver({"STAGE_BEFORE_UPDATE": "on"}, engine_cls=EngineWithPrepareOnly)
+    self.assertIsInstance(run.error, RuntimeError)
+    self.assertIn("lacks ['wait_for_device']", str(run.error))
+    self.assertEqual(run.calls, [("pathways_initialize",), ("initialize",)])
+
+  def test_fused_update_worker_is_refused_at_the_first_boundary(self):
+    run = run_driver({"STAGE_BEFORE_UPDATE": "on"}, worker_cls=FusedUpdateTrainerWorker)
+    self.assertIsInstance(run.error, RuntimeError)
+    self.assertIn("google/tunix#2719", str(run.error))
+    self.assertEqual([c[0] for c in run.calls if c[0] in ("prepare", "wait_for_device")], [])
+    self.assertEqual(run.calls[-1], ("update", 0, 4))  # The worker's own update, inside step 0's last fwd_bwd.
+
+  def test_fused_update_worker_runs_with_the_switch_off(self):
+    run = run_driver({"STAGE_BEFORE_UPDATE": "off"}, worker_cls=FusedUpdateTrainerWorker)
+    self.assertIsNone(run.error)
+    # The worker's update inside the last fwd_bwd, then the driver's own update(), which has nothing to apply.
+    updates = [c[1:] for c in run.calls if c[0] == "update"]
+    self.assertEqual(updates, [(0, 4), (1, 0), (1, 4), (2, 0), (2, 4), (3, 0)])
+
+  def test_other_values_are_refused_before_pathways(self):
+    for value in ("yes", "ON", "1", ""):
+      run = run_driver({"STAGE_BEFORE_UPDATE": value})
+      self.assertIsInstance(run.error, ValueError, value)
+      self.assertIn("STAGE_BEFORE_UPDATE must be on or off", str(run.error))
+      self.assertEqual(run.calls, [], value)
+
+
+class CheckerTest(unittest.TestCase):
+  """Each check fails on a run that is wrong in the way it guards against."""
+
+  @classmethod
+  def setUpClass(cls):
+    super().setUpClass()
+    cls.off = run_driver({"STAGE_BEFORE_UPDATE": "off"})
+    cls.on = run_driver({"STAGE_BEFORE_UPDATE": "on"})
+
+  def test_original_loop_check_rejects_a_staged_run(self):
+    failures = original_loop_failures(self.on, 3, 4)
+    self.assertTrue(any("engine calls differ" in f for f in failures), failures)
+    self.assertTrue(any("STAGE_BEFORE_UPDATE lines" in f for f in failures), failures)
+
+  def test_original_loop_check_rejects_the_wrong_micro_batch(self):
+    regular = _fwd_bwd(self.off)[0][3]
+    run = types.SimpleNamespace(**vars(self.off))
+    run.calls = [c[:3] + (regular,) + c[4:] if c[:3] == ("fwd_bwd", 1, 3) else c for c in self.off.calls]
+    failures = original_loop_failures(run, 3, 4)
+    self.assertTrue(any("engine calls differ from the original loop at call" in f for f in failures), failures)
+
+  def test_original_loop_check_rejects_identical_micro_batches(self):
+    regular = _fwd_bwd(self.off)[0][3]
+    run = types.SimpleNamespace(**vars(self.off))
+    run.calls = [c[:3] + (regular,) + c[4:] if c[0] == "fwd_bwd" else c for c in self.off.calls]
+    failures = original_loop_failures(run, 3, 4)
+    self.assertTrue(any("cannot tell them apart" in f for f in failures), failures)
+
+  def test_staged_check_rejects_a_run_that_never_stages(self):
+    failures = staged_failures(self.off, self.off, 3, 4)
+    self.assertTrue(any("prepared inputs trained" in f for f in failures), failures)
+    self.assertTrue(any("prepare_fwd_bwd_input calls" in f for f in failures), failures)
+
+  def test_staged_check_rejects_a_different_micro_batch(self):
+    run = types.SimpleNamespace(**vars(self.on))
+    run.calls = [("fwd_bwd", 1, 0, "0" * 12, True, False) if c[:3] == ("fwd_bwd", 1, 0) else c for c in self.on.calls]
+    failures = staged_failures(run, self.off, 3, 4)
+    self.assertTrue(any("fwd_bwd sequence" in f for f in failures), failures)
+
+  def test_original_loop_check_rejects_swapped_payloads(self):
+    regular, final = _fwd_bwd(self.off)[0][3], _fwd_bwd(self.off)[3][3]
+    swap = {regular: final, final: regular}
+    run = types.SimpleNamespace(**vars(self.off))
+    run.calls = [c[:3] + (swap[c[3]],) + c[4:] if c[0] == "fwd_bwd" else c for c in self.off.calls]
+    failures = original_loop_failures(run, 3, 4)
+    self.assertTrue(any("must be the final payload" in f for f in failures), failures)
+
+  def test_staged_check_rejects_an_extra_wait(self):
+    run = types.SimpleNamespace(**vars(self.on))
+    i_update = _index(self.on.calls, "update", 2)
+    run.calls = self.on.calls[:i_update] + [("wait_for_device", 2, 4)] + self.on.calls[i_update:]
+    failures = staged_failures(run, self.off, 3, 4)
+    self.assertTrue(any("wait_for_device at steps" in f for f in failures), failures)
+
+  def test_staged_check_rejects_a_wrong_logged_wait(self):
+    run = types.SimpleNamespace(**vars(self.on))
+    run.log = [line.replace("device_wait_s=0.250", "device_wait_s=0.000") for line in self.on.log]
+    failures = staged_failures(run, self.off, 3, 4)
+    self.assertTrue(any("is wrong or out of place" in f for f in failures), failures)
+
+  def test_staged_check_rejects_staging_after_the_update(self):
+    run = types.SimpleNamespace(**vars(self.on))
+    calls = list(self.on.calls)
+    i_prep = next(i for i, c in enumerate(calls) if c[0] == "prepare" and c[1] == 0)
+    i_update = next(i for i, c in enumerate(calls) if c[0] == "update" and c[1] == 0)
+    calls.insert(i_update, calls.pop(i_prep))  # update now precedes prepare
+    calls.insert(i_update, calls.pop(i_prep))  # and wait_for_device
+    run.calls = calls
+    failures = staged_failures(run, self.off, 3, 4)
+    self.assertTrue(any("step 0: want last fwd_bwd < prepare" in f for f in failures), failures)
+
+
+if __name__ == "__main__":
+  unittest.main()