[parambole]

TL;DR

• Speed (trainer only, on a bench): one optimizer step takes 98.31 s instead of 130.80 s (-24.8%): 3,000 vs 2,255 tok/s/chip, 94% of the 3,193 tok/s/chip pretraining target counting padding slots, about 85% counting real tokens. Bench = Add an RL trainer bench under experimental/rl_trainer_bench AI-Hypercomputer/maxtext#5701: the Qwen3.5-397B-A17B trainer alone on 128 TPU v7x chips, random weights, synthetic sequences.
• What: four trainer changes, numbered as in the table below, on by default only in mlperf_gbs1024_submission.sh: (1) 1 libtpu flag + 6 MaxText settings; (2) 2 precision keys (same switch); (3) a sharded log-prob head: the GRPO loss gathers the vocab table once per micro-batch and each device scores only its own tokens; (4) a staged step boundary: the next step's first micro-batch is copied to the chips while the current step's last one computes. (5) The recipe's trainer defaults move to the bench cell, and MaxText is pinned to atwigg/mlperf + 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) AI-Hypercomputer/maxtext#5718 (the pin is not behind a switch).
• Correctness: changes 2 and 3 passed real-gradient A/B checks on TPU (same pod). Change 1 has speed evidence only. Change 4 passed order checks on TPU; its real-gradient check has not run, and its -7.66 s was measured through a bench-only switch ([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) AI-Hypercomputer/maxtext#5719) that issues the same trainer calls. This PR's orchestrator has run staged on TPU only in a reduced test rig. A faster bundle of 4 MoE settings failed its check and is left out.
• To run: 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) AI-Hypercomputer/maxtext#5718 is merged and this PR is re-pinned to its merge commit e526ea02 (baa52b4c); build the image from this branch and run mlperf_gbs1024_submission.sh start with TUNIX_IMAGE set to the image by digest. The recipe's default image lacks this code and stops at startup.
• Status: draft. No end-to-end RL run has used this stack, and the exact rendered configuration (GBS 1,024, the MaxText pin, the Oct-13 Pathways build) has not run on TPU. Pending on TPU before un-drafting: (1) a real-gradient A/B of change 1 against the team recipe, (2) a real-gradient A/B of change 4, (3) one end-to-end smoke run.

Terms

• Optimizer step: one weight update. On the bench: about 1,152 trajectories packed into 9 micro-batches of 64 rows x 65,536 tokens (at 256 chips: 5 micro-batches of 128 rows), each run forward and backward, then one update. Times are the mean of steps 1-4; step 0 is warm-up.
• tok/s/chip: token slots per optimizer step (micro-batches x rows x 65,536) / chips / seconds, padding included. At 128 chips: 294,912 / seconds per optimizer step.
• Bench cell: mesh fsdp 32 x context 4 x expert 2, REMAT_CONTEXT=device, log-prob chunk 1024, MoE buffer factor 2.75 with dropless retry: a micro-batch whose MoE buffers overflow reruns with a dropless program, so no token is dropped.
• Team recipe: the team's trainer settings before this PR, run at the bench cell. The recipe's old default (fully dropless buffers, context offload, 4 token chunks, no retry, log-prob chunk 512) was not measured here.
• Same pod / cross-pod: A/B arms that ran back to back in one TPU job / in separate jobs.

What changes

#	Change	What it does	Switch (default in the submission recipe)	Bench, s per optimizer step	Numerics on TPU
1	Tuned settings: 1 libtpu flag, 6 MaxText settings	Attention all-gathers K and V over the context axis (not a ring), with 2 Splash kernel options; GDN context parallelism shards the sequence and keeps its carries on device; MoE runs 2 token chunks, not 4; XLA stops offloading copies to the SparseCores	TRAINER_TUNED_FLAGS=true	130.80 -> 122.00, cross-pod (a)	speed only (gate 1)
2	2 precision keys	GDN state composition runs 3 bf16 passes, not 6. Gradients summed by an all-reduce alone are cast to f32 after it	same switch	107.59 -> 106.24, same pod, on top of change 3	PASS
3	Sharded log-prob head	Gathers the vocab table once per micro-batch, scores only each device's own tokens, sums the head-parameter gradients in f32 once	TUNIX_LOGPS_SHARDED_HEAD=local	122.00 -> 107.59, same pod	PASS
4	Staged step boundary	Copies step k+1's first micro-batch to the chips while step k's last one computes; step k's update is held back (deferred) until the copy has started	PIPELINE_STEP_BOUNDARY=1	105.97 -> 98.31, same pod, through the bench-only switch (b)	order checks PASS; real-gradient check pending (gate 2)
5	Recipe defaults and MaxText pin	Moves the recipe's trainer defaults to the bench cell. Pins MaxText to atwigg/mlperf plus AI-Hypercomputer/maxtext#5718	TRAINER_TUNED_FLAGS=true for the defaults	not measured: every row ran at the bench cell	no A/B against the old default; padding exclusion, now on by default, has CPU parity only and no TPU on/off A/B

