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..86cdf794c1
--- /dev/null
+++ b/src/maxtext/experimental/rl_trainer_bench/README.md
@@ -0,0 +1,295 @@
+<!--
+ 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. 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
+
+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. `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
+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 |
+|---|---|
+| `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_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 |
+| `tests/` | CPU unit tests of the renderer, the payload builder and the driver's step loop |
+
+## 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 > launch-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. 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
+
+```bash
+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 `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_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
+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 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
+`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 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
+`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 `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
+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 took about two minutes each of host time in our runs, 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 --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.
+
+## 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
+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_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
+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/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.
+
+## 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_bench.sh b/src/maxtext/experimental/rl_trainer_bench/aot_bench.sh
new file mode 100755
index 0000000000..7e7fdb3e1e
--- /dev/null
+++ b/src/maxtext/experimental/rl_trainer_bench/aot_bench.sh
@@ -0,0 +1,47 @@
+#!/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.
+#
+#   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_bench.sh [extra MaxText key=value ...]
+#
+# 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)"
+: "${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_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-bench-$(date -u +%Y%m%dT%H%M%SZ).log}
+mkdir -p "$(dirname "$LOG")"
+echo "log: $LOG"
+{ 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" \
+  -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..d4b65072d8
--- /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_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]
+"""
+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..dc4b096037
--- /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_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),
+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/bench_driver.py b/src/maxtext/experimental/rl_trainer_bench/bench_driver.py
new file mode 100644
index 0000000000..68a96d98ad
--- /dev/null
+++ b/src/maxtext/experimental/rl_trainer_bench/bench_driver.py
@@ -0,0 +1,564 @@
+#!/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.
+
+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
+import os
+import shlex
+import sys
+import time
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - [TrainerBench] %(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["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.
+  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"))
+  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
+  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())
+  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
+  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}")
+
+  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
+    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],
+    )
+    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:
+      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/check_numerics_cpu.py b/src/maxtext/experimental/rl_trainer_bench/check_numerics_cpu.py
new file mode 100644
index 0000000000..6b98b59a9c
--- /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 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 (--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
+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 --launch-render /r/<launch-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 bench_driver as d  # pylint: disable=g-import-not-at-top
+
+
+def orch_argv_from_render(path, fsdp=None):
+  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:
+    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("--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.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)
+  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..8d5ce39676
--- /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 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
+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 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
+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/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_bench.py b/src/maxtext/experimental/rl_trainer_bench/render_bench.py
new file mode 100644
index 0000000000..7faf5868ad
--- /dev/null
+++ b/src/maxtext/experimental/rl_trainer_bench/render_bench.py
@@ -0,0 +1,470 @@
+#!/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 (`--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`).
+
+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
+    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
+    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_bench.sh instead of rendering.
+"""
+
+from __future__ import annotations
+
+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"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 "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/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/bench_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_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_')))}",
+  ]
+  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("--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__)), "bench_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 (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 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)")
+  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_bench.sh instead of rendering")
+  return p
+
+
+def main(argv=None):
+  a = build_parser().parse_args(argv)
+  if a.launch_render:
+    if a.train_yaml or a.orch_yaml:
+      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 --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:
+      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_bench.sh b/src/maxtext/experimental/rl_trainer_bench/render_bench.sh
new file mode 100755
index 0000000000..40993ac922
--- /dev/null
+++ b/src/maxtext/experimental/rl_trainer_bench/render_bench.sh
@@ -0,0 +1,62 @@
+#!/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 (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:
+#   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)
+# 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 ..."               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}"
+: "${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_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:-}" \
+  --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..615acedca6
--- /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_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
+(TRAINER_ARGV), the orchestrator argv (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_bench 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/bench_driver.py" in l)
+  toks = shlex.split(line)
+  env = {}
+  for t in toks:
+    if t in ("exec", "python", "/tmp/bench_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("TRAINER_ARGV", ""))),
+      "orch_argv": flagmap(shlex.split(env.pop("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..d140a9d41a
--- /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 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.
+"""
+
+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 bench_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_bench_test.py b/src/maxtext/experimental/rl_trainer_bench/tests/render_bench_test.py
new file mode 100644
index 0000000000..f793346664
--- /dev/null
+++ b/src/maxtext/experimental/rl_trainer_bench/tests/render_bench_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_bench.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_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", "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 (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/bench_driver.py" in l)
+  env = {}
+  for t in shlex.split(line):
+    k, eq, v = t.partition("=")
+    if eq:
+      env[k] = v
+  return env
+
+
+class RenderBenchTest(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 = ["--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)
+    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/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["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["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["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["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["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):
+      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", "--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)
+    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/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()
diff --git a/src/maxtext/experimental/rl_trainer_bench/tests/testdata/launch_dryrun_fixture.yaml b/src/maxtext/experimental/rl_trainer_bench/tests/testdata/launch_dryrun_fixture.yaml
new file mode 100644
index 0000000000..4673b7a0f8
--- /dev/null
+++ b/src/maxtext/experimental/rl_trainer_bench/tests/testdata/launch_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
+                )