Repository navigation
[DO NOT MERGE] perf(rl_trainer_bench): opt-in STAGE_BEFORE_UPDATE step boundary (-7.2% step time on the 397B bench; stacked on #5701) - #5719
Conversation
A standalone timing harness for the RL trainer program of a Tunix GRPO e2e run: the same MaxText and libtpu flags, mesh, GRPO loss and chunked log-prob pass on one trainer slice, driven through Tunix's TrainerWorker with a synthetic micro-batch whose MoE routing load is calibrated to a production run, and without the orchestrator or rollout engines. render_e2e.py turns the DRY_RUN output of the e2e launcher into the bench JobSet, keeping the trainer JobSet's scheduling fields and Pathways containers and replacing only names, maxRestarts and the client command, with knobs for MaxText flags, libtpu flags, trainer and orchestrator argv, trainer slice and profiler window. Also included: an AOT memory check, a structured render diff, a MaxText config diff, CPU checks of the payload and the numerics knobs, a static HLO collective-overlap summary, and CPU unit tests of the renderer (on a hand-written fixture) and the payload builder.
…e DRY_RUN context switch
The bench is standalone; the e2e names only meant "derived from the e2e
launch's DRY_RUN render". Renames, no behaviour change:
- E2E_RENDER -> LAUNCH_RENDER; --e2e-render -> --launch-render
(render_bench.py and check_numerics_cpu.py)
- render_e2e.{sh,py} -> render_bench.{sh,py}
- e2e_driver.py -> bench_driver.py (also /tmp/bench_driver.py in the pod)
- E2E_TRAINER_ARGV / E2E_ORCH_ARGV -> TRAINER_ARGV / ORCH_ARGV
- aot_e2e.sh -> aot_bench.sh (default log logs/aot-bench-<time>.log)
- tests/render_e2e_test.py -> tests/render_bench_test.py;
tests/testdata/e2e_dryrun_fixture.yaml -> launch_dryrun_fixture.yaml
- heredoc E2E_DRIVER_EOF -> BENCH_DRIVER_EOF; log tag [E2EBench] ->
[TrainerBench]; RenderE2ETest -> RenderBenchTest
README states that the bench is standalone and that LAUNCH_RENDER is the
full launcher's DRY_RUN output. Prose "e2e" meaning the full RL run is kept.
STAGE_BEFORE_UPDATE=on (default off) times the step boundary of a trainer that stages the next step's first micro-batch before the previous step's update, as Tunix's TrainerWorker.stage_fwd_bwd does (google/tunix#2877). After step k's last fwd_bwd and before its update, step k+1's first micro-batch goes through the engine's prepare_fwd_bwd_input, so its host-to-device copy runs while the device still computes step k's last micro-batch. The driver then waits for the device (wait_for_device), runs the update, the drain and the metrics read as before, and step k+1's first fwd_bwd takes the prepared input. Off, the loop is the original one: bench_driver.py only gains lines (no existing line changes), and each added line is either behind the switch or leaves `staged` None. On, the engine must have prepare_fwd_bwd_input and wait_for_device (#5718). Without them the driver raises right after the trainer worker initializes, before the compile. A trainer worker that runs the update inside the step's last fwd_bwd (Tunix's fused update, google/tunix#2719) leaves no update to stage before; the driver sees the engine's train_step move and raises at the first boundary. Any value other than on/off raises before Pathways starts. Same-pod A/B on the 397B recipe (128 chips, MICRO_STEPS=9, steps 1-4): 107.727 -> 100.007 s/step. The step's second fwd_bwd call loses its 7.5-7.8 s extra; regular micro-batches and the step tail are unchanged. tests/stage_before_update_test.py runs main() against stand-ins for Pathways, JAX, Tunix and the engine. Off reproduces the original loop's engine calls, on trains the same micro-batches in the same order with one update per step, and an engine without the two methods or a fused-update worker is refused.
There was a problem hiding this comment.
Code Review
This pull request introduces a standalone RL trainer benchmark harness (rl_trainer_bench) designed to time and test the RL trainer program of a Tunix GRPO end-to-end run independently. It includes a driver script, rendering utilities, ahead-of-time memory checking tools, static schedule analyzers, and associated unit tests. The code review feedback focuses on improving robustness, such as handling edge cases in random choice population sizes, safely accessing metric attributes, initializing loop variables to prevent static analysis warnings, validating YAML inputs, and ensuring safe string splitting during log parsing. All review comments are valid and should be addressed.
| 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 [] |
There was a problem hiding this comment.
When real_len is small (e.g., in testing or scaled-down configurations) or n_seg is large, the population size of np.arange(64, real_len, 64) can be smaller than n_seg - 1. This causes rng.choice to raise a ValueError because replace=False is specified. We should add a fallback to a smaller step size (e.g., 1) or allow replacement with replace=True if the population is still too small.
| 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 [] | |
| for r in range(rows - n_pad_rows): | |
| n_seg = max(1, len(np.unique(seg0[r][seg0[r] > 0]))) | |
| if n_seg > 1: | |
| population = np.arange(64, real_len, 64) | |
| if len(population) < n_seg - 1: | |
| population = np.arange(1, real_len) | |
| if len(population) < n_seg - 1: | |
| cuts = np.sort(rng.choice(population, size=n_seg - 1, replace=True)) if len(population) > 0 else [] | |
| else: | |
| cuts = np.sort(rng.choice(population, size=n_seg - 1, replace=False)) | |
| else: | |
| cuts = [] |
| tot = dens.sum() + (wm.eps or 0.0) | ||
| if wm.min_denom is not None: | ||
| tot = max(tot, wm.min_denom) |
There was a problem hiding this comment.
Directly accessing wm.eps and wm.min_denom can raise an AttributeError if the underlying WeightedMetric implementation in MaxText changes or if a different version is used. Using getattr with a default value is safer and more robust.
| tot = dens.sum() + (wm.eps or 0.0) | |
| if wm.min_denom is not None: | |
| tot = max(tot, wm.min_denom) | |
| tot = dens.sum() + (getattr(wm, "eps", None) or 0.0) | |
| min_denom = getattr(wm, "min_denom", None) | |
| if min_denom is not None: | |
| tot = max(tot, min_denom) |
| 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. |
There was a problem hiding this comment.
Initializing train_step_at_start before the loop ensures it is always defined in the scope, which avoids potential static analysis warnings (like possibly-used-before-assignment) and allows us to remove the inline pylint disable comment later.
| 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. | |
| staged = None # STAGE_BEFORE_UPDATE: step k+1's first micro-batch as prepare_fwd_bwd_input returned it | |
| train_step_at_start = None | |
| 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. |
| 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 |
There was a problem hiding this comment.
Since train_step_at_start is now initialized before the loop, we can safely remove the # pylint: disable=possibly-used-before-assignment comment.
| if engine.train_step != train_step_at_start: # pylint: disable=possibly-used-before-assignment | |
| if engine.train_step != train_step_at_start: |
| 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) |
There was a problem hiding this comment.
If train_yaml or orch_yaml is empty or invalid, yaml.safe_load can return None or a non-dictionary type, which will cause a cryptic TypeError later when accessing keys. It is safer to validate that both loaded objects are dictionaries.
| 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) | |
| 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) | |
| if not isinstance(train, dict) or not isinstance(orch, dict): | |
| sys.exit("Error: train-yaml and orch-yaml must be valid YAML dictionaries") |
| mems = re.findall(r"Memory analysis: CompiledMemoryStats\(([^)]*)\)", s) | ||
| for name, c, m in zip(("fwd_bwd", "accumulate", "update"), costs, mems): | ||
| d = dict(kv.split("=") for kv in m.split(", ")) |
There was a problem hiding this comment.
If m contains any elements without an = sign (e.g., due to trailing commas or format changes), kv.split("=") will return a single-element list, causing dict() to raise a ValueError. Adding a check to ensure "=" is present in kv makes the parsing more robust.
| 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(", ")) | |
| mems = re.findall(r"Memory analysis: CompiledMemoryStats\(?([^)]*)\)", s) | |
| for name, c, m in zip(("fwd_bwd", "accumulate", "update"), costs, mems): | |
| d = dict(kv.split("=", 1) for kv in m.split(", ") if "=" in kv) |
TL;DR
STAGE_BEFORE_UPDATE(defaultoff), to Add an RL trainer bench under experimental/rl_trainer_bench #5701'sbench_driver.py. With iton, step k+1's first micro-batch is prepared on the mesh (prepare_fwd_bwd_input) before step k'supdate(). Its host-to-device copy then runs while the device still computes step k's last micro-batch, instead of on an idle device after the drain.onneedsprepare_fwd_bwd_inputandwait_for_devicefrom feat(training_engine): add prepare_fwd_bwd_input and wait_for_device to overlap the next step's input copy (-7.2% s/step, 397B RL trainer bench, 128 v7x) #5718. If the engine lacks them, the driver stops before the compile with aRuntimeErrorthat names them. It also refuses a Tunix whose worker runs the update inside the lastfwd_bwd(perf(orchestrator): fuse final microbatch fwd_bwd+update RPC and defer get_metrics off weight-sync path google/tunix#2719), since nothing could be staged before it there.a03bcd4.How it works
In #5701's loop, step k ends with its last
fwd_bwd, thenupdate(), then the throttler drain (the end of the step ruler), then the metrics read. Step k+1's firstfwd_bwdthen maps its payload and copies it to the device while the device has nothing else to do. With replayed routing that payload is large: here the loggedrouted_expertsalone is(64, 65536, 60, 10)int16, 5.03 GB across the slice.The switch moves that mapping and copy earlier, and adds one host wait (
wait_for_device) beforeupdate(). The update, the drain and the metrics read stay where they are.sequenceDiagram participant D as bench driver (host) participant T as trainer (devices) alt STAGE_BEFORE_UPDATE=off (stock) D->>T: fwd_bwd(step k, last micro-batch) D->>T: update(), drain, read metrics Note over T: idle D->>T: fwd_bwd(step k+1, first): map and copy, then compute else STAGE_BEFORE_UPDATE=on D->>T: fwd_bwd(step k, last micro-batch) D->>T: prepare_fwd_bwd_input(step k+1, first): copy overlaps the last micro-batch D->>T: wait_for_device() D->>T: update(), drain, read metrics D->>T: fwd_bwd(prepared): compute starts at once endThe trainer sees the same micro-batches, in the same order, with the same
apply_optimizerflags and oneupdateper step. Each staged boundary logs one line after step k'scompleted stepline:STAGE_BEFORE_UPDATE step=k staged=bench_<k+1>_0 stage_s=… device_wait_s=….Results
MEASURED. Same pod and cell; the two arms ran back to back, switch off first. Steps 1-4; step 0 includes first execution and is not scored.
fwd_bwdcall minus a regular one, sScoring:
Details are below.
Changes
bench_driver.py: eight inserted blocks, and no existing line changes (diffagainst70d6444: 8 insert hunks, +60 lines, 0 removed):on/offraises before Pathways initializes;prepare_fwd_bwd_inputandwait_for_device, and raising aRuntimeErrorthat names feat(training_engine): add prepare_fwd_bwd_input and wait_for_device to overlap the next step's input copy (-7.2% s/step, 397B RL trainer bench, 128 v7x) #5718 if not;staged = None;train_step;fwd_bwdand beforeupdate(), three things:train_stepalready moved (a fused update);prepare_fwd_bwd_inputof the next step'sregularpayload (finalwhenMICRO_STEPSis 1);wait_for_device();STAGE_BEFORE_UPDATE step=…log line.tests/stage_before_update_test.py(new, 20 tests, NumPy only), described under Testing.README.md: a new "Step boundary: STAGE_BEFORE_UPDATE" section; the Tests section and the Files table mention the new test.render_bench.sh,render_bench.py:STAGE_BEFORE_UPDATEis added to theDRIVER_ENVknob list. These are comments only:DRIVER_ENValready passes anyK=Vto the driver.Why a stacked PR. #5701's head branch,
sivaibhav/rl-trainer-bench, lives in the forkultrons/maxtext(GitHub reportsisCrossRepository: true), so a PR into that branch cannot be opened in this repository.atwigg/mlperf. It carries Add an RL trainer bench under experimental/rl_trainer_bench #5701's three commits unchanged (9b706a3,d6db22f,70d6444, which is Add an RL trainer bench under experimental/rl_trainer_bench #5701's head today) plus one commit,a03bcd4.Testing
python3 -m pytest src/maxtext/experimental/rl_trainer_bench/tests(Python 3.12.8, NumPy 2.1.3, pytest 9.1.1)• on vs off: same calls, digests, request ids and STEP lines, with one prepare and one wait per boundary, and the logged wait matching the engine's;
• refusals: no hooks, one hook, fused-update worker, bad value;
• 9 checker self-tests, each showing a check fail on a wrong run
• stages
final: 1 test fails• trains staged and plain: 3 fail
• stages after
update(): 8 fail• prepared input unused: 2 fail
• never stages: 6 fail
• fused-update guard removed: 1 fails (the fused-worker test)
• #5701's stock head driver: 10 fail (3 on, 4 refusal, 3 checker tests); its 3 off tests pass
The PR driver passes 20 of 20.
70d6444), through the same harness (one-off script, not in the PR), 3 shapes (STEPS×MICRO_STEPS 3×4, 3×1, 1×3) × switch unset/offTrainerWorker; tiny Qwen3.5 MoE on 4 CPU devices with router replay andretry_when_tokens_dropped. Stand-ins replace only Pathways, the argv parsers, the warm-up payload and the tokenizertraining_engine/files (one declared log-path normalization, see details). The fused-update guard fired only in the deliberately misordered driver. Each run covers:• off is identical to stock (trace, log, bit-identical params);
• on vs off: same micro-batches, order and updates, with bit-identical params and STEP lines;
• a pending overflow check at every prepare;
• five broken drivers and two refusals caught
DRIVER_ENV="STAGE_BEFORE_UPDATE=on"on the bench's DRY_RUN fixture gives… MICRO_STEPS=9 STAGE_BEFORE_UPDATE=on exec python /tmp/bench_driver.py, and the embedded driver is byte-identical to this filebench_driver.pyhas exactly the same 12 messages as at70d6444(none new); the new test scores 10.00/10.• pyink 24.10.1 (
--pyink-indentation=2 --line-length=122): clean on the inserted line ranges and on the new file. The rest ofbench_driver.pyis left as #5701 wrote it.• codespell: clean.
• #5701's 3-line rename (
70d6444);• the docstring and refusal-message text;
• pyink's wrapping of two log calls;
• the fused-update guard, which was added after the run. On that image's non-fused worker the guard is a comparison that never fires; the CPU proof above exercises it on the real engine.
Risks
on. Without it, the switch refuses after the trainer worker has initialized (the model is built, nothing is compiled), with a message that names the two methods and feat(training_engine): add prepare_fwd_bwd_input and wait_for_device to overlap the next step's input copy (-7.2% s/step, 397B RL trainer bench, 128 v7x) #5718. Off needs nothing new: the tests run off against an engine without the methods (SOURCE; MEASURED on CPU).l3m, 128 chips), staging after the update left 7.43 s of boundary gap, against 0.02 s when staging before it.fwd_bwd, next to the last micro-batch and the update. Arm 2 ran all 5 steps at this cell without running out of memory (MEASURED). A cell with less headroom may not fit.wait_for_device()blocks the host beforeupdate()is dispatched. At this cellretry_when_tokens_droppedis on, so the stockupdate()already waits for the last micro-batch's overflow flag first (SOURCE). The extra wait therefore costs nothing here: the tail is 12.556 vs 12.548 s. Without dropless replay, the stockupdate()would be dispatched while the last micro-batch still runs, and the wait could delay it. That case is unmeasured.worker._trainer, as the driver already does for_configand_throttler, and passes MaxText'sPreparedFwdBwdInputthroughTrainRequest.payload; theTrainerWorkerhands that to the engine unchanged (SOURCE). Once [DO NOT MERGE] perf(orchestrator): opt-in pipelined step boundary that stages the next micro-batch before the deferred update google/tunix#2877 lands, the driver could callTrainerWorker.stage_fwd_bwdinstead.TrainerWorker.fwd_bwd(apply_optimizer=True)runsupdate()itself, so no update is left to stage before. The driver detects that from the engine'strain_stepand raises at the first boundary (MEASURED on CPU with a fused-worker stand-in). With the switch off, Add an RL trainer bench under experimental/rl_trainer_bench #5701's loop behaves there as it does without this PR.LOGP_GATESon, so loss and gradient norm are 0 in both arms (Add an RL trainer bench under experimental/rl_trainer_bench #5701's README explains why). That makes X3 (equal loss denominators) vacuous. The inputs were matched instead through the per-micro-batchmult_prob_error_max(X3b). Numerics with moving weights come from the CPU proofs (bit-identical parameters).Related
prepare_fwd_bwd_input/wait_for_device, the MaxText engine methods this switch calls.TrainerWorker.stage_fwd_bwdbefore the deferred update). This switch emulates that boundary on the bench.Full evidence
The run
parambole-l3bon bookingr-06ca7; one pod; arms in ordero275-hd(switch off), theno275-hd-l3(switch on).a9f8754768ca633068545a2dbd2b32c0. The arms differ only inSTAGE_BEFORE_UPDATE: their AOT env files differ only in the arm comment and that value.tpu7x:4x4x8).gcr.io/tpu-vm-gke-testing/sivaibhav/trellis-experimental@sha256:b4647c188b25638e644c7a6fc53fb8e0751741525a277a1bcf07dca92ba9a6ac; Pathways server662f33b0, proxye714351c.TrainerWorkerhas no fused update.num_moe_token_chunks=2,context_parallel_strategy=all_gather,sa_fuse_reciprocal=true,sa_use_base2_exp=false,gdn_cp_state=device,moe_route=device,gdn_cp_mode=seq,enable_checkpointing=false.TUNIX_LOGPS_SHARDED_HEAD=local).ragged_buffer_factor=2.75with dropless replay,MICRO_STEPS=9,STEPS=5,LOGP_GATESon.training_engine/, overlaid from2d821dfc1(an earlier push of feat(training_engine): add prepare_fwd_bwd_input and wait_for_device to overlap the next step's input copy (-7.2% s/step, 397B RL trainer bench, 128 v7x) #5718's branch);maxtext_engine.pymd5e0afbda1.c79bcfd: the two patch series differ only in docstring text (SOURCE).e2e_driver.pyatd6db22f(md5e78c082f) plus seven inserted blocks (md54e893e74).40069a1a) differs from it in five ways: the three lines70d6444renamed (log tag,TRAINER_ARGV,ORCH_ARGV), the docstring paragraph, the refusal message, pyink's wrapping of twologging.infocalls, and the fused-update guard (two blocks behind the switch, added after the run).gs://parambole-397b-replay-nap/win-tpu/bench/r-06ca7/parambole-l3b/o275-hd.log(md5caa615cda21cc0265cee34db70eec5d1, 150,322 B)…/o275-hd-l3.log(md5e2542d86d49ebe209a8c108bbf6a3d0a, 151,081 B)Step ruler lines, verbatim (driver tag
[E2EBench]before #5701's rename)dispatch s[1]minus the median ofdispatch s[2:]. It was 7.695 / 7.636 / 7.535 / 7.762 s with the switch off and −0.050 / −0.054 / −0.045 / −0.040 s with it on.Pre-registration
PREREG-L3B.json(md5bffdb963…), written at 12:39:58Z, before either arm ran.PREREG-L3B-AMENDMENT-CHECK.json, written at 13:27Z, still before the run. It added downgrade-only guards to the scorer, the expected benign P5 miss, and the order controls quoted under Risks.score_l3b.py(md50e98c7ec), outputscore.json(md55e8efe38).device_wait_s8.0 [2.0, 10.95], scored as > 0.5;stage_s1.5 [0, 8.5]Exact checks, all hit:
online and one stage per step 0-3, each stagingbench_<k+1>_0micro_denominator(vacuous here: all 0 withLOGP_GATESon); 0 drops and 0 replays in both armsmult_prob_error_maxin both arms (the input check that is not vacuous)device_wait_s> 0.5 s at all 4 stagesReading: HIDES-ON-BENCH, under both the pre-registered rule and the guarded rule.
CPU proofs
Real-stack matrix,
l3b_cpuproof.py, in a new directory on cmcs (~/u10-eng/PRL3BENCH):TrainerWorker(md510c859d8, no fused update), running a 1-layer, 4-expert, top-2 Qwen3.5 MoE on 4 host devices with forced routing andretry_when_tokens_dropped.Drivers: this PR's (
40069a1a) against #5701's head (a2fdaf1a). Checks:=off, equals stock. The engine trace with input digests is identical, the log is identical modulo time, and the final parameters are bit-identical. Run without and with every micro-batch replayed dropless.fwd_bwdsequence, digests andapply_optimizerflags. The prepared input arrives exactly at micro 0 of steps 1..S-1. Each prepare sits between the step's lastfwd_bwdand itsupdate, withwait_for_devicein between, and its digest equals the off run's next first micro-batch. STEP lines and parameters are bit-identical. Shapes: S4×M3 (with and without replay), S3×M9, S3×M1.final, train twice, stage after the update, staged input unused, never stage;=yes, are refused with 0 trainer calls;e0afbda1(the engine that ran on TPU)VERDICT.jsonmd5bb0cd3fc). The fused-update guard fired only in the misordered driver (update before stage), which it stopped at its first boundary; it never fired in a correct run16e52650(#5718 head'straining_engine/files, with #5568'smetrics.py)VERDICT.jsonbytes,bb0cd3fc). The unmodified scorer reads 15/18, and its 3 misses are K1 log comparisons that differ only in that pathThe same matrix on the driver before the fused-update guard (
877b3980) also read 18/18 on both engines. The newer engine logs a TensorBoard path that contains the run directory, so its scorer gets one declared normalization line for that directory name. The same scorer still reads the first engine's matrix unchanged.Reproduce on #5701's renderer
DRIVER_ENV="STAGE_BEFORE_UPDATE=on", using an image whose MaxText has feat(training_engine): add prepare_fwd_bwd_input and wait_for_device to overlap the next step's input copy (-7.2% s/step, 397B RL trainer bench, 128 v7x) #5718 and whose Tunix does not fuse the update.completed steplines for steps 1-4 and theSTAGE_BEFORE_UPDATE step=lines.Our A/B ran both arms in one pod with our own two-arm wrapper around the same driver, to keep pod-to-pod variation out of the delta.