(a) The measured runs also set moe_route=device. It did nothing in the bench image's MaxText, and this PR leaves it at its default.
(b) 105.97 s is the change 1-3 program without staging, in the same pod. Staged through a bench-driver switch (AI-Hypercomputer/maxtext#5719, do not merge) that makes the same trainer calls in the same order, not through this PR's orchestrator.

The 6 MaxText settings of change 1 are context_parallel_strategy=all_gather, sa_fuse_reciprocal=true, sa_use_base2_exp=false, gdn_cp_mode=seq, gdn_cp_state=device and num_moe_token_chunks=2, next to the libtpu flag --xla_tpu_enable_offloading_copy_to_sparsecore=false. The 2 precision keys of change 2 are gdn_cp_matmul_precision=high and cast_grads_after_all_reduce=true.

How it works

Where each change acts in one optimizer step

flowchart TD
  B1["Step boundary - change 4<br/>this step's first micro-batch was staged<br/>during the previous step"]
  subgraph MB["Each of the 9 micro-batches, forward and backward"]
    direction TB
    MIX["Token mixing, attention or GDN by layer<br/>attention, change 1 - K and V all-gathered<br/>over the context axis, 2 Splash options<br/>GDN, change 1 - sequence-sharded, carries on device<br/>GDN, change 2 - state composed in 3 bf16 passes"]
    MOE["MoE feed-forward in every layer<br/>change 1 - 2 token chunks instead of 4"]
    HEAD["Log-prob head and GRPO loss<br/>change 3 - sharded head"]
    GRAD["Gradient reduction across chips<br/>change 2 - all-reduce-only bf16 gradients<br/>cast to f32 after the all-reduce"]
    MIX --> MOE --> HEAD --> GRAD
  end
  UPD["Optimizer update - change 4<br/>deferred until the next step's<br/>first micro-batch is staged"]
  B2["Next step boundary"]
  XLA["Compiler, every trainer program - change 1<br/>copies not offloaded to the SparseCores"]
  B1 --> MB --> UPD --> B2
  XLA -.-> MB

Loading

One optimizer step: changes 1-3 act inside each of the 9 micro-batches, change 4 at the step boundary. Bench savings: change 1 -8.79 s (speed only), change 3 -14.41 s (PASS), change 2 -1.36 s (PASS), change 4 -7.66 s (numerics pending); 130.80 -> 98.31 s overall.

Change 3: the log-prob head, before and after

flowchart TB
  IN["One micro-batch reaches the head<br/>hidden states, 64 rows x 65,536 tokens<br/>rows over fsdp x expert, tokens over 4 context ranks"]
  subgraph OLD["Before - compute_chunked_logps under GSPMD"]
    direction TB
    OA["Scan 64 chunks of 1,024 tokens<br/>every context rank computes every chunk<br/>4x the matmuls"]
    OB["Every chunk, forward and rematerialized backward<br/>all-gathers the 2 GB vocab table"]
    OC["Every backward chunk<br/>all-reduces the 2 GB bf16 table gradient"]
    OD["Per micro-batch<br/>129 table all-gathers, 65 table all-reduces<br/>533 TFLOP of head matmuls per device"]
    OA --> OB --> OC --> OD
  end
  subgraph NEW["After - sharded head, shard_map with a custom VJP"]
    direction TB
    NA["Each head parameter all-gathered once<br/>in the forward, reused by the backward"]
    NB["Each device scans only its own tokens<br/>1 row x 16,384 tokens in 16 steps of 1,024"]
    NC["Head-parameter gradients summed in f32<br/>on each device, then reduced once"]
    ND["Per micro-batch<br/>2 table all-gathers, 2 table all-reduces<br/>133 TFLOP of head matmuls per device"]
    NA --> NB --> NC --> ND
  end
  IN -->|"switch off or unset"| OLD
  IN -->|"TUNIX_LOGPS_SHARDED_HEAD=local"| NEW

Loading

Per micro-batch, from an AOT compile of the team recipe at the bench cell (128 chips, 256 devices): the old head gathers the 2 GB table 128 times, the new one once; each count in the diagram adds one gather and one all-reduce for the input embedding.

The GRPO loss needs the log-probability of every next token: the hidden state times the vocab table (4,096 x 248,320 in bf16, 2 GB), in chunks of 1,024 tokens. A TPU profile of the old head, on a variant of the team recipe, put it at 1,871.5 ms per micro-batch per TensorCore, against 199.0 ms for the pretraining loss head of the same model. The new head runs the same per-chunk math; its one numerical change is that the vocab-table and final-norm-scale gradients are summed in f32 once, not in bf16 chunk by chunk (XLA compiles that f32 reduce-scatter as an all-reduce plus a slice). It falls back to the old head, logging [logps-head] sharded head off ..., in four cases: no mesh set with jax.set_mesh, no token layout that fits the mesh, a compute_final_logits that reads no model state, or a head parameter with no dimension divisible by the device count.

Change 4: the step boundary, before and after

gantt
    title One step boundary on the bench, in seconds
    dateFormat ss
    axisFormat %S
    todayMarker off
    section Before
    Last micro-batch of step k, 10.75 s : 00, 11
    Update of step k, 1.6 s : 11, 13
    Chips idle 7.7 s, next input copied in : crit, 13, 21
    First micro-batch of step k+1 : 21, 32
    section Staged
    Last micro-batch of step k : 00, 11
    Next input copied in, overlapped : active, 00, 08
    Deferred update of step k : 11, 13
    First micro-batch of step k+1, at once : 13, 24

Loading

Bench, same pod. Bars are rounded to whole seconds. The copy bar is drawn as long as the gap it removes, because the bench does not time the copy itself.

On the bench the gap is mapping and copying step k+1's first micro-batch to the chips (its router-replay routing alone is 5.03 GB: synthetic on the bench, recorded by the rollouts in a run): the call that waits on it took 18.36-18.46 s instead of 10.75 s, and 10.71 s when staged. In the reduced rig, where the orchestrator also closes the step, the chips idled 12.46 s per boundary with the flags off and 0.02 s with them on.

Reduced rig: a 128-chip test setup, not part of this PR: an earlier revision of this PR's orchestrator and trainer worker (see "Corrections to commit messages") drives one MaxText trainer over gRPC on recorded trajectory groups instead of rollouts (3 micro-batches per step, learning rate 0, no weight sync or checkpoint).

Results

On the bench (see TL;DR) at the bench cell, with synthetic sequences of the real shape. 0 MoE tokens dropped and 0 dropless retries in every run below. Each row adds one change.

Setup	s per optimizer step	Change	How it was compared
Team recipe	130.80	-	mean of 3 pods (130.78-130.81)
+ change 1: libtpu flag, 6 MaxText settings (a)	122.00	-8.79	cross-pod
+ change 3: sharded log-prob head	107.59	-14.41	same pod
+ change 2: 2 precision keys	106.24	-1.36	same pod
same program, in the boundary pod	105.97	reference	same pod as the next row
+ change 4: staged step boundary (b)	98.31	-7.66	same pod; a repeat read 98.26

(a), (b): see What changes. The change 1-3 program measured 106.24, 106.21 and 106.04 s in three pods (mean 106.16). Deltas use unrounded means.

• Total: -32.48 s per optimizer step (-24.8%). The stage deltas come from different pods and add up to 32.22 s.
• The boundary saving is only the boundary. A regular micro-batch takes 10.754 s with staging and 10.755 s without (medians).
• 256 chips (fsdp 64 x context 4 x expert 2, 5 micro-batches of 128 rows, changes 1-3 on): 71.47 -> 56.12 s per optimizer step with the staged boundary, same pod. It is the same work as at 128 chips (560 real rows, in 640 row slots instead of 576). With the boundary on at both sizes, 128 -> 256 chips is 1.75x; perfect scaling would be 49.2 s.

tok/s/chip	Team recipe	This PR	Pretraining target
All token slots, padding included	2,255 (70.6% of target)	3,000 (93.9%)	3,193
Real tokens only	about 2,039 (63.8%)	about 2,712 (84.9%)	3,193

The pretraining target is the MaxText pretraining run of this model on the same kind and number of chips (128 v7x), with the same mesh and micro-batch shape; its rows had no padding. About 9.6% of the bench's token slots are padding (7.0% in each of the first 8 micro-batches, 30.25% in the last), so both rows are shown.

Correctness

Bench speed runs cannot see gradient errors: the bench input's rollout log-probs are synthetic, so the loss's log-prob gates (sequence error threshold, truncated importance sampling) drop every sequence of a random-weight model, and loss and grad_norm are 0. A real-gradient A/B turns the gates off (LOGP_GATES=off) and gives every sequence segment a seeded random advantage; all arms run in one pod from the same weights and inputs. The forward log-prob metric is the mean absolute per-token gap between the trainer's log-probs and those rollout log-probs (sampler_is/token_logdiff_absmean). Pass bounds for changes 2 and 3, written before each run (relative, except the forward metric): loss 1e-5, forward metric 5e-5, entropy 1e-4, grad_norm 2e-3; fail above 5e-4, 1e-3 and 1e-2 for the last three. Parentheses give the same quantity between two reference runs: the noise floor.

Change	Check on TPU	Result	Verdict
1. libtpu flag and 6 MaxText settings	None against the team recipe. They ran with real gradients as the reference arms of the head A/B	loss and grad_norm finite; two runs agree to 5.3e-5	speed only (gate 1)
2. 2 precision keys	Same pod, real gradients, step 0: one run without the keys, two with them	grad_norm 1.47e-5 relative (1.92e-5); forward log-prob metric 1.1e-6 (3.6e-7); entropy 4.2e-9 (8.4e-9); loss identical as printed	PASS, measured with the head on
3. Sharded log-prob head	Same pod, real gradients, step 0: head off, on, off again	loss bit-identical; grad_norm +2.78e-4 relative to the mean of the head-off runs (5.32e-5); forward metric +3.7e-6 (3.8e-6); entropy -1.7e-8 relative	PASS
4. Staged step boundary	Exact checks on the bench and in the reduced rig; no real-gradient A/B yet	Same micro-batches in the same order, one update per step, every micro-batch of step k+1 after update k, staged before every update; regular micro-batch and forward metric unchanged	order PASS; numerics pending (gate 2)
Left out: a bundle of 4 MoE settings	Same pod, real gradients	forward metric 9.74e-4, above the 5e-4 fail bound; grad_norm -0.80%, about 420x the spread between two reference runs	FAIL, not in this PR

• The bit-identical loss does not test the head: at an importance ratio of 1 the GRPO loss depends only on advantages and masks, while the gradient goes through the head's log-probs. The +2.78e-4 grad_norm shift fits the head's one numerical change: on CPU only the table and norm-scale gradients differ from the old head's, both closer to an all-f32 reference (0.93x and 0.13x of the old head's error).
• The 4 MoE settings (prefuse_moe_weights=false, moe_accumulate_wi_dlhs=true, moe_accumulate_chunk_wgrad=true, decoder_layer_input=device) fail the same way without the last one, so the cause is prefuse_moe_weights=false and/or the two accumulate flags. The launcher tests fail if the switch sets any of them.

