[parambole]

Do not merge. This change now ships in #2878, together with every other tunix change of the 128-chip best run (MaxText side: AI-Hypercomputer/maxtext#5718). This PR stays open for its review history and as the pin of the file measured on TPU (a29991d).

TL;DR

• Why: today's GRPO log-prob head re-gathers the 2 GB vocab table for every 1,024-token chunk, all-reduces its bf16 gradient per backward chunk, and runs every chunk on all 4 context-parallel ranks. A TPU profile puts it at 1,871.5 ms per micro-batch per TensorCore, against 199.0 ms for the pretraining loss head of the same model.
• What it does: adds an opt-in shard_map version of that head in compute_per_token_logps, for packed rows (segment_ids set). It gathers the table once per micro-batch and keeps it for the backward, each device scans only its own tokens, and the table gradient is summed in f32 and reduced once. Off by default: set TUNIX_LOGPS_SHARDED_HEAD=local in the process that traces the trainer step (on Pathways, the client / trainer driver); with it off, the traced program is the same as today's (jaxpr hashed and compared on CPU).
• Result: on the MLPerf Qwen3.5-397B-A17B trainer bench (Add an RL trainer bench under experimental/rl_trainer_bench AI-Hypercomputer/maxtext#5701, one 128-chip v7x slice), head off then on in one pod: the bench step (one optimizer step: about 1,152 trajectories packed into 9 micro-batches) drops from 122.00 to 107.59 s (-11.8%) and a regular micro-batch (8 of the 9) from 12.527 to 10.937 s (-12.7%). On newer images (pod 10: the Pathways tag this PR's base moves the recipe to, a newer client image and MaxText build): 120.34 to 106.60 s (-11.4%).
• Numerics: on CPU, 0 of 262,144 tokens differ from today's head in log-prob or entropy. On TPU grad_norm is 2-3e-4 relative higher, consistent with summing the table gradient in f32 once instead of in bf16 per chunk.
• Main limits: it engages only under jax.set_mesh, which MaxText's training engine uses; tunix's own RL stack uses with mesh: and falls back to today's head. It has not run end to end with real weights, and it needs a full-size f32 gradient accumulator per device (see Risks).

How it works

flowchart TD
  D["compute_per_token_logps: packed rows, chunk_size above 0"]
  subgraph TODAY["Today, or switch off: compute_chunked_logps under GSPMD"]
    A1["1,024-token chunks, each on all 4 context-parallel ranks"]
    A2["each chunk: all-gather the 2 GB table (64 fwd + 64 bwd)"]
    A3["head matmuls: 533.3 TFLOP per device per micro-batch"]
    A4["each backward chunk: all-reduce the bf16 table gradient (64)"]
  end
  subgraph LOCAL["TUNIX_LOGPS_SHARDED_HEAD=local: shard_map with a custom VJP"]
    B2["one-token shifts: ppermute between neighbouring shards"]
    B1["all-gather the head parameters once, kept for the backward"]
    B3["scan only local tokens: 1x16384 in 16 steps of 1024"]
    B4["same per-chunk math: 133.3 TFLOP per device per micro-batch"]
    B5["backward: per-device f32 gradient accumulator (4.07 GB)"]
    B6["one psum_scatter of the gradient"]
  end
  D -->|switch off| A1
  D -->|switch on| B2
  D -.->|switch on but a checked precondition fails| A1
  A1 --> A2 --> A3 --> A4
  B2 --> B1 --> B3 --> B4 --> B5 --> B6

Loading

If a checked precondition fails, the helper logs a [logps-head] sharded head off warning and today's path runs; some inputs raise instead, such as explicit-axis meshes (limitations 2 and 9). When it engages, it logs [logps-head] sharded head local: ....

Results

measurement	where	head off	head on	change
bench step (9 micro-batches)	TPU, pod 5	122.00 s	107.59 s	-14.41 s (-11.8%)
regular micro-batch	TPU, pod 5	12.527 s	10.937 s	-1.590 s (-12.7%)
bench step, newer images	TPU, pod 10	120.34 s	106.60 s	-13.74 s (-11.4%)
grad_norm at step 0, real gradients	TPU, pod 7	0.0557634532 and 0.0557664186	0.0557804443	+2.78e-4 relative (off vs off: 5.32e-5)
per-token log-probs and entropy vs today, 4 rows x 65,536 tokens	CPU, 8 devices			0 of 262,144 differ
table all-gathers / all-reduces per micro-batch; head matmuls per device	AOT compile, baseline cell	129 / 65 (head 128 / 64, input embedding 1 / 1); 533.3 TFLOP	2 / 2; 133.3 TFLOP	4.00x fewer head FLOPs

• TPU cell: arm B, the tuned cell (the baseline cell + the 7 MaxText keys and 1 libtpu flag of 397B v7x RL trainer: tuned settings, sharded log-prob head, staged step boundary (130.8 -> 98.3 s per optimizer step, trainer bench) #2878), on one 128-chip v7x slice (256 devices, 2 per chip): trainer mesh fsdp 32 x context 4 x expert 2, 64 packed rows x 65,536 tokens per micro-batch, log-prob chunk 1024, random weights, 0 MoE drops. The baseline cell is the trainer of the measured MLPerf submission-recipe launch (arm A in "How to reproduce").
• Terms: a pod is one JobSet run whose arms ran back to back on one slice (in pods 5 and 10 only the switch differs). The regular micro-batch time is the median of dispatch entries 3-9 over steps 1-4 (step 7 of "How to reproduce").
• Speed mode and real gradients: in every speed arm, including pods 5 and 10, the log-prob gates drop every sequence of the random-weight model (kept_frac 0), so loss and grad_norm are 0. Pod 7 keeps every sequence with seeded random advantages (real gradients) and measured -14.23 s/step (-11.7%) over steps 1-2.

Changes

• tunix/rl/common.py
  • compute_per_token_logps: the only edit to today's path. Packed, chunked calls (segment_ids set, chunk_size > 0) first try maybe_compute_sharded_chunked_logps; on None they run compute_chunked_logps as before. Today's path pads packed log-probs and entropy with one zero at the front; the sharded head's output is already aligned, so that pad is skipped.
  • sharded_head_mode() (new): parses the switch. local, 1, true, on, yes turn it on; unset, empty, 0, false, off, no leave it off; anything else raises ValueError.
  • _token_partition() (new): mesh axes for the batch and sequence dims, from the flax logical-axis rules (activation_batch, activation_length) when valid, else from the mesh order.
  • _head_parameters(), _use_dtypes() (new): trace compute_final_logits once, abstractly, to find the parameters it reads and the dtype to gather each in.
  • maybe_compute_sharded_chunked_logps() (new): checks the preconditions (else warns and returns None), then runs a shard_map with head_local, a custom VJP (forward: gather once, scan local tokens; backward: f32 accumulation, one psum_scatter), and the neighbour ppermute one-token shifts. One long function; its per-chunk math duplicates compute_chunked_logps's.
• tests/rl/sharded_logps_head_test.py (new): 19 tests on an 8-device CPU mesh: parity with compute_chunked_logps, engagement, fallbacks and switch parsing.

Testing

• Unit tests (CPU, run by hand): the new module 19/19, tests/rl/common_test.py 83/83. Broken copies of the head (input or output boundary token dropped, accumulator overwrite, logits x1.001) each fail 5-7 tests.
• Numerics gate (CPU, with positive controls): the real algo_core.grpo_loss_fn and MaxText head on 8 CPU devices. At 4 rows x 65,536 tokens, 0 of 262,144 tokens differ; the loss, dL/dh (the gradient for the hidden states) and 68 other gradient leaves are exactly equal, and only the table and norm-scale gradients differ, both closer to an all-f32 reference; every deliberately broken head fails.
• Default off: with the switch off, the jaxpr of value_and_grad through compute_per_token_logps hashes identically to unmodified tunix.
• TPU compiles on a CPU host: AOT of the baseline cell for tpu7x-256; a head-only compile with libtpu 0.0.51.dev20261010 at chunk 1024 and 4096.
• TPU, one 128-chip v7x slice: same-pod A/Bs in pods 5 and 10, a cross-pod pair (pods 3 and 4), and a real-gradient A/B in pod 7 (off, on, off again). On the pinned images of "How to reproduce" (every pod except 10), the head-on program ran 4 times in speed mode, in pods 4, 5 and 6, spanning 0.45 s/step.
• Not run: CI (it runs only on PRs into main), a single-process pytest tests/rl/ with the new module, pylint (pyink is clean on the changed lines), any end-to-end RL run.

Risks and limitations

• Scope: it engages only under jax.set_mesh. tunix's own RL stack (with mesh:) falls back to today's head with a warning, so it gets no speedup yet (limitation 1).
• Bench only: every TPU number for the GRPO head comes from the trainer bench (random weights, synthetic micro-batches) and one model; the A/Bs are on one 128-chip mesh (it also ran, speed only, on a 256-chip slice at context 4 and context 8; "Not tested").
• Memory: each device holds a full-size f32 gradient accumulator (4.07 GB = 3.79 GiB for the table) and keeps the gathered table from forward to backward; neither shrinks as chips are added. At the baseline cell (AOT) the main program's memory fell by 6.6 GiB, and its dropless retry variant (which reruns a micro-batch without dropping tokens when the MoE buffer overflows) rose by 3.84 GiB; both fit (limitation 3).
• Robustness gaps (limitations 2 and 4 to 9): none gives a wrong result on the measured path. Off it, two can return a wrong result with no error: a contrived head that fools the dtype detection (a test head was off by up to 1.19 nats) and a direct call with hidden states longer than the tokens. The other five: explicit-axis meshes raise; custom axis rules listing two sequence axes out of mesh order fail to compile; a misspelled switch is rejected only on packed, chunked calls; the switch is read at trace time; the test module sets the device count process-wide.

Related

• 397B v7x RL trainer: tuned settings, sharded log-prob head, staged step boundary (130.8 -> 98.3 s per optimizer step, trainer bench) #2878 turns on, in the MLPerf submission recipe, the 7 MaxText keys and 1 libtpu flag that both TPU arms ran with (so head off vs head on isolates this PR), plus the two GDN-CP keys of "Other MaxText levers".
• [DO NOT MERGE] Compile compute_chunked_logps at chunk_size > 1024 under context parallelism #2876: makes today's head compile at compute_logps_chunk_size above 1024 under context parallelism.
• 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 and [DO NOT MERGE] perf(orchestrator): opt-in pipelined step boundary that stages the next micro-batch before the deferred update #2877: opt-in flags that shorten the trainer's idle time between steps; separate from this PR.
• Add an RL trainer bench under experimental/rl_trainer_bench AI-Hypercomputer/maxtext#5701: the RL trainer bench behind every TPU number for the GRPO head.

TPU profile of today's head

Compared op by op with a MaxText pretraining run of the same model at 3,193 tok/s/chip (its own loss head, num_vocab_tiling=2), a TPU profile puts the GRPO head at +1.67 s per micro-batch per TensorCore. The profiled cell is the baseline cell plus two SparseCore all-gather/reduce-scatter latency-multiplier flags and the TensorCore ragged-sort MoE keys (unprofiled median 13.21 s per micro-batch, against 13.505-13.522 s for the baseline cell).

piece	GRPO head today	pretraining loss head
table all-gather	545.6 ms: 128 calls (64 fwd + 64 bwd), 4.27 ms each	26.0 ms
table-gradient all-reduce	660.8 ms: 64 calls, 10.3 ms each (bf16, 64-way over fsdp x expert)	8.1 ms
vocab matmuls	507.7 ms: 256 calls, 533.6 TFLOP (4.0x)	135.6 ms: 133.3 TFLOP
norm, log-softmax, entropy, glue	110.2 ms	19.4 ms
other head ops	47.2 ms	9.9 ms
head total	1,871.5 ms	199.0 ms

A paper sum from this profile (not a model) estimated -1.3 to -1.6 s per micro-batch; the same-pod A/B, run on arm B, measured -1.590 s.

Design details: gather dtype, chunking, fallbacks, explicit shifts

Each chunk runs today's math: compute_final_logits, f32 logits, then selective_log_softmax and entropy.

• Gather dtype: a parameter used only after a cast to a narrower float (for example an fp32 table cast to bf16 for the dot) is gathered in that dtype. At bf16 weights the bench gathers the norm scale (4096,) and the table (4096, 248320) in bf16.
• Chunking: chunk_size becomes tokens per scan step per device, so per-device logits memory per step is unchanged. Chunks are dynamic_slices of the local block, not today's reshape + swapaxes.
• Fallback (the helper logs [logps-head] sharded head off, using compute_chunked_logps: <reason> and returns None, and today's path runs): no mesh in context (jax.set_mesh); no token layout, because some mesh axis of size > 1 shards neither batch nor sequence; a head parameter that does not divide over the mesh; or a compute_final_logits that reads no model state. Explicit-axis meshes raise instead (limitation 2). Unpacked rows and switch-off calls take today's path.
• Why explicit ppermute shifts: while testing this, the XLA SPMD partitioner (CPU, jax/jaxlib 0.11.1) miscompiled jnp.concatenate([x[:, 1:], 0]).at[:, k].set(0) on a context-sharded [B, T] array: the first token of every later context block kept its unshifted value (got 4097, want 4098). The shift alone, the set alone, and the shift then jnp.where are correct. It hit the 8-device fsdp2 x context2 x expert2 mesh (4 rows x 1 boundary = 4 wrong tokens) and an 8-way context-only mesh (4 rows x 7 boundaries = 28 wrong tokens), not a 2-device one. This PR does both shifts inside the shard_map, so its correctness never depends on how the partitioner shifts a sequence-sharded array.

TPU evidence: speed (pods 3, 4, 5, 6 and 10)

Setup: the TPU cell of "Results" (arm B), plus ragged_buffer_factor 2.75 and 9 micro-batches per bench step, which #5701 calibrates to one optimizer step of the recipe (about 1,152 trajectories packed into 64-row micro-batches). Client image trellis-experimental@sha256:b4647c18... (pod 10: newer images) with this PR's tunix/rl/common.py (md5 a5bfda56) installed over the image's copy. "Head off" arms ran this PR's file with TUNIX_LOGPS_SHARDED_HEAD=off, except the pod-3 reference (stock file). Step 7 of "How to reproduce" defines both tok/s/chip rulers.

Two more measurements of the head effect, besides pods 5 and 10 in "Results":

comparison	head off: s/step / micro-batch	head on: s/step / micro-batch	change: s/step	change: per micro-batch
cross-pod (pod 3 off, pod 4 on; no TPU host shared)	121.74 / 12.492	107.54 / 10.919	-14.19 (-11.7%)	-1.573 (-12.6%)
same pod (pod 7), real gradients, steps 1-2, off = mean of 2 arms	121.82 / 12.519	107.59 / 10.936	-14.23 (-11.7%)	-1.583 (-12.6%)

• Pod 5 ran six arms back to back: + decoder input on device, + chunk 4096, head off, head on, + GDN-CP keys (gdn_cp_matmul_precision=high, cast_grads_after_all_reduce=true), head on + profiler. Within-arm spread over steps 1-4: 0.09 s (off), 0.31 s (on). tok/s/chip rises 14.5% (micro-batch ruler) and 13.4% (step ruler). MoE dropped / replayed: 0 / 0. Head on logged two engaged [logps-head] lines and no fallback; head off logged none.
• Forward check, pod 5: sampler_is/token_logdiff_absmean (trainer per-token log-prob against the payload's rollout log-prob, mean over the 40 regular payloads of steps 0-4) reads 12.2219079 (off) and 12.2219074 (on): 5.1e-7 apart, against a within-arm range of about 2e-5. Speed arms are forward only (loss and grad_norm are 0, see "Results"); pod 7 checks gradients.
• Across pods: on the pinned images the head-on program ran 4 times in 3 pods in speed mode: 107.54 (pod 4), 107.59 and 107.48 (pod 5; the 107.48 run had the profiler on for one step-0 call), 107.94 (pod 6), a range of 0.45 s/step (0.42%); regular micro-batch 10.919 to 10.949 s. Head off: 121.74 (pod 3, stock file) and 122.00 (pod 5). The three pinned-image head effects (-14.19 to -14.41 s/step) are about 30x that 0.45 s range. Bands of ±1.0 s around pod 4 (head on, 107.54) and pod 3 (head off, 121.74), written down before pods 5 and 6 ran (106.54-108.54 and 120.74-122.74), held every head-on and head-off run in pods 5 and 6.
• Pod 10, newer images (order C, B): Pathways server and proxy tagged raiden_20261013212200 (unsanitized_server@sha256:7cd4dbed..., unsanitized_proxy_server@sha256:8679cb5c...), client gcr.io/cloud-tpu-multipod-dev/atwigg/trellis-experimental:1013b (sha256:f30c7b2f...). Same tunix/rl/common.py as b4647c18 (md5 be286d30), different MaxText build (854 resolved config keys instead of 934, no embed_router axis rule, gdn_pin_in_proj_output_layout=True), so only the A/B inside pod 10 is comparable. s/step, steps 1-4: off 120.337 / 120.319 / 120.307 / 120.393, on 106.585 / 106.610 / 106.588 / 106.633; regular micro-batch 12.357 to 10.813 s (-1.544 s, -12.5%). tok/s/chip (micro-batch / step ruler): 2,652 / 2,451 off, 3,031 / 2,766 on. 0 MoE drops, 0 replays; diff_config.py shows only the 8 path and name keys between the arms (keys, ragged_buffer_factor, chunk, mesh and batch as in arm B); head on logged two engaged lines, no fallback. Forward check: 12.2219074 (off), 12.2219072 (on), within-arm ranges 2.3e-5 and 1.7e-5. Here the 9th dispatch entry reads about two regular micro-batches (21.65 s on, 24.74 s off) and the time after the last dispatch shrinks by one (head on: 1.60 s, against 12.52 s in pod 4); the median of entries 3-9 still picks a regular entry.

TPU evidence: numerics with real gradients (pod 7)

Numerics mode: LOGP_GATES=off keeps every sequence (kept_frac 1); ADVANTAGES=random gives a seeded N(0, 1) advantage per segment, the same in every arm, with the same random init and payload. Three arms in one pod: head off, head on, head off again (the replicate). The step-0 pass rules were written before the run, and at step 0 the weights are identical; steps 1-2 follow one and two updates with each arm's own gradient, so they are report-only.

step	loss (all 3 arms)	grad_norm: off / on / off again	on vs mean(off), rel	off vs off, rel	entropy: on vs mean(off), rel
0	0.0304453493	0.0557634532 / 0.0557804443 / 0.0557664186	+2.78e-4	5.32e-5	-1.7e-8
1	0.0304453493	0.0557874702 / 0.0557994284 / 0.0557889491	+2.01e-4	2.65e-5	-1.8e-7
2	0.0304453493	0.0557963252 / 0.0558119081 / 0.0557952039	+2.89e-4	2.01e-5	-3.6e-7

step-0 gate (written before the run)	measured	PASS if	FAIL if	verdict
kept_frac	1 / 1 / 1	> 0 in every arm	0	PASS
loss, rel	0 (bit-identical)	<= 1e-5	> 1e-3	PASS
forward check token_logdiff_absmean, on minus mean(off)	3.7e-6 (off vs off: 3.8e-6)	<= 5e-5	> 5e-4	PASS
entropy, rel	1.7e-8	<= 1e-4	> 1e-3	PASS
grad_norm, rel	2.78e-4 (off vs off: 5.32e-5)	<= 2e-3	> 1e-2	PASS

• The bit-identical loss does not test the head: at an importance ratio of 1 the GRPO loss value depends only on advantages and masks. The gradient still flows through the head's log-probs.
• The forward matches within the replicate noise (log-prob check and entropy).
• grad_norm is higher with the head on at all three steps: 5.2x the off-vs-off spread at step 0, then 7.6x and 14.4x at steps 1-2, after the weights diverged. This is the head's one numerical change: the table gradient is summed in f32 once, not in bf16 over 64 all-reduces per micro-batch. On CPU only the table and norm-scale gradients differ from today's (relative Frobenius difference 1.63e-3 / 5.34e-3), both closer to an all-f32 reference ("CPU validation").
• Limits: the TPU run logs only the global norm, so attributing the difference to the table gradient rests on the CPU result, and no TPU run has shown this gate failing on a deliberately wrong head (on CPU every deliberately broken head fails the gate).

Other MaxText levers measured on top of the head (not part of this PR)

Each arm adds MaxText settings to head on ("GDN-CP keys": see the pod 5 bullet above). Changes are in s/step against the same pod's reference: head on at 107.59 (pod 5) or 107.94 (pod 6), head on + GDN-CP keys at 106.21 (pod 9). Differences under about 0.3 s/step are step-to-step noise.

• Pod 5: GDN-CP keys 106.24 (-1.36; real-gradient A/B on TPU: every gate passes, grad_norm within 1.5e-5 relative against a 1.9e-5 replicate spread); decoder_layer_input=device (cell: offload) 107.11 (-0.49); --compute_logps_chunk_size=4096 (cell: 1024) 107.84 (+0.25, no gain). Forward checks within noise (+1.2e-6, +1.0e-6, +1.6e-6). The chunk-4096 arm shows the sharded head compiles and runs at chunk 4096 on TPU, where today's head fails TPU compile ("CPU validation", item 4).
• Pods 6 and 9: prefuse off (MaxText and trainer), moe_accumulate_wi_dlhs=true, moe_accumulate_chunk_wgrad=true and decoder_layer_input=device: 105.72, replicate 105.85 (-2.22, -2.09, pod 6); without decoder_layer_input=device 106.69 (-1.24, pod 6); with GDN-CP keys 104.01, replicate 103.94 (-2.19, -2.26, pod 9). All three shift the forward check by +9.7e-4, and a real-gradient A/B on TPU measured grad_norm 0.80% lower (about 420x the replicate spread) with or without decoder_layer_input=device, so these settings are not used.
• Pod 9, weight_dtype=float32 (cell: bfloat16), with GDN-CP keys: 108.05 (+1.85); forward check +7.5e-4; fp32 confirmed by the compiled program's arguments (5.80 GiB per device, against 2.92 GiB for this pod's bf16 lever arm).

CPU validation: numerics gate, unit tests, compiles

1. Numerics gate through the real loss, with positive controls. Unmodified tunix algo_core.grpo_loss_fn and MaxText TunixMaxTextAdapter head (decoder norm + logits_dense); tunix's GRPOAdapter + SequencePackedBatchAssembler pack the micro-batch, with completion masks interleaving model and environment turns, an overlong rollout, and seq-mask-tis 0.999/1.002 keeping some sequences and dropping others. Submission-recipe GRPO config; fsdp2 x context2 x expert2 on 8 CPU devices, cp-as-ep rules, chunk 1024, 4 packed rows of T tokens; pass rules written before any treatment arm ran. A halo is the boundary token a sequence shard gets from its neighbour for the one-token shift.

check	T = 8,192	T = 65,536 (32 scan steps per device)
today vs today, two processes	174/174 arrays bit-identical	174/174 bit-identical
per-token log-probs and entropy vs today	0 tokens differ	0 of 262,144 differ
loss, every aux metric, dL/dh, 68 other gradient leaves	exactly equal	exactly equal
table / norm-scale gradient vs today (bf16 summation order), relative Frobenius	1.6e-3 / 5.2e-3 (bounds 3.9e-3 / 1.6e-2)	1.6e-3 / 5.3e-3
error vs an all-f32 reference, ratio to today's	table 0.92, scale 0.48	table 0.93, scale 0.13
positive controls (each must fail): input halo dropped, accumulator overwrite, logits x1.001, dL/dh offset	all fire	all fire, as does a fifth control (output halo dropped). Both halo controls change exactly 4 tokens, all at index 32,768 (the context boundary)
comparator self-test on deliberately wrong inputs		9/9

The gate ran on an earlier revision whose local and switch-off training jaxprs hash identically to this commit's (jaxpr of value_and_grad through compute_per_token_logps at the unit-test cell; a one-line mutant changes the hash). Switch-off also hashes identically to unmodified tunix.

2. Unit tests, 19/19. Parity with compute_chunked_logps: f32 and bf16; uneven and single-step chunks; log-probs only; the no-gradient forward with and without entropy; the mesh-order layout without logical rules; a sequence-only layout over 8 devices (7 shard boundaries). Also engagement, fallbacks and switch parsing. Broken heads (input halo dropped, output halo dropped, accumulator overwrite, logits x1.001) each fail 5-7 tests; a no-gradient-path bug fails exactly the 2 forward-only tests.

3. AOT at the baseline cell (arm A): a CPU-hosted compile for tpu7x-256 (libtpu 0.0.46; GRPO + router replay + fp8 fake-quant + ragged_buffer_factor 2.75 + the dropless retry program). No TPU run used the head at this cell (they used the arm B/C cell). It compiled an earlier revision that shifted with a global concatenate instead of the ppermutes; the head scan and table collectives are the same code. Between the columns MaxText differs by one commit (#5568, engine and metrics files only) and tunix, apart from this change, only outside rl/common.py and algo_core.py. The cell's xla_tpu_max_hbm_size_mib=84992 (83.0 GiB) caps temporaries + arguments + outputs (80.05 GiB today, 73.43 GiB with local); 94.74 GiB is the per-device HBM, the fit limit including state and code. Counts and FLOPs: see "Results".

	today	local
collective bytes per device, whole fwd_bwd	994 GB	609 GB
fwd_bwd temp / total (with arguments, outputs, state, code), of 94.74 GiB	71.35 / 86.62 GiB	64.73 / 80.01 GiB
dropless retry program temp	45.34 GiB	49.18 GiB (+3.84, about the 3.79 GiB f32 accumulator; fits)
s32 token-id reshape from the chunking transpose (the op that fails TPU compile at chunk 2048/4096)	present	absent

XLA rewrote the head's f32 reduce-scatter into one all-reduce plus a dynamic-slice. Instruction by instruction the head's FLOPs fall exactly 4.00x; the only other FLOP change is one extra rematerialization of a decoder projection the compiler chose (+30.8 TFLOP, about 27 ms at peak).

4. Head-only compile with the TPU-matched compiler: libtpu 0.0.51.dev20261010 reproduces the TPU compiler's pass and fail verdicts on this cell. The real compute_per_token_logps runs under value_and_grad with a stand-in head at the real dims and shardings (hidden [64, 65536, 4096] bf16, table [4096, 248320] bf16) on 256 tpu7x devices. Today's head passes at chunk 1024 and fails at 4096 with the TPU's exact signature (Reshape should have supported layout ... s32[4,1,4,1024] ... compute_per_token_logps/transpose); local passes at 1024 (0 token-id reshapes, 0 all-to-alls) and at 4096, as does the TPU run at 4096.

How to reproduce the TPU A/B (pins, commands, expected numbers)

The bench (AI-Hypercomputer/maxtext#5701) runs the exact trainer program of an e2e GRPO launch (same MaxText and libtpu flags, mesh, GRPO loss and chunked log-prob pass) on one trainer slice, with random weights and a synthetic micro-batch whose MoE routing is calibrated to a production run. Its input is the e2e launcher's DRY_RUN output. The A/B is three bench runs; this PR's effect is C vs B:

arm	trainer program	how the bench sets it
A, baseline cell	the trainer of the MLPerf launch in step 3, as rendered (ragged_buffer_factor 2.75, logps chunk 1024)	no knob
B, tuned cell	A + 7 MaxText keys + 1 libtpu flag (the settings of #2878)	REPLACE_MAXTEXT_FLAGS, EXTRA_MAXTEXT_FLAGS, LIBTPU_EXTRA
C, tuned cell + this PR	B + this PR's tunix/rl/common.py + TUNIX_LOGPS_SHARDED_HEAD=local	IMAGE (a derived image), DRIVER_ENV

Pinned inputs

what	pin
bench	AI-Hypercomputer/maxtext#5701 at its head 70d64441eba957fe5fe934cc4346a26df4706719 (branch sivaibhav/rl-trainer-bench of ultrons/maxtext), directory src/maxtext/experimental/rl_trainer_bench
e2e launcher (defines the trainer cell)	google/tunix ac001198872b0212351b4c58bc8932720487bf14 (atwigg/mlperf, merge of #2858): tunix/experimental/examples/recipes/mlperf_gbs1024_submission.sh
this PR	#2873 at a29991d3fef51325fe4a306e3590ca52fdad80b3. Its one runtime file is tunix/rl/common.py, md5 a5bfda56676e92d65261f3e24303101b, 77,339 bytes
client image (A, B, and the base of C)	gcr.io/tpu-vm-gke-testing/sivaibhav/trellis-experimental@sha256:b4647c188b25638e644c7a6fc53fb8e0751741525a277a1bcf07dca92ba9a6ac. tunix is an editable install at /app. The image's tunix/rl/common.py is md5 be286d30a3071d180676ef56b5c89782, byte-identical to the file in this PR's base (atwigg/mlperf 8d245c561)
Pathways resource manager and workers	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
colocated Python sidecar (checkpointing; pinned by the launcher)	gcr.io/cloud-tpu-multipod-dev/yixuannwang_google_com-runner@sha256:6682fc0f364d2246a8768265c8c06ef50e213d86e1639d367543c78a90a8f0f2
slice and mesh	one tpu7x:4x4x8 (128 chips = 256 devices, 32 hosts of 4 chips); trainer mesh fsdp 32 x context 4 x expert 2, custom_mesh_and_rule=cp-as-ep; micro-batch 64 packed rows x 65,536 tokens

Use the client image by digest: it carries a specific MaxText and tunix build.

Requirements

• GKE with JobSet and Kueue: a tpu7x 4x4x8 slice for the 32 Pathways worker pods and a CPU node pool for the Pathways head pod (recipe default cpu-np); pull access to the images above and a registry the cluster can pull arm C's image from.
• A GCS folder the JobSet's service account (SERVICE_ACCOUNT, default xpk-sa) can write: each arm saves a final checkpoint there after its timed steps, about 5 TiB (740 GiB of parameters + 4.3 TiB of optimizer state). Delete it afterwards.
• git, curl, the gcloud CLI, Docker, kubectl and Python 3 with PyYAML (the DRY_RUN needs only bash and python3); about 25-35 minutes of slice time per arm, plus scheduling. No Hugging Face access is needed: the image ships the Qwen/Qwen3.5-397B-A17B tokenizer in /opt/hf_cache and sets HF_HUB_OFFLINE=1.

Set these once (bash; in zsh, first run setopt interactivecomments, because the commands contain # comments):

export NS=my-namespace                                  # your Kubernetes namespace
export CTX=my-kube-context                              # your kubectl context for that cluster
export OUTPUT_PREFIX=gs://my-bucket/rl-trainer-bench     # each arm writes to $OUTPUT_PREFIX/<arm name>
export REGISTRY=us-docker.pkg.dev/my-project/my-repo     # where arm C's image is pushed

1. Check out the bench and the launcher at the pinned commits

mkdir -p tb-repro && cd tb-repro

# The bench: only its directory of PR #5701, at the PR head.
git init -q maxtext-5701
git -C maxtext-5701 sparse-checkout set src/maxtext/experimental/rl_trainer_bench
git -C maxtext-5701 fetch -q --depth=1 --filter=blob:none \
  https://github.com/AI-Hypercomputer/maxtext.git 70d64441eba957fe5fe934cc4346a26df4706719
git -C maxtext-5701 checkout -q FETCH_HEAD
export BENCH=$PWD/maxtext-5701/src/maxtext/experimental/rl_trainer_bench

# The e2e launcher whose DRY_RUN defines the trainer cell.
git init -q tunix-launcher
git -C tunix-launcher fetch -q --depth=1 https://github.com/google/tunix.git ac001198872b0212351b4c58bc8932720487bf14
git -C tunix-launcher checkout -q FETCH_HEAD

Check (md5 -q on macOS):

md5sum "$BENCH/render_bench.sh" "$BENCH/render_bench.py" "$BENCH/bench_driver.py"
# 96ddb23bb408daf781de5052b7265745  render_bench.sh
# 012ec6d3476bd2dbe28573196a1b70de  render_bench.py
# a2fdaf1ab1bfb7486e710585486c4d28  bench_driver.py

2. Build the client image for arm C

The bench image's tunix/rl/common.py is byte-identical to this PR's base, so copying the PR's file over it applies the PR. Only the client image (the main container, whose bench_driver.py traces the head) changes; IMAGE swaps it and leaves the Pathways containers alone.

gcloud auth configure-docker "gcr.io,${REGISTRY%%/*}"   # pull the base image, push to $REGISTRY
mkdir -p head-image
curl -fsSL -o head-image/common.py \
  https://raw.githubusercontent.com/google/tunix/a29991d3fef51325fe4a306e3590ca52fdad80b3/tunix/rl/common.py
cat > head-image/Dockerfile <<'EOF'
# The #5701 bench client image plus tunix/rl/common.py from google/tunix#2873 @ a29991d3f.
FROM gcr.io/tpu-vm-gke-testing/sivaibhav/trellis-experimental@sha256:b4647c188b25638e644c7a6fc53fb8e0751741525a277a1bcf07dca92ba9a6ac
# Refuse an unexpected base: the bench image's file is tunix 335c2abee's common.py.
RUN echo "be286d30a3071d180676ef56b5c89782  /app/tunix/rl/common.py" | md5sum -c -
COPY common.py /app/tunix/rl/common.py
RUN echo "a5bfda56676e92d65261f3e24303101b  /app/tunix/rl/common.py" | md5sum -c - \
 && rm -f /app/tunix/rl/__pycache__/common.*.pyc
EOF
docker build -t "$REGISTRY/trellis-experimental:logps-head-a29991d3f" head-image
docker push "$REGISTRY/trellis-experimental:logps-head-a29991d3f"
# The digest of the pushed image in $REGISTRY (not of any other copy of it).
export HEAD_IMAGE=$(docker inspect --format '{{range .RepoDigests}}{{println .}}{{end}}' \
  "$REGISTRY/trellis-experimental:logps-head-a29991d3f" | grep -m1 "^$REGISTRY/trellis-experimental@sha256:")
echo "${HEAD_IMAGE:?no digest in REGISTRY: did the push succeed?}"    # <registry>/trellis-experimental@sha256:...

Both md5sum -c lines print /app/tunix/rl/common.py: OK, and the build deletes the image's stale common.cpython-312.pyc. The base image is linux/amd64 only (elsewhere add --platform linux/amd64; the verified builds ran on x86-64 Linux). Optional checks:

docker run --rm -e JAX_PLATFORMS=cpu --entrypoint python "$HEAD_IMAGE" -c \
  "import tunix.rl.common as c; print(c.__file__, hasattr(c, 'sharded_head_mode'))"
# /app/tunix/rl/common.py True         (the stock bench image prints False)

curl -fsSL -o head-image/sharded_logps_head_test.py \
  https://raw.githubusercontent.com/google/tunix/a29991d3fef51325fe4a306e3590ca52fdad80b3/tests/rl/sharded_logps_head_test.py
docker run --rm -e JAX_PLATFORMS=cpu -v "$PWD/head-image/sharded_logps_head_test.py:/t.py:ro" \
  --entrypoint python "$HEAD_IMAGE" /t.py
# Ran 19 tests ... OK                 (about a minute on CPU; the stock bench image errors on all 19)

3. Render the e2e launch with DRY_RUN

Step 1 of the bench's README, with the measured launch's environment. It writes YAML to stdout and applies nothing; all three arms use it. Keep it local: it describes your environment.

(cd tunix-launcher && env -i HOME="$HOME" PATH="$PATH" KUBECONFIG=/dev/null DRY_RUN=true \
  JOB_PREFIX=tb K8S_NAMESPACE="$NS" PRIORITY_CLASS=low \
  TUNIX_IMAGE=gcr.io/tpu-vm-gke-testing/sivaibhav/trellis-experimental@sha256:b4647c188b25638e644c7a6fc53fb8e0751741525a277a1bcf07dca92ba9a6ac \
  PATHWAYS_SERVER_IMAGE=us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/gke/datenglin/unsanitized_server@sha256:662f33b0ca6cd236a12534e4c78a8b7d715868669366738a015401925a39cb86 \
  PATHWAYS_PROXY_IMAGE=us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/gke/datenglin/unsanitized_proxy_server@sha256:e714351cdde50d306dd62f127cffabcdaf98dff80ad22e8b4c66c5a1d1df6ddb \
  BATCH_SIZE=72 RAGGED_BUFFER_FACTOR=2.75 RETRY_WHEN_TOKENS_DROPPED=true REMAT_CONTEXT=device \
  MOE_EXCLUDE_PADDING=true LOG_MOE_DROPS=true COMPUTE_LOGPS_CHUNK_SIZE=1024 \
  bash tunix/experimental/examples/recipes/mlperf_gbs1024_submission.sh start) > launch-dryrun.yaml
grep -c '^kind: JobSet' launch-dryrun.yaml    # 130: orchestrator, trainer, 128 rollout engines

• Set exactly (trainer cell): the three images, BATCH_SIZE=72 (GBS 1,152: lr 2.121e-6, clip 0.0589), RAGGED_BUFFER_FACTOR=2.75 with RETRY_WHEN_TOKENS_DROPPED=true (recipe default: dropless, -1), REMAT_CONTEXT=device (default offload), MOE_EXCLUDE_PADDING=true, LOG_MOE_DROPS=true, COMPUTE_LOGPS_CHUNK_SIZE=1024 (default 512). The rest are recipe defaults.
• Scheduling only: K8S_NAMESPACE; PRIORITY_CLASS (must exist on your cluster; the measured runs used low); if your cluster differs, KUEUE_QUEUE (default multislice-queue), TPU_RESERVATION (default ghostfish-tb2qe8ysltros), SERVICE_ACCOUNT (default xpk-sa), HEAD_NODEPOOL (default cpu-np). Any JOB_PREFIX works: it survives only in four orchestrator paths, which the bench parses but never writes to.
• env -i, KUBECONFIG=/dev/null: even with DRY_RUN=true the recipe runs kubectl config use-context (to its own cluster), kubectl config set-context --current --namespace=... and kubectl get workload; with an empty kubeconfig these fail, the recipe ignores that, and the output is byte-identical. env -i keeps variables such as MAXTEXT_EXTRA_FLAGS or LIBTPU_INIT_ARGS out of the render.

4. Render the three bench JobSets

export LAUNCH_RENDER=$PWD/launch-dryrun.yaml K8S_NAMESPACE="$NS" STEPS=5 MICRO_STEPS=9 COMPUTE_LOGPS_CHUNK_SIZE=1024
OURS_REPLACE="num_moe_token_chunks=2 context_parallel_strategy=all_gather sa_fuse_reciprocal=true sa_use_base2_exp=false gdn_cp_mode=seq"
OURS_EXTRA="gdn_cp_state=device moe_route=device"
SC_FLAG="--xla_tpu_enable_offloading_copy_to_sparsecore=false"

# A: the baseline cell as rendered
NAME=tb-base bash "$BENCH/render_bench.sh"
# B: the tuned cell
NAME=tb-tuned REPLACE_MAXTEXT_FLAGS="$OURS_REPLACE" EXTRA_MAXTEXT_FLAGS="$OURS_EXTRA" \
  LIBTPU_EXTRA="$SC_FLAG" bash "$BENCH/render_bench.sh"
# C: the tuned cell + this PR
NAME=tb-tuned-head REPLACE_MAXTEXT_FLAGS="$OURS_REPLACE" EXTRA_MAXTEXT_FLAGS="$OURS_EXTRA" \
  LIBTPU_EXTRA="$SC_FLAG" IMAGE="$HEAD_IMAGE" DRIVER_ENV="TUNIX_LOGPS_SHARDED_HEAD=local" \
  bash "$BENCH/render_bench.sh"

Each command writes only rendered/<NAME>.yaml. REPLACE_MAXTEXT_FLAGS edits the client's MAXTEXT_EXTRA_FLAGS in place, failing unless the key occurs once: num_moe_token_chunks 4→2, context_parallel_strategy ring→all_gather, sa_fuse_reciprocal false→true, sa_use_base2_exp true→false, gdn_cp_mode head→seq. EXTRA_MAXTEXT_FLAGS appends gdn_cp_state=device and moe_route=device (absent from the baseline cell). LIBTPU_EXTRA reaches the client, the pathways-proxy (where XLA compiles) and the pw-node workers. IMAGE replaces only the main (client) container; DRIVER_ENV is the inline environment of python /tmp/bench_driver.py, which traces the trainer. STEPS/MICRO_STEPS give 5 steps (step 0 is warm-up; 1-4 are timed) of 9 fwd_bwd calls; COMPUTE_LOGPS_CHUNK_SIZE=1024 is already in the DRY_RUN and kept explicit. Leave TRAINER_SLICE, the payload knobs (PAYLOAD, ROUTING*, PAD_FRACTION, FINAL_PAD_FRACTION, PAYLOAD_SEED, EP_SIZE, ADVANTAGES, LOGP_GATES) and the profiler knobs unset, as the measured runs did. Optional arm A2 separates the libtpu flag from the 7 keys: NAME=tb-base-sc LIBTPU_EXTRA="$SC_FLAG" bash "$BENCH/render_bench.sh".

5. Check the renders before you spend the slice

python3 "$BENCH/structdiff_render.py" rendered/tb-base.yaml rendered/tb-tuned.yaml          # 15 differences
python3 "$BENCH/structdiff_render.py" rendered/tb-tuned.yaml rendered/tb-tuned-head.yaml  # 6 differences

B vs A lists the 7 MaxText keys, the libtpu flag at the client, proxy and worker sites, and the JobSet name, gang id, output folder and worker id. C vs B must list exactly these 6 (long lines cut):

[leaf] .metadata.labels.scheduling.x-k8s.io/gang-id: tb-tuned -> tb-tuned-head
[leaf] .metadata.name: tb-tuned -> tb-tuned-head
[leaf] .spec.replicatedJobs[proc].template.spec.template.spec.containers[main].image: gcr.io/tpu-vm-gke-testing/sivaibhav/trellis-experimental@sha256:b4647c18... -> <HEAD_IMAGE>
[trainer_argv] --maxtext_output_directory: ['gs://my-bucket/rl-trainer-bench/tb-tuned'] -> ['gs://my-bucket/rl-trainer-bench/tb-tuned-head']
[trainer_argv] --worker_id: ['tb-tuned'] -> ['tb-tuned-head']
[env] TUNIX_LOGPS_SHARDED_HEAD: <absent> -> local
6 differences

6. Run each arm

One arm at a time, in any order, for each NAME in tb-base, tb-tuned and tb-tuned-head (the bench README's commands):

NAME=tb-base
kubectl --context "$CTX" -n "$NS" apply -f rendered/$NAME.yaml
kubectl --context "$CTX" -n "$NS" logs -f --tail=-1 -c main \
  -l jobset.sigs.k8s.io/jobset-name=$NAME,jobset.sigs.k8s.io/replicatedjob-name=proc 2>&1 | tee $NAME.log
# when the log shows EXIT_CODE=0:
kubectl --context "$CTX" -n "$NS" delete jobset $NAME

If the pod is still pending, logs returns at once; rerun it. Keep --tail=-1, or a late attach with a label selector starts at the last 10 lines and misses the config and [logps-head] lines. Each attach replays the whole log, so tee overwrites $NAME.log. The client log shows, in order (times from arm C in pod 4; ranges over the measured runs):

1. Trainer worker initialized in ... on 256 devices. 295.0 s (6-379 s).
2. Bench micro-batch: ... completion_ids=(64, 65536) ...
3. [logps-head] sharded head local: ... twice, arm C only (step 7).
4. Trainer compile finished in ... 539.0 s (52-584 s; a program compiled before is faster).
5. Realistic payloads built in ... 182.1 s of host time, then PAYLOAD_ROUTING regular pad_fraction=0.0700 ... and final ... pad_fraction=0.3025.
6. completed step: 0 to completed step: 4, about 108 s each, each followed by STEP_LOSS step=N loss=0 grad_norm=0 ... (expected: Caveats, "Speed only").
7. Saving checkpoint at step 5. ... Finished save ... 91.8 s.
8. Program End: ..., EXIT_CODE=0 (the bench's PID-1 wrapper). A stack trace at interpreter shutdown (IFRT proxy server disconnected, then Py_FinalizeEx) is expected; the measured runs exited 0.

Then delete the checkpoints: gcloud storage rm -r "$OUTPUT_PREFIX/$NAME".

7. Read the result

One line per step; verbatim from arm C in pod 4 (the measured runs used the earlier driver, tagged [E2EBench]; at #5701's head the tag is [TrainerBench]):

[E2EBench] completed step: 1, seconds: 107.469, micro_steps: 9, seconds/micro_step: 11.941, fwd_bwd dispatch s: [0.015, 18.547, 10.859, 10.921, 10.917, 10.919, 10.921, 10.924, 10.922]

seconds is the whole step, from the first fwd_bwd dispatch until its device work has drained, including the optimizer update; use steps 1-4 (step 0 includes first execution). fwd_bwd dispatch s is how long each of the 9 calls blocked the driver: entry 1 returns at once (the previous step drained), entry 2 is longer, and from entry 3 on they settle to one regular micro-batch.

Two throughput rulers; do not mix them. A micro-batch is 64 rows x 65,536 tokens on 128 chips = 32,768 tokens per chip. A bench step runs 8 regular payloads (7.0% padding), then one final payload (30.25% padding); both rulers count padded rows.

• Regular micro-batch ruler (the one quoted in this PR): 32,768 / median of dispatch entries 3-9 over steps 1-4 (28 values).
• Step ruler: 9 x 32,768 / mean seconds of steps 1-4; it also pays the update and step boundary.

bench_readout.py (save it next to the logs) computes both, plus the head and forward checks:

python3 bench_readout.py tb-base.log tb-tuned.log tb-tuned-head.log

bench_readout.py (stdlib only)

#!/usr/bin/env python3
"""bench_readout.py LOG [LOG ...]: the two throughput rulers (step 7) from RL trainer bench client logs.

Skips step 0. Regular micro-batch time = median of fwd_bwd dispatch entries 3..K of the later steps. Rows, row length
and device count come from the log; tpu7x has 2 JAX devices per chip. Duplicate log lines are dropped.
"""
import re
import statistics
import sys

STEP = re.compile(r"completed step: (\d+), seconds: ([\d.]+), micro_steps: (\d+), seconds/micro_step: [\d.]+, "
                  r"fwd_bwd dispatch s: \[([^\]]*)\]")
SHAPE = re.compile(r"Bench micro-batch: .*?completion_ids=\((\d+), (\d+)\)")
DEVICES = re.compile(r"Trainer worker initialized in [\d.]+s on (\d+) devices")
HEAD = re.compile(r"\[logps-head\] sharded head (local:|off,)")
FWD = re.compile(r"STEP_METRICS step=(\d+) .*?sampler_is/token_logdiff_absmean=\[([^\]]*)\]")


def readout(path, devices_per_chip=2):
  steps, rows, length, devices, head, fwd = {}, None, None, None, {}, {}
  for line in open(path, errors="replace"):
    m = STEP.search(line)
    if m:
      n = int(m.group(1))
      rec = (float(m.group(2)), int(m.group(3)), [float(x) for x in m.group(4).split(",")])
      if steps.setdefault(n, rec) != rec:
        sys.exit(f"{path}: two different lines for step {n}")
    if (m := SHAPE.search(line)) and rows is None:
      rows, length = int(m.group(1)), int(m.group(2))
    if (m := DEVICES.search(line)) and devices is None:
      devices = int(m.group(1))
    if m := FWD.search(line):
      fwd[int(m.group(1))] = [float(x) for x in m.group(2).split(",")][:-1]  # drop the final micro-batch
    if m := HEAD.search(line):
      head[line.strip()] = "engaged" if m.group(1) == "local:" else "FELL BACK"  # duplicate captures count once
  timed = [n for n in sorted(steps) if n > 0]
  if not timed or rows is None or devices is None:
    sys.exit(f"{path}: need step lines after step 0, the Bench micro-batch line and the device count")
  tok = rows * length / (devices / devices_per_chip)
  k = steps[timed[0]][1]
  step_mean = statistics.mean(steps[n][0] for n in timed)
  regular = [x for n in timed for x in steps[n][2][2:]]
  mb = statistics.median(regular)
  print(f"{path}\n  steps {timed[0]}-{timed[-1]}: s/step " + " / ".join(f"{steps[n][0]:.3f}" for n in timed) +
        f"  mean {step_mean:.3f}")
  print(f"  regular micro-batch: median {mb:.4f} s of {len(regular)} dispatch entries (3..{k} of each step), "
        f"range {min(regular):.3f}-{max(regular):.3f}")
  print(f"  tokens per chip per micro-batch: {rows} x {length} / {devices // devices_per_chip} chips = {tok:,.0f}")
  print(f"  tok/s/chip  regular micro-batch ruler {tok / mb:,.0f}   step ruler ({k} x {tok:,.0f} / s/step) "
        f"{k * tok / step_mean:,.0f}")
  if fwd:
    v = [x for n in sorted(fwd) for x in fwd[n]]
    print(f"  forward check: sampler_is/token_logdiff_absmean mean {statistics.mean(v):.7f} over {len(v)} regular "
          f"micro-batches (steps {min(fwd)}-{max(fwd)}), range {max(v) - min(v):.1e}")
  vals = list(head.values())
  print(f"  [logps-head] lines: {vals.count('engaged')} engaged, {vals.count('FELL BACK')} fell back")


if __name__ == "__main__":
  for p in sys.argv[1:]:
    readout(p)

Its output on arm C's log from pod 4:

  steps 1-4: s/step 107.469 / 107.528 / 107.487 / 107.683  mean 107.542
  regular micro-batch: median 10.9190 s of 28 dispatch entries (3..9 of each step), range 10.859-10.925
  tokens per chip per micro-batch: 64 x 65536 / 128 chips = 32,768
  tok/s/chip  regular micro-batch ruler 3,001   step ruler (9 x 32,768 / s/step) 2,742
  forward check: sampler_is/token_logdiff_absmean mean 12.2219076 over 40 regular micro-batches (steps 0-4), range 2.3e-05
  [logps-head] lines: 2 engaged, 0 fell back

Proof that the head ran (arm C): exactly two such lines during the trainer compile (presumably one per traced program, fwd_bwd and its dropless retry variant; not verified), and no sharded head off line:

[logps-head] sharded head local: tokens batch=('data', 'fsdp', 'expert') seq=('context',) (logical-axis-rules), 1x16384 local tokens in 16 steps of 1024; gathers [((4096,), 'bfloat16'), ((4096, 248320), 'bfloat16')] over ('fsdp', 'context', 'expert')

That is 1 row x 16,384 tokens per device (64 rows over fsdp 32 x expert 2; 65,536 tokens over context 4), and both head parameters are gathered once, in bf16. Arms A and B print no [logps-head] line; neither does arm C if its client lacks this PR's file (the stock image ignores the switch).

Other checks: python3 "$BENCH/diff_config.py" tb-tuned.log tb-tuned-head.log should list only 8 path and name keys of the resolved MaxText configs (base_output_directory, checkpoint_dir, dump_hlo_gcs_dir, dump_jaxpr_gcs_dir, managed_mldiagnostics_dir, metrics_dir, run_name, tensorboard_dir). grep -c 'MOE_DROP_AUDIT_STEP .*dropped_assignments_used=0 dropless_replays=0' <log> gives 5 in every arm: no token dropped, no micro-batch replayed.

Expected numbers

Measured 2026-10-13/14 with the step 7 rulers, steps 1-4. The A/B of this PR is the pod 5 pair (B, then C, in one JobSet on one slice); the other rows show how far the same program moves between pods.

arm	pod	s/step (steps 1-4)	mean	regular micro-batch	tok/s/chip, micro-batch ruler	tok/s/chip, step ruler
A, baseline cell (under the bench's own wrapper)	1	130.872 / 130.807 / 130.720 / 130.815	130.804	13.522 s	2,423	2,255
A, baseline cell	2	130.785 / 130.754 / 130.812 / 130.869	130.805	13.505 s	2,426	2,255
A, baseline cell	8	130.723 / 130.719 / 130.766 / 130.910	130.780	13.511 s	2,425	2,255
A2, baseline cell + libtpu flag (optional)	3	126.434 / 126.407 / 126.360 / 126.253	126.364	13.018 s	2,517	2,334
B, tuned cell	5	121.949 / 122.035 / 122.042 / 121.980	122.002	12.527 s	2,616	2,417
B, tuned cell	3	121.759 / 121.773 / 121.776 / 121.634	121.736	12.492 s	2,623	2,423
C, tuned cell + this PR	5	107.595 / 107.640 / 107.415 / 107.723	107.593	10.937 s	2,996	2,741
C, tuned cell + this PR	4	107.469 / 107.528 / 107.487 / 107.683	107.542	10.919 s	3,001	2,742
C, tuned cell + this PR	6	108.061 / 107.902 / 107.935 / 107.853	107.938	10.949 s	2,993	2,732

Every cross-pod B/C pairing gives -13.80 to -14.46 s/step and -1.54 to -1.61 s per regular micro-batch. Pods 5 and 6 ran with enable_checkpointing=false (the final checkpoint comes after the timed steps), and pod 5's B ran this PR's file with TUNIX_LOGPS_SHARDED_HEAD=off, the same path as unset; otherwise their B and C arms match the step 4 renders in every field that defines the trainer program, and diff_config.py on pod 5's B and C shows only the 8 path and name keys.

How close to expect: within a run the 28 dispatch entries span 0.07-0.08 s. Across pods the same program moved 0.03 s per step for the baseline cell (pods 1, 2 and 8), 0.27 s for B (two pods) and 0.40 s for C (three pods); the regular micro-batch time moved 0.04 s (B) and 0.03 s (C). Your B and C will usually run in different pods: on the pinned images, expect C vs B within about 0.6 s per step of pod 5's -14.4 s, as the cross-pod pairings were (pod 10, newer images: -13.74).

How the measured runs differed from these steps

None of these changes the trainer program. (1) Driver: e2e_driver.py at #5701 d6db22f99 (md5 e78c082f) instead of bench_driver.py at 70d64441 (md5 a2fdaf1a); they differ in 3 lines, the log tag ([E2EBench] → [TrainerBench]) and the two argv variable names (E2E_TRAINER_ARGV/E2E_ORCH_ARGV → TRAINER_ARGV/ORCH_ARGV). (2) Wrapper: a multi-arm wrapper instead of the bench's PID-1 wrapper; timing is computed inside the driver, and the baseline cell measured 130.804 s per step under the bench's wrapper (pod 1) and 130.805 / 130.780 under the multi-arm one (pods 2 and 8). (3) This PR's file was copied over the image's at pod start with md5 guards (raw in pod 4, gzip+base64 in pods 5-7) instead of baked in: same bytes (md5 a5bfda56), same path, stale .pyc removed. (4) The switch was exported before the driver (pod 4) or set inline on the driver line (pods 5-7) instead of DRIVER_ENV; both reach the driver's environment. (5) Keys, flag, images: keys re-appended, the flag inserted mid-list and image tags swapped for these digests after rendering, with the same result: the same 74 keys and values, flag sets (17 client, 21 proxy, 19 worker) and images. (6) Cluster fields (namespace, JobSet-level slice-topology annotation, activeDeadlineSeconds, names): scheduling only.

Rendered from a DRY_RUN made as in step 3, all three arms match the launched manifests field by field: the MaxText flags as a set and as last-wins values, every libtpu flag at the three sites, the trainer argv (53 items, apart from the output folder and worker id), the orchestrator argv (88 items, apart from the four JOB_PREFIX paths), the 35 other client variables, and every image and YAML field apart from (1)-(6).

Without a registry, install the file at pod start as the measured runs did: render arm C without IMAGE (keep DRIVER_ENV); in the client script, before the driver heredoc, write the file, exit 3 unless it is md5 a5bfda56 and /app/tunix/rl/common.py is md5 be286d30, copy it over that path and delete /app/tunix/rl/__pycache__/common.*.pyc. The client script is one exec argument, limited to 131,071 bytes: pod 4's single-arm script fit with the file raw (114.5 KB); pods 5-7's multi-arm script needed gzip+base64.

Caveats

• Speed only: with the e2e orchestrator's flags the log-prob gates drop every sequence of a random-weight model (sample_mask/kept_frac=0), so loss and grad_norm are exactly 0 in every speed arm. For gradients use the numerics mode, as pod 7 did: DRIVER_ENV="TUNIX_LOGPS_SHARDED_HEAD=local LOGP_GATES=off ADVANTAGES=random" for C, the same without the switch for B.
• tok/s/chip counts padded rows: at 7.0% padding per regular payload (30.25% in the final one), 2,996 is about 2,786 on real tokens.
• The baseline cell's --xla_tpu_enable_sparse_core_reduce_scatter_v2=true is logged by the Pathways proxy as deprecated on every compile, in all arms; whether it still takes effect is unknown.
• Profiling (optional, never in a timed run): PROFILER_SKIP=12 PROFILER_STEPS=2 traces fwd_bwd calls 12 and 13 (regular payloads of step 1; calls count from 0, 9 per step) to $OUTPUT_PREFIX/<NAME>/tensorboard/. A step that overlaps the window carries the profiler's cost.

Limitations and follow-ups

Line numbers are at a29991d3f (tunix/rl/common.py md5 a5bfda56676e92d65261f3e24303101b, test md5 060eac1d95da6bfa14da01d6f1f95d76). None of these gives a wrong result on the measured path: the MaxText training engine (jax.set_mesh) with bf16 weights.

1. Engages only under jax.set_mesh. The head reads jax.sharding.get_abstract_mesh() (common.py:1188-1190), which tunix's own RL stack (RLEngine, alias RLCluster; legacy with mesh: at rl_cluster.py:1293-1312) leaves empty, so there it falls back (... no mesh in context (jax.set_mesh)). Follow-up: when it is empty, use pxla.thread_resources.env.physical_mesh, which tunix/sft/peft_trainer.py already reads.
2. Explicit-axis meshes raise instead of falling back, inside jax.shard_map (common.py:1393, ValueError: in_specs passed to shard_map ... does not match the specs of the input). jax 0.11.1's jax.make_mesh builds explicit axes unless axis_types is given; tunix's create_mesh (tunix/utils/mesh.py:114,119) and the bench's MaxText mesh use Auto. The docstring (common.py:1159-1162) promises a fallback for any failed precondition and omits the fourth (:1216); errors while tracing the head abstractly (in _head_parameters and the abstract leaf list) also propagate. Follow-up: fall back on explicit axes (or jax.sharding.reshard first), turn precondition errors into fallbacks, list all four fallbacks.
3. Memory that does not shrink with the mesh (sizes in Risks, "Memory"). Until the one reduce-scatter, the backward holds each head gradient at full gathered shape in float32, and the forward keeps the gathered parameters, including the 2 GB bf16 table, as residuals (head_local_fwd). The arm B/C TPU runs fit. Follow-up (optional): gather again in the backward instead of keeping the residual.
4. A misspelled switch raises only on packed, chunked calls (sharded_head_mode(), common.py:1036-1043, reached via :644-657); elsewhere TUNIX_LOGPS_SHARDED_HEAD=locla is silently ignored. Follow-up: parse it on every call.
5. The switch is read at trace time and is not in any jit cache key, so changing it after the first trace does nothing in that process (the unit tests call jax.clear_caches()). Follow-up: a typed config field next to compute_logps_chunk_size, passed as a static argument, would close 4 and 5.
6. dtype detection can be fooled by a contrived head. _use_dtypes (common.py:1101-1119) sees direct consumers and inner jaxprs only. If a head passes the table unchanged out of an inner call (for example jax.jit(lambda w: w)(table)) and uses that in f32, while another use casts it to bf16, the table is gathered in bf16 and the f32 use silently gets bf16-rounded weights: a test head built this way was off by up to 1.19 nats on 251 of 256 tokens, with no warning. Nothing is narrowed at bf16 weights; in the one fp32-weight TPU arm (pod 9) only the table was, to the bf16 its dot uses. Follow-up: narrow only when every use is a top-level convert_element_type.
7. No length check on direct calls of maybe_compute_sharded_chunked_logps (common.py:1149): dynamic_slice clamps, so hidden states 8 positions longer than the tokens return a result without an error. From the source, compute_per_token_logps cannot trigger this with the models in this repo. Follow-up: fall back when hidden_states.shape[:2] != input_tokens.shape.
8. The test module sets the device count for the whole process: tests/rl/sharded_logps_head_test.py:19 overwrites XLA_FLAGS with --xla_force_host_platform_device_count=8 at import. The CI command, a single-process pytest tests/rl/ (.github/workflows/cpu-tests.yml:76), imports it last, so every RL test there sees 8 CPU devices, not the 4 rl_utils_test.py set before; on a backend already up with fewer devices, _mesh() fails with an unhelpful reshape error. Follow-up: run the module in its own pytest process (its own CI step, --ignored in the tests/rl/ step) with tunix.tests.test_common.safe_set_n_cpu_devices(8) in setUpClass; in the shared process that call raises once JAX is up with fewer devices.
9. Custom axis rules listing two sequence axes out of mesh order fail to compile. _token_partition keeps the rule's order and neighbour passes it to jax.lax.ppermute, which in jax 0.11.1 raises at lowering when it differs from the mesh order (jax_raise_on_ppermute_sort_diff, on by default). From the source; not run. MaxText's built-in rules are in mesh order, and the TPU runs used one sequence axis. Follow-up: sort both axis tuples into mesh order.

Tests worth adding: sharded head parameters (the tests replicate them, MaxText shards the table; a one-off CPU check matched); the with mesh: fallback and its warning; tied embeddings, whose two table cotangents must add (a one-off CPU check matched); the two unreached fallbacks, asserting the warning text; one test per fix above; the backward's entropy-only branch (never exercised), its neither-output branch (matched in a one-off CPU check) and GRPO's case, entropy requested but kept out of the gradient (covered by the CPU gate and TPU runs, not by unit tests); and in every parity test, checks of engagement, layout and gather dtypes over every output.

Not tested

• Gradients on the newer images: pod 10 (Pathways raiden_20261013212200, the tag this PR's base 8d245c561 / build(raiden): update Raiden wheel to dev20261013212200 and Pathways images to raiden_20261013212200 #2862 moves the recipe to) ran only the speed A/B and the forward check. The gradient A/B (pod 7) ran only on raiden_20261009183300_cl996655201 with client b4647c18.
• End-to-end RL with real weights: no RL run with the real checkpoint, live rollouts and rewards has used this head; under tunix's own RL stack it would fall back (limitation 1).
• Other meshes and slice sizes: the A/Bs are on fsdp 32 x context 4 x expert 2 on one 128-chip v7x slice, at chunk 1024 and 4096. Speed only, with no head-off control: on a 256-chip slice, fsdp 64 x context 4 x expert 2 (128-row micro-batches) and fsdp 32 x context 8 x expert 2 both compiled and ran with the head engaged (1x16384 local tokens in 16 steps of 1024 and 2x8192 local tokens in 8 steps of 1024) and 0 MoE drops. No other model.
• fp32 head weights, beyond speed: the one fp32 arm (pod 9, weight_dtype=float32) engaged and gathered the norm scale in float32 and the table in bfloat16, the logits dot's dtype (gathers [((4096,), 'float32'), ((4096, 248320), 'bfloat16')]). With no head-off fp32 arm, its forward check (+7.5e-4 against bf16 weights) does not test the head; no fp32 gradient was compared.
• A frozen head (for example LoRA): the custom VJP ignores the perturbed flags, so the backward builds and reduce-scatters the f32 accumulator even for an untrained head. Whether XLA removes that dead work was not checked.
• The reproduction steps as written: verified offline (each render matches its launched manifest field by field; the derived image builds and passes the unit tests), not re-run on TPU.

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!