CPU tests on the pushed code (Linux, CPU JAX):

• Launcher: 35 of 35 at the moe_route commit (no later commit changes a recipe). Sharded log-prob head: 40 tests OK, 1 skipped by design (head files unchanged since).
• Orchestrator, queue and worker suites, at "requirements: pin MaxText ...": 1,315 passed; the 28 failures and 1 collection error also fail on the base (atwigg/mlperf at the merged commit 4c7823f20). The cleanup commit then changed 4 of those test files (names, comments and docstrings; 19 diagnostic prints removed; no behavior change) and only comments and docstrings in rl_program.py and trainer_worker.py. The 7 step-boundary test files pass 549 of 549, the same count before and after.
• Shown to fail on deliberately wrong code: the head before the mesh-order fix fails both out-of-order tests (log-probs off by up to 10.65 and 13.99); deleting the head's non-finite guard fails the NaN test; the moe_route test fails on the previous recipe; the measured-cell tests fail on the old recipe. For in-order layouts like the bench's, the mesh-order fix leaves the traced value_and_grad program unchanged (same jaxpr hash, head on and off).
• Step boundary, on [DO NOT MERGE] perf(orchestrator): opt-in pipelined step boundary that stages the next micro-batch before the deferred update #2877's branch: 599 tests; 31 of 31 deliberately broken hooks caught. With the flags off, the tests compare every trainer call with a frozen copy of the old loop. 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) AI-Hypercomputer/maxtext#5718's tests train bit-identically through this PR's TrainerWorker in the staged order.
• The counts above were taken before the last commit (the step clock). It adds one test, with the boundary off and on; its message reports that the test fails on the previous code in the two cases where the boundary is on and a step waits for its groups.

How to run end to end

Not yet run end to end; see "Gates before un-drafting".

1. Merge 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) AI-Hypercomputer/maxtext#5718.
2. Re-pin (done in baa52b4c): both MaxText lines of requirements/maxtext_requirements.txt (maxtext and maxtext-vllm-adapter) moved from 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) AI-Hypercomputer/maxtext#5718's head commit c79bcfd16 to its squash-merge commit e526ea02 on atwigg/mlperf. The two commits have the same tree (e1330ef9), so the installed code is unchanged.
3. Build, push and launch. Create raiden_wheels/ (not in git; the Dockerfile copies it) and put the pinned Raiden wheel in it, tpu_sync_jax-0.0.1.dev20261013212200, sha256 c3a34a8dc94e6dbabd4aa5e54bde2de6e1c3c5f616bf04c71cc507241d64ece0; a wheel there is installed without a checksum check. Or leave it empty and build on a GCE host whose service account can read the wheel bucket: the bucket refuses unauthenticated reads, so with neither the build fails at the Raiden step. Then:

bash build_docker.sh --maxtext --raiden --deepswe   # --maxtext adds the MaxText vLLM adapter
docker tag tunix_base_image "gcr.io/<project>/<user>/trellis:<tag>"
docker push "gcr.io/<project>/<user>/trellis:<tag>"   # prints the image digest
export TUNIX_IMAGE="gcr.io/<project>/<user>/trellis@sha256:<digest>"
export WANDB_API_KEY="<key>" MAXTEXT_OUTPUT_DIR="gs://<bucket>/trellis/maxtext"
JOB_PREFIX="<prefix>" K8S_NAMESPACE="<namespace>" \
  bash tunix/experimental/examples/recipes/mlperf_gbs1024_submission.sh start

• Pin by digest. Client containers pull with imagePullPolicy: Always, so a moving tag can mix code across restarts. An empty TUNIX_IMAGE falls back to gcr.io/cloud-tpu-multipod-dev/atwigg/trellis:latest (built 2026-09-27), which stops at startup: its orchestrator rejects the new flags and its MaxText does not know 13 of the trainer's keys.
• Launch from a clean shell. Recipe settings are ${VAR:-default}, so an exported value wins. Silent changes: an exported MAXTEXT_EXTRA_FLAGS replaces the recipe's base MaxText list (mesh rule, kernels, buffer factor, chunks; the retry, padding, audit and tuned keys are still appended); TUNIX_LOGPS_SHARDED_HEAD inside TRAINER_EXTRA_ENV overrides the switch; COMPUTE_LOGPS_CHUNK_SIZE=0 turns the head off; a boundary switch set to anything but 1, true or True (for example TRUE) is off.

Switch	Default in mlperf_gbs1024_submission.sh	What it turns on	To turn it off
TRAINER_TUNED_FLAGS	true	Changes 1 and 2, and the bench cell: RAGGED_BUFFER_FACTOR=2.75 with RETRY_WHEN_TOKENS_DROPPED=true, 2 MoE token chunks, REMAT_CONTEXT=device, MOE_EXCLUDE_PADDING=true, LOG_MOE_DROPS=true, COMPUTE_LOGPS_CHUNK_SIZE=1024. Each can be overridden	false: the trainer renders as before this PR (fully dropless buffers, context offload, 4 token chunks, no retry, log-prob chunk 512)
RAGGED_BUFFER_FACTOR	2.75 when TRAINER_TUNED_FLAGS=true, else -1	The buffer factor of the bench cell, with dropless retry	-1: fully dropless buffers, context offload, 4 chunks, no retry, padding or audit flags. Not run on TPU with these settings
TUNIX_LOGPS_SHARDED_HEAD	local	Change 3, in the trainer JobSet only	0 (or false, off, no). An empty value falls back to local in this recipe
PIPELINE_STEP_BOUNDARY	1 (on for 1, true or True)	Change 4: --pipeline_step_boundary, plus --stage_next_microbatch through STAGE_NEXT_MICROBATCH, which follows it unless set	false turns both off. STAGE_NEXT_MICROBATCH=false alone keeps the pipelined boundary without staging

Every other 397B v7x recipe keeps these switches off; the MaxText pin applies to every image built from this branch. STAGE_NEXT_MICROBATCH without PIPELINE_STEP_BOUNDARY is refused. Keep router replay on (RETURN_ROUTED_EXPERTS=true, the default): on the bench at 256 chips, with it off every micro-batch overflowed the 2.75 buffer and was retried dropless, at 156.0 s per optimizer step against 72.7 s with replay on in the same pod (a setup without the 2 precision keys, so not the 71.47 s arm in Results).

What the step-boundary flags do

• --pipeline_step_boundary packs step k+1's first micro-batch while step k's last ones train, and sends it right after step k's trainer-metrics fetch, checkpoint (if one is due) and weight-sync staging; step k's logging and other bookkeeping follow the send. Same trainer calls in the same order, nothing copied early.
• --stage_next_microbatch (needs the flag above) sends that micro-batch to the trainer while the chips still compute step k's last one (TrainerWorker.stage_fwd_bwd); with 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) AI-Hypercomputer/maxtext#5718 the trainer maps it and starts its copy to the chips (prepare_fwd_bwd_input).
• --stage_before_update (default on, no env var) stages between step k's last fwd_bwd(apply_optimizer=True, defer_update=True) and apply_deferred_update(). Under dropless retry update() returns only after the last micro-batch is done, so staging after it hides nothing: staged that way, the rig still idled 7.43 s per boundary.

While an update is deferred, the worker refuses anything that needs the step closed (another micro-batch, an update, metrics, a checkpoint, weight sync, actor log-probs, eval). A stage that fails on its way to the trainer falls back to an unstaged send and turns staging off for the run; a stage that raised on the trainer ends the run, as a failed train call does. Without AI-Hypercomputer/maxtext#5718's methods the boundary runs pipelined but unstaged, with one warning.

Where staging applies in a run. In a 10-step GBS-1,024 run, staging can happen at 6 of the 9 step boundaries (after optimizer steps 1-6). It is refused after steps 7-9, whose checkpoints the MLPerf validation scores (val_start_step 7): a slow stage there could delay the checkpoint's timestamp. It also needs step k+1's first micro-batch packed in time; if rollouts are late, the update is still deferred but nothing is staged.

What to check in the first run

What	Where	Healthy	Problem
The image has this PR	the run starts	the orchestrator accepts --pipeline_step_boundary, the trainer the tuned keys	an older image stops at argument or config parsing
The log-prob head engaged	trainer log, warm-up compile	two [logps-head] sharded head local: ... 1x16384 local tokens in 16 steps of 1024 lines, as on the bench	[logps-head] sharded head off, using compute_chunked_logps: <reason>, or no [logps-head] line (head code missing, chunk size 0 or packing off)
Staging engaged	orchestrator log	at each eligible boundary, [pipeline] STAGED next step's first microbatch on the trainer before the update in X s; the last microbatch then still computed for Y s	[pipeline] stage_next_microbatch is set but the trainer cannot stage ... (no AI-Hypercomputer/maxtext#5718 in the image) or [pipeline] The trainer cannot defer a step's update .... A Y of about 2.7-2.9 s means the chips were already idle
Router replay reaches the trainer	trainer log, MOE_DROP_AUDIT_STEP step=N dropped_assignments_used=U dropless_replays=R	dropless_replays near 0	a retry on every micro-batch, or a recompile after warm-up (a third [logps-head] line): check that the rollouts return routing. On the random-weight bench, missing routing made every micro-batch retry
No token was dropped	the same lines	dropped_assignments_used=0	any other value

Gates before un-drafting

This PR stays a draft until all three checks have run on TPU and passed.

1. Change 1 against the team recipe. A same-pod real-gradient A/B of the libtpu flag and the six MaxText settings (the first commit's seven keys minus moe_route, now left out) against the team recipe, with pass bounds written down before the run.
2. The staged step boundary with real gradients, at 128 chips. Staging on against two staging-off runs in the same pod, on the bench through the bench-driver switch ([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) AI-Hypercomputer/maxtext#5719). Pass if loss, grad_norm, entropy and the forward log-prob metric at steps 0, 1 and 2 stay inside the band between the two staging-off runs. This PR's orchestrator path is exercised on TPU only by gate 3.
3. One end-to-end smoke run from an image built at the pin: checkpointing on, real rollouts, and if possible a forced overflow and dropless retry at 128 chips.

What no run has covered yet:

• The pinned MaxText: the bench ran its own image's MaxText (305 commits the pin does not have) with 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) AI-Hypercomputer/maxtext#5718's methods overlaid, on the Oct-9 Pathways build. On the recipe's Pathways build (raiden_20261013212200) only the head A/B was re-measured (120.34 -> 106.60 s, same pod). It ran MaxText 60425bfbd, the pin without 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) AI-Hypercomputer/maxtext#5718's engine-only commits (the same model, kernel and config code), within the 84,992 MiB budget. Not measured at the pin: the 2 precision keys, the staged boundary, GBS 1,024 and checkpointing.
• The orchestrator path on TPU ran only in the reduced rig. No staged TPU run has retried a micro-batch; a retry inside update() after a stage is CPU-tested only.
• Router replay on this rollout stack: no run has checked that vLLM, tpu-inference and the MaxText vLLM adapter at the pin return routed experts on TPU. The router-replay row of "What to check in the first run" is the first place to see it.
• The fully dropless layout fit the 84,992 MiB HBM budget before the tuned keys and the head, and has not been compiled with them.

Commits in this PR

Commit	What it does	How it was checked
recipes(397b v7x): trainer flags measured on the RL trainer bench	TRAINER_TUNED_FLAGS: the libtpu flag and seven MaxText keys (the six settings of change 1 plus moe_route=device, which a later commit removes)	Bench: 130.8 -> 126.4 s with the flag (another pod), 121.7 s with the keys (same pod)
recipes(397b v7x): two precision keys join TRAINER_TUNED_FLAGS	gdn_cp_matmul_precision=high, cast_grads_after_all_reduce=true	107.59 -> 106.24 s, same pod; real-gradient A/B PASS
Merge atwigg/mlperf 4c7823f20	Base sync; brings in #2633	Adds exactly the base's change
rl: opt-in sharded log-prob head for packed rows (TUNIX_LOGPS_SHARDED_HEAD)	#2873's head, squashed, plus sequence axes in mesh order and three tests	122.00 -> 107.59 s, same pod; real-gradient A/B PASS (on #2873's file a29991d3f)
orchestrator: opt-in staged step boundary (PIPELINE_STEP_BOUNDARY, STAGE_NEXT_MICROBATCH)	#2877's change, squashed	105.97 -> 98.31 s, same pod, via the bench-driver switch; rig idle 12.46 -> 0.02 s
recipes(gbs1024 submission): turn on the sharded log-prob head and the staged step boundary	Head and staged boundary on by default, in this recipe only	Render tests, incl. off switches and other recipes
recipes(gbs1024 submission): default to the trainer cell measured on the RL trainer bench	The default trainer is the bench cell	Render tests for each default and override
requirements: pin MaxText to atwigg/mlperf plus AI-Hypercomputer/maxtext#5718	Both MaxText lines at c79bcfd16: 60425bfbd plus six commits	git: the six touch only MaxText's training engine, README and tests
tests(orchestrator): describe measured costs in plain words; drop diagnostic prints	Plain-word test names, comments and docstrings; 19 diagnostic prints removed; comment and docstring edits in rl_program.py and trainer_worker.py (the no-overlap latch cannot fire on Pathways); no behavior change	549 of 549 before and after
recipes(397b v7x): leave moe_route out of TRAINER_TUNED_FLAGS	moe_route stays at its default, as the bench effectively ran	Launcher 35 of 35; the new test fails on the previous recipe
orchestrator: start a step's clock before waiting for its groups	With --pipeline_step_boundary, a step that waits for its groups starts its clock before that wait, as without the flag, so step_time and the MLPerf train_step_time built from it include the wait. No wall time or scored timestamp changes. Found in review	New test, boundary off and on; per the commit message it fails on the previous code when the boundary is on and a step waits

A dry-run render at the moe_route commit (no later commit changes a recipe) gives a trainer whose MaxText settings equal the 98.31 s bench arm's on all 75 shared keys, with all 17 libtpu flags equal. Only the bench has enable_checkpointing=false and the inert moe_route=device. The batch differs (GBS 1,024 here with BATCH_SIZE=64, learning rate 2e-6, clip 0.0625; GBS 1,152 on the bench), and so do the Pathways images and the MaxText build.

Superseded PRs (do not merge)

Each is marked DO NOT MERGE and opens with a note that points here.

PR	What it was	Status / why kept
#2873	The sharded log-prob head	Its review history, and the exact file measured on TPU (a29991d3f)
#2877	The staged step boundary	Its review history and the reduced-rig write-up
#2876	Lets the old head compile at log-prob chunks above 1024 under context parallelism	Reference only; with the sharded head, chunk 4096 gave no gain (107.84 vs 107.59 s)
#2887	Integration build of every tunix change of the 128-chip run	Superseded by this PR
AI-Hypercomputer/maxtext#5719	The bench-driver switch used to measure the staged step boundary	Bench only
AI-Hypercomputer/maxtext#5723	MaxText integration build	Superseded by AI-Hypercomputer/maxtext#5718, which merges first

Known limitations and follow-ups

Known limitations

• Staging cannot turn itself off on Pathways. It should stop when a stage did not hide (device wait under 0.05 s), but an idle device's wait reads 2.7-2.9 s there. If a stage ever outlasts the last micro-batch, the update waits with no warning.
• Padding exclusion (moe_exclude_padding_from_routing=true) is now on by default in this recipe (the 256-chip recipe keeps it off). It has CPU parity only (MoE: exclude padding tokens from routing (moe_exclude_padding_from_routing) AI-Hypercomputer/maxtext#5672); its TPU on/off numerics A/B has not run. Every bench row, the team recipe included, ran with it on, so no measurement here isolates it.
• Head scope. It engages only under jax.set_mesh, which MaxText's training engine uses; tunix's own RL stack (with mesh:) falls back with a warning, and Explicit-axis meshes raise. The switch is read at trace time.
• Memory. The head keeps a full-size f32 table-gradient accumulator per device (4.07 GB = 3.79 GiB) and the gathered table from forward to backward; for the team recipe (AOT) the main program fell by 6.6 GiB and the dropless retry program rose by 3.84 GiB, and both fit. Staging holds about 20 MB per device of inputs on the trainer and 5.11 GB more orchestrator host memory (measured on an earlier 128-chip test rig, [DO NOT MERGE] perf(orchestrator): opt-in pipelined step boundary that stages the next micro-batch before the deferred update #2877).
• Precision keys. Measured only with the head on; the split of the -1.36 s between the two keys is not measured.
• Splash options. sa_fuse_reciprocal=true and sa_use_base2_exp=false reverse the values [v7x] 397B 256 recipe: restore v7x trainer config after mlperf_base move #2435 set in this recipe (fix(recipes): add TPU v7x XLA/Pathways flags and Qwen3.5 GDN/Splash settings to mlperf_35b_128_v7x.sh #2509 lists sa_fuse_reciprocal=false among Qwen3.5 stability flags for the 35B recipe); gate 1 covers this.
• REMAT_POLICY=full with the switch on fails MaxText's config validation at startup (gdn_cp_state=device needs custom remat).
• Batch size. The bench ran GBS 1,152, the recipe runs 1,024; team logs show the packer still makes 9 micro-batches of 64 rows at GBS 1,024, the last one 53-95% empty (not an A/B).
• End-to-end gain of the boundary is small while rollouts dominate a step: [DO NOT MERGE] perf(orchestrator): opt-in pipelined step boundary that stages the next micro-batch before the deferred update #2877 estimated about 0-41 s per 26-minute team run (not measured).

Follow-ups (not blocking review)

1. Make a missing 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) AI-Hypercomputer/maxtext#5718 fail loudly at bring-up when STAGE_NEXT_MICROBATCH is set; today the orchestrator warns once, at the first step boundary, and runs unstaged.
2. The tuned keys replace the same keys in a caller's MAXTEXT_EXTRA_FLAGS, while an explicit NUM_MOE_TOKEN_CHUNKS or libtpu flag is kept. Pick one rule and pin it with a test.
3. LOG_MOE_DROPS=true is only a default. The drop audit is the run's record that no tokens were dropped (moe_dropped_assignments is written only with it), so a scored run should require it.
4. RAGGED_BUFFER_FACTOR is compared to -1 as a string, so -1.0 counts as a finite factor and MaxText refuses the config at trainer start: loud, but after bring-up instead of at render time.
5. Two test gaps: a recipe that changes which layouts get the 1,024-token log-prob chunk, and one that adds the padding mask and drop audit under TRAINER_TUNED_FLAGS=false with a retry override, both pass the launcher tests.

Repro on the bench

Follow "How to reproduce" in #2873 (bench pin, images, DRY_RUN, render, run, read), with these changes:

1. DRY_RUN this branch with the bench images and BATCH_SIZE=72 (GBS 1,152). It renders the final trainer flags (minus the inert moe_route=device), so render the bench without REPLACE_MAXTEXT_FLAGS, EXTRA_MAXTEXT_FLAGS and LIBTPU_EXTRA; REPLACE_MAXTEXT_FLAGS refuses the four keys the render sets twice (context_parallel_strategy, sa_fuse_reciprocal, sa_use_base2_exp, gdn_cp_mode). Checked against the measured arm offline, not re-run on TPU.
2. Client image: the bench image with this PR's tunix/rl/common.py over its own: use this PR's file (md5 0ca8d1ad11b769de5c18414d7874d9d0) in place of the guide's a29991d3f download and md5 check. Add 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) AI-Hypercomputer/maxtext#5718's maxtext_engine.py and inflight_throttler.py changes. The driver needs the STAGE_BEFORE_UPDATE switch from [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) AI-Hypercomputer/maxtext#5719, which raises without those methods and refuses a tunix that runs the update inside the last fwd_bwd (perf(orchestrator): fuse final microbatch fwd_bwd+update RPC and defer get_metrics off weight-sync path #2719, in this branch's base).
3. Two arms, one pod: STEPS=5 MICRO_STEPS=9, DRIVER_ENV="TUNIX_LOGPS_SHARDED_HEAD=local STAGE_BEFORE_UPDATE=on", then the same with STAGE_BEFORE_UPDATE=off, back to back.
4. Read steps 1-4 of the completed step lines and the STAGE_BEFORE_UPDATE step=k ... lines. Expect about 105.97 s off and 98.31 s on, two [logps-head] sharded head local lines, and 0 MoE drops.

Images of the measured runs:

• client gcr.io/tpu-vm-gke-testing/sivaibhav/trellis-experimental@sha256:b4647c188b25638e644c7a6fc53fb8e0751741525a277a1bcf07dca92ba9a6ac
• Pathways server us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/gke/datenglin/unsanitized_server@sha256:662f33b0ca6cd236a12534e4c78a8b7d715868669366738a015401925a39cb86
• Pathways proxy us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/gke/datenglin/unsanitized_proxy_server@sha256:e714351cdde50d306dd62f127cffabcdaf98dff80ad22e8b4c66c5a1d1df6ddb

The measured runs installed #2873's head file (a29991d3f, md5 a5bfda56676e92d65261f3e24303101b) and an overlay of AI-Hypercomputer/maxtext#5718's engine methods at pod start, with md5 checks. This PR's head file differs from a29991d3f in docstrings and comments, a clearer error for an unrecognized switch value, one moved pure-Python statement, and the mesh-order line.

Evidence: per-step numbers, numerics, the staged boundary on TPU

Per-step numbers, s per optimizer step. Pods are lettered to show which arms ran in the same pod.

Arm	Pod	Steps 1-4	Mean
Team recipe	A	130.872 / 130.807 / 130.720 / 130.815	130.80
Team recipe	B	130.785 / 130.754 / 130.812 / 130.869	130.81
Team recipe	C	130.723 / 130.719 / 130.766 / 130.910	130.78
+ libtpu flag only	D	126.434 / 126.407 / 126.360 / 126.253	126.36
+ 6 MaxText settings (and the flag)	D	121.759 / 121.773 / 121.776 / 121.634	121.74
same, repeat	E	121.949 / 122.035 / 122.042 / 121.980	122.00
+ sharded log-prob head	E	107.595 / 107.640 / 107.415 / 107.723	107.59
+ 2 precision keys	E	106.329 / 106.310 / 106.126 / 106.185	106.24
same program, reference for the boundary	F	106.001 / 105.944 / 105.993 / 105.950	105.97
+ staged step boundary	F	98.297 / 98.303 / 98.327 / 98.319	98.31
+ staged step boundary, repeat	F	98.271 / 98.271 / 98.262 / 98.236	98.26

• The libtpu flag alone was -4.43 s against the team recipe (cross-pod) and the MaxText settings -4.63 s on top (same pod D); the tables use pod E's repeat, which shares a pod with the head arms. Every arm from the 6 MaxText settings on also carried the inert moe_route=device.
• 98.31 sits inside the band written down before the run (98.48, range 97.5-99.8).
• Regular micro-batch, median, down the five stages: 13.51, 12.53, 10.94, 10.78 and 10.75 s.
• 256 chips, same pod: boundary off 71.683 / 71.465 / 71.374 / 71.372, on 56.127 / 56.116 / 56.122 / 56.095. Predicted -11.7 s (range -12.4 to -7.6), measured -15.36: a miss on the good side, as the stage hid completely.

Sharded log-prob head, real gradients. One pod, arms head off, head on, head off again, same random init and input. Step 0 is gated (bounds and results in Correctness; kept fraction 1 in every arm, loss bit-identical at 0.0304453493); steps 1-2 follow updates and are reported only. grad_norm is higher with the head at all three steps: +2.78e-4, +2.01e-4 and +2.89e-4 relative, against off-vs-off spreads of 5.32e-5, 2.65e-5 and 2.01e-5. That fits the f32 summation; the TPU run logs only the global norm, so the attribution rests on CPU. On CPU (real GRPO loss and MaxText head, 8 devices, 4 rows x 65,536 tokens), 0 of 262,144 tokens differ in log-prob or entropy, the loss, dL/dh and 68 other gradient leaves are exactly equal, and 5 deliberately broken heads are caught.

Staged boundary, bench, final 128-chip pair (one pod; arms in the order on, off, on again):

• Every boundary logged STAGE_BEFORE_UPDATE step=k staged=... stage_s=0.007-0.041 device_wait_s=10.49-10.72: the stage call returned at once, and the update then waited about one micro-batch for the device.
• Forward log-prob metric 12.2219073 off, 12.2219076 on (within-arm range about 2e-5). 0 drops, 0 retries.
• Exact checks ran on an earlier same-pod pair without the 2 precision keys (107.73 -> 100.01 s, same bench switch) and all passed: same micro-batches in the same order (per-micro-batch input fingerprint), one update per step, staging lines and call counts as expected. The call sequence, including apply_optimizer, is pinned by the bench switch's CPU tests.

Reduced rig (defined under "Change 4"): an earlier revision of this PR's StandardRLProgram and TrainerWorker, with timing instrumentation, drives the trainer with the team recipe's settings and dropless retry. Input: 16 recorded trajectory groups from the first optimizer step of a team 397B RL run, each served twice. All arms in one pod, 5 steps each, medians:

Arm	Flags	Boundary device idle, s	Step time, s
A0	off	12.46	64.45
A0p	off, repeat	11.72	61.85
A1	--pipeline_step_boundary	11.34	62.26
A2	+ --stage_next_microbatch --no-stage_before_update	7.43	59.23
A3	+ --stage_next_microbatch (default order)	0.02	50.82

• In every arm, each step trained the same micro-batches in the same order (content digests), with the same summed loss denominator (2,040,178), and every micro-batch of step k+1 entered after update k. In A3 every stage reached the trainer before the update. After warm-up, no arm ran a backend compile of 5 s or more.
• Written down before the run (predicted value, range): A0 idle 12.0 s (9-15), A3 0.016 s (0-3), step time A0 64.6 s (59.6-70.6), A3 52.6 s (48.1-57.6). All hit. A stricter estimator puts A3's idle at 0.026 s.
• The rig's scorer printed an overall FAIL with 74 of 79 checks passed. No ordering, data or decision check missed: the 2 misses are grad_norm bands of two non-default arms, and 3 loss checks could not be scored because at learning rate 0 the flags-off loss does not repeat within 2% from step to step.

Corrections to commit messages

Pushed history is not rewritten, so the corrections are recorded here. The numbers in the messages are right; some sentences claim more than was measured, or have gone stale.

• recipes(397b v7x): trainer flags measured on the RL trainer bench lists moe_route=device, which the moe_route commit removes (it did nothing on the bench). It names the source runs in Speed up Qwen3.5-397B-A17B training on TPU v7x AI-Hypercomputer/maxtext#5409 by number: read that as settings of the MaxText pretraining recipe for this model.
• recipes(397b v7x): two precision keys join TRAINER_TUNED_FLAGS: since the moe_route commit the block holds six keys plus these two, and the exact-set test expects eight; the pin is now c79bcfd16 (the keys mean the same there).
• rl: opt-in sharded log-prob head ...: each head parameter is gathered once, in the forward, and the backward reuses it (not "once per pass"). There are four fallbacks, not three. The old head's 129 and 65 collectives each include one for the input embedding, and XLA compiles the "reduce-scattered once" step as an all-reduce plus a slice.
• orchestrator: opt-in staged step boundary ...: the reduced rig fed recorded trajectory groups from a real 397B RL run's first optimizer step, not synthetic ones, and ran an earlier revision (26b06c504, before the base merge that brought [mlperf] Add intra-prompt rollout over-generation, straggler cancellation, and 35B recipe fixes #2633). Its boundary code is the same apart from comments, docstrings, one help string and the keyword order in one constructor, but the rig never ran it together with [mlperf] Add intra-prompt rollout over-generation, straggler cancellation, and 35B recipe fixes #2633's changes. The "~7.7 s per step" is the bench's gap; the orchestrator's own idle in the rig was 12.46 s, cut to 0.02-0.03 s. "When a stage measures no overlap, staging latches off" does not hold on Pathways (see Known limitations).
• recipes(gbs1024 submission): turn on the sharded log-prob head ...: the 98.31 s run used the same head switch, but its boundary ran through the bench driver's switch ([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) AI-Hypercomputer/maxtext#5719), not PIPELINE_STEP_BOUNDARY; the bench has no orchestrator. At that commit the default trainer was still fully dropless with log-prob chunk 512; the next commit changed that.
• recipes(gbs1024 submission): default to the trainer cell ...: "the same MaxText and libtpu flags as the bench's fastest 128-chip run" held at that commit; since the moe_route commit they match on the 75 shared keys (see Commits in this PR).

[gemini-code-assist]

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces significant performance optimizations for the Qwen3.5-397B-A17B RL trainer on TPU v7x. By tuning MaxText settings, implementing a sharded log-prob head, and optimizing the step boundary pipeline, the trainer achieves a substantial reduction in optimizer step time. These changes are designed to improve throughput and efficiency for large-scale RL training workloads, with performance validated on 128 TPU v7x chips.

Highlights

• Performance Optimization: Reduced optimizer step time by approximately 25% (130.8s to 98.3s) through tuned MaxText settings, precision adjustments, and staged step boundaries.
• Sharded Log-Prob Head: Implemented a sharded log-prob head to reduce memory overhead and compute time by gathering vocab tables once per micro-batch.
• Staged Step Boundary: Introduced staged step boundaries to overlap data copying with computation, effectively reducing device idle time.
• Recipe & Dependency Updates: Updated trainer defaults and pinned MaxText to atwigg/mlperf to ensure stability and support for the new performance features.

New Features

🧠 You can now enable Memory (public preview) to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces pipelining for the step boundary and support for staging the next step's first micro-batch on the trainer, alongside a sharded chunked log-prob head (TUNIX_LOGPS_SHARDED_HEAD) to optimize distributed post-training. The review feedback highlights several critical and medium-severity issues: a potential AttributeError from using jax.shard_map instead of jax.experimental.shard_map, multiple violations of the 'One well-lit path' style guide rule due to speculative attribute access and duck-typing on structured schemas, missing or overly broad type annotations (list[Any]), and type-checking bypasses that should be resolved by extending the AbstractRLEngine interface.