Add an RL trainer bench under experimental/rl_trainer_bench by ultrons · Pull Request #5701 · AI-Hypercomputer/maxtext

ultrons · 2026-10-13T18:51:14Z

Description

This PR adds a standalone bench for the RL trainer program of a Tunix GRPO e2e run, under src/maxtext/experimental/rl_trainer_bench/. It runs the exact trainer program the e2e run compiles (same MaxText flags, libtpu flags, mesh, GRPO loss and chunked log-prob pass) on one trainer slice under the same Pathways setup, with a synthetic micro-batch whose MoE routing load is calibrated to a production run. There is no orchestrator, no rollout engines and no weight sync, so a MaxText flag, libtpu flag or slice-size variant can be A/B tested and profiled in one short job on one slice.

We built it because maxtext.experimental.maxtext_engine.train runs MaxText's pre-training loss on MaxText's synthetic loader, which is a different program from the RL trainer (no GRPO loss, no vocabulary-chunked compute_per_token_logps, a different batch layout), so its timings and compile results did not transfer to the RL step.

How it works

render_bench.py takes the DRY_RUN output of the e2e launcher (e.g. DRY_RUN=true mlperf_gbs1024_submission.sh start), picks the trainer and orchestrator JobSets by their --process_main, and writes a bench JobSet that keeps every scheduling field and the Pathways containers of the e2e trainer. Only the names, maxRestarts and the client command change. The client runs bench_driver.py, which builds the trainer worker, GRPO algorithm, BatchConfig and assembler through Tunix's own parsers and factories from the e2e argv, compiles against Tunix's TRAINER_WARMUP_COMPILE payload, and times STEPS optimizer steps of MICRO_STEPS fwd_bwd calls. Knobs cover MaxText flags (append or replace in place), libtpu flags (written to client, proxy and workers), trainer and orchestrator argv, the trainer slice (FSDP fills it), the logps chunk and the profiler window (in fwd_bwd calls). Everything environment-specific (namespace, output prefix, image, the e2e render itself) is an input; the README uses placeholders.

The driver calls the trainer worker as the orchestrator does. When the worker applies the optimizer inside the last fwd_bwd(apply_optimizer=True) and reports updated (google/tunix#2719), the driver does not call update() again; an older worker gets the explicit call. When the worker has get_step_timing, the driver logs one TRAINER_STEP_TIMING line per step in the orchestrator's format, so the per-micro-batch device timeline of a bench run reads the same as an e2e run's.

The payload keeps the warm-up payload's shapes and fields and replaces its contents (which at production shapes are mostly padding with a single token id and no replayed routing) with varied token ids, tail padding and replayed top-k routing calibrated per layer to a measured load profile. The default table comes from a Qwen3.5-397B-A17B production run and can be overridden with ROUTING_TARGETS or stretched with ROUTING_SCALE.

Also included: an AOT memory check that takes its flags and mesh from the same e2e render, a structured diff of two renders, a diff of resolved MaxText configs, CPU checks of the payload and of the numerics knobs (LOGP_GATES=off, ADVANTAGES=random, which make gradient norms comparable between variants on random weights), and a static summary of collective overlap in a scheduled HLO dump.

Placement

We used experimental/ because it already holds the standalone engine script and the RL experiments, whereas the top-level benchmarks/ directory holds launch recipes for MaxText's own training loop. Most of what the driver calls is Tunix code (run_trainer_node, run_deepswe_dist, trainer_warmup, TrainerWorker), so a reviewer may prefer it next to the Tunix recipes; we are happy to move it.

Caveats

Random weights, a synthetic (calibrated) payload, a static AOT check, and a host-side gap after each step that the e2e trainer does not have; the README has the details.

Tests

python3 -m pytest src/maxtext/experimental/rl_trainer_bench/tests runs 26 CPU-only tests that need PyYAML and NumPy alone (all pass). They exercise the renderer on a hand-written DRY_RUN fixture (JobSet selection, names, libtpu flags in all three places, in-place replacement and its failures, the trainer-slice fields checked with the structured diff, profiler and driver knobs, the AOT env), the payload builder on a small stand-in payload (shapes and dtypes, padding, distinct replayed experts, per-layer load targets, final micro-batch padding, seeded advantages, metric summaries), and the end-of-step helpers (no second update() after a fused one, the explicit update() for older workers, the TRAINER_STEP_TIMING line and its failure path). They live next to the code, so the default testpaths = tests does not collect them.

pre-commit run --files <changed files> passes (codespell, pylint, pyink, yamllint).

We also rendered three configurations from a real DRY_RUN output with this code and got files byte-identical to those of our previous internal version (up to the renamed driver file and argv variable names), and check_numerics_cpu.py --launch-render <render> passes in a trainer image whose Tunix matches the one it was written for. The driver, including the fused-update and step-timing handling, has run on TPU in its internal form; this generalized version has not been run on a cluster. The generalized AOT path (aot_bench.sh taking its flags and mesh from a real DRY_RUN render, then aot_summarize.py) ran to completion on a CPU host.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
I have necessary comments in my code, particularly in hard-to-understand areas.
~~I have run end-to-end tests tests and provided workload links above if applicable.~~ Not applicable: a standalone bench under experimental/ that changes no training code; the internal form of the driver ran on TPU (see Tests).
I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

A standalone timing harness for the RL trainer program of a Tunix GRPO e2e run: the same MaxText and libtpu flags, mesh, GRPO loss and chunked log-prob pass on one trainer slice, driven through Tunix's TrainerWorker with a synthetic micro-batch whose MoE routing load is calibrated to a production run, and without the orchestrator or rollout engines. render_e2e.py turns the DRY_RUN output of the e2e launcher into the bench JobSet, keeping the trainer JobSet's scheduling fields and Pathways containers and replacing only names, maxRestarts and the client command, with knobs for MaxText flags, libtpu flags, trainer and orchestrator argv, trainer slice and profiler window. Also included: an AOT memory check, a structured render diff, a MaxText config diff, CPU checks of the payload and the numerics knobs, a static HLO collective-overlap summary, and CPU unit tests of the renderer (on a hand-written fixture) and the payload builder.

…e DRY_RUN context switch

gemini-code-assist

Code Review

This pull request introduces a standalone benchmark harness for the Tunix GRPO RL trainer, featuring a driver, a JobSet renderer, ahead-of-time memory and CPU verification tools, and corresponding unit tests. The code review feedback highlights several critical robustness and portability improvements, including resolving potential runtime errors in 'rng.choice' and '.index()' calls, preventing resource leaks by using 'with' statements for file operations, and explicitly specifying UTF-8 encoding to avoid platform-dependent file decoding issues.

gemini-code-assist · 2026-10-13T18:55:42Z

+  real_mask = np.zeros((rows, budget), dtype=bool)
+  for r in range(rows - n_pad_rows):
+    n_seg = max(1, len(np.unique(seg0[r][seg0[r] > 0])))
+    cuts = np.sort(rng.choice(np.arange(64, real_len, 64), size=n_seg - 1, replace=False)) if n_seg > 1 else []


Using rng.choice with replace=False on np.arange(64, real_len, 64) can raise a ValueError if the population size (the number of 64-token chunks) is smaller than n_seg - 1. This can happen with smaller test shapes or custom configurations where real_len is small.

Consider adding a fallback (such as using np.linspace or a smaller step size) when the population size is insufficient.

Suggested change

cuts = np.sort(rng.choice(np.arange(64, real_len, 64), size=n_seg - 1, replace=False)) if n_seg > 1 else []

pop = np.arange(64, real_len, 64)

if len(pop) >= n_seg - 1:

cuts = np.sort(rng.choice(pop, size=n_seg - 1, replace=False)) if n_seg > 1 else []

else:

cuts = np.sort(rng.choice(np.arange(1, real_len), size=n_seg - 1, replace=False)) if n_seg > 1 and real_len > n_seg else np.linspace(0, real_len, n_seg + 1, dtype=int)[1:-1]

gemini-code-assist · 2026-10-13T18:55:42Z

+  """The JobSet-like documents of a DRY_RUN output: YAML documents split on `---`; chunks that do not parse as a
+  mapping (launcher log lines, kubectl output) are skipped."""
+  docs = []
+  with open(path) as f:


Omitting the encoding parameter when opening files can cause UnicodeDecodeError on platforms where the default system encoding is not UTF-8. It is recommended to explicitly specify encoding="utf-8".

Suggested change

with open(path) as f:

with open(path, encoding="utf-8") as f:

gemini-code-assist · 2026-10-13T18:55:42Z

+  i = cmd.index("(\n") + 2
+  j = cmd.index(RUNTIME_MAIN)
+  k = cmd.index("\n)", j)


Using .index() directly can raise a generic ValueError: substring not found if the input command format deviates from the expected structure (e.g., missing (\n or \n)). Wrapping this in a try-except block with a more descriptive error message will improve debuggability and maintainability.

try: i = cmd.index("(\n") + 2 j = cmd.index(RUNTIME_MAIN) k = cmd.index("\n)", j) except ValueError as e: raise ValueError( f"Failed to parse command. Expected to find '(\\n', '{RUNTIME_MAIN}', and '\\n)' in: {cmd}" ) from e

gemini-code-assist · 2026-10-13T18:55:42Z

+  md.setdefault("labels", {})["scheduling.x-k8s.io/gang-id"] = a.name
+  js["spec"].setdefault("failurePolicy", {})["maxRestarts"] = 0
+
+  with open(a.driver) as f:


Omitting the encoding parameter when opening files can cause UnicodeDecodeError on platforms where the default system encoding is not UTF-8. It is recommended to explicitly specify encoding="utf-8".

Suggested change

with open(a.driver) as f:

with open(a.driver, encoding="utf-8") as f:

gemini-code-assist · 2026-10-13T18:55:42Z

+    extras = [f for f in env["value"].split() if f.split("=", 1)[0] not in base_keys]
+    env["value"] = " ".join(libtpu + extras)
+
+  with open(a.out, "w") as f:


Omitting the encoding parameter when opening files can cause UnicodeDecodeError on platforms where the default system encoding is not UTF-8. It is recommended to explicitly specify encoding="utf-8".

Suggested change

with open(a.out, "w") as f:

with open(a.out, "w", encoding="utf-8") as f:

gemini-code-assist · 2026-10-13T18:55:42Z

+  for line in open(path):
+    if cur is None:
+      m = re.match(r"^(ENTRY )?%([\w.\-]+) .*\{\s*$", line)  # headers start at column 0, instructions are indented
+      if m:
+        cur = m.group(2)
+        comps[cur] = []
+      continue
+    if line.rstrip() == "}":  # multi-line backend configs also start lines with "}}"
+      cur = None
+      continue
+    m = re.match(r"^\s*(ROOT )?%([\w.\-]+) = (.*)$", line)
+    if m:
+      comps[cur].append((m.group(2), m.group(3)))


Opening a file without a with statement can lead to resource leaks as the file descriptor is not guaranteed to be closed promptly. Additionally, omitting the encoding parameter can cause UnicodeDecodeError on platforms where the default system encoding is not UTF-8. Use a with statement and specify encoding="utf-8".

Suggested change

for line in open(path):

if cur is None:

m = re.match(r"^(ENTRY )?%([\w.\-]+) .*\{\s*$", line) # headers start at column 0, instructions are indented

if m:

cur = m.group(2)

comps[cur] = []

continue

if line.rstrip() == "}": # multi-line backend configs also start lines with "}}"

cur = None

continue

m = re.match(r"^\s*(ROOT )?%([\w.\-]+) = (.*)$", line)

if m:

comps[cur].append((m.group(2), m.group(3)))

with open(path, encoding="utf-8") as f:

for line in f:

if cur is None:

m = re.match(r"^(ENTRY )?%([\w.\-]+) .*\{\s*$", line) # headers start at column 0, instructions are indented

if m:

cur = m.group(2)

comps[cur] = []

continue

if line.rstrip() == "}": # multi-line backend configs also start lines with "}}"

cur = None

continue

m = re.match(r"^\s*(ROOT )?%([\w.\-]+) = (.*)$", line)

if m:

comps[cur].append((m.group(2), m.group(3)))

gemini-code-assist · 2026-10-13T18:55:42Z

+  for line in open(path, errors="replace"):
+    m = re.search(r"Config param ([A-Za-z0-9_]+): (.*)$", line.rstrip("\n"))
+    if m:
+      out[m.group(1)] = m.group(2).strip()


Opening a file without a with statement can lead to resource leaks as the file descriptor is not guaranteed to be closed promptly. Additionally, omitting the encoding parameter can cause UnicodeDecodeError on platforms where the default system encoding is not UTF-8. Use a with statement and specify encoding="utf-8".

Suggested change

for line in open(path, errors="replace"):

m = re.search(r"Config param ([A-Za-z0-9_]+): (.*)$", line.rstrip("\n"))

if m:

out[m.group(1)] = m.group(2).strip()

with open(path, encoding="utf-8", errors="replace") as f:

for line in f:

m = re.search(r"Config param ([A-Za-z0-9_]+): (.*)$", line.rstrip("\n"))

if m:

out[m.group(1)] = m.group(2).strip()

gemini-code-assist · 2026-10-13T18:55:42Z

+
+
+def parse(path):
+  with open(path) as f:


Omitting the encoding parameter when opening files can cause UnicodeDecodeError on platforms where the default system encoding is not UTF-8. It is recommended to explicitly specify encoding="utf-8".

Suggested change

with open(path) as f:

with open(path, encoding="utf-8") as f:

gemini-code-assist · 2026-10-13T18:55:42Z

+
+
+def summarize(path, cap_gib):
+  with open(path, errors="replace") as f:


Omitting the encoding parameter when opening files can cause UnicodeDecodeError on platforms where the default system encoding is not UTF-8. It is recommended to explicitly specify encoding="utf-8".

Suggested change

with open(path, errors="replace") as f:

with open(path, encoding="utf-8", errors="replace") as f:

The bench is standalone; the e2e names only meant "derived from the e2e launch's DRY_RUN render". Renames, no behaviour change: - E2E_RENDER -> LAUNCH_RENDER; --e2e-render -> --launch-render (render_bench.py and check_numerics_cpu.py) - render_e2e.{sh,py} -> render_bench.{sh,py} - e2e_driver.py -> bench_driver.py (also /tmp/bench_driver.py in the pod) - E2E_TRAINER_ARGV / E2E_ORCH_ARGV -> TRAINER_ARGV / ORCH_ARGV - aot_e2e.sh -> aot_bench.sh (default log logs/aot-bench-<time>.log) - tests/render_e2e_test.py -> tests/render_bench_test.py; tests/testdata/e2e_dryrun_fixture.yaml -> launch_dryrun_fixture.yaml - heredoc E2E_DRIVER_EOF -> BENCH_DRIVER_EOF; log tag [E2EBench] -> [TrainerBench]; RenderE2ETest -> RenderBenchTest README states that the bench is standalone and that LAUNCH_RENDER is the full launcher's DRY_RUN output. Prose "e2e" meaning the full RL run is kept.

TRAINER_TUNED_FLAGS=true now also appends gdn_cp_matmul_precision=high and cast_grads_after_all_reduce=true to the trainer's MaxText flags, in the same block as the seven keys: - gdn_cp_matmul_precision=high: under gdn_cp_mode=seq, the float32 matmuls that compose the GDN state across the sequence shards run 3 bf16 passes instead of 6. - cast_grads_after_all_reduce=true: a gradient that the forward/backward pass sums by an all-reduce alone, and that is not already float32, leaves it in the parameter's dtype and is cast to float32 after the all-reduce. Both keys have these meanings at this branch's MaxText pin (60425bf) and in the bench image's MaxText (3792fdef8). Bench (AI-Hypercomputer/maxtext#5701, 128 v7x chips), measured with the sharded log-prob head of #2873 on in both arms: 107.59 -> 106.24 s/step in one pod, and 107.71 -> 106.41 in a second pod that trained on real gradients. In that pod the gradient norm moved 1.47e-5 relative (two runs with the keys differ by 1.92e-5), the forward log-prob metric moved 1.1e-6 (3.6e-7), and the loss was identical as printed. Four MoE settings measured alongside them (prefuse_moe_weights=false, moe_accumulate_wi_dlhs=true, moe_accumulate_chunk_wgrad=true, decoder_layer_input=device) moved the gradient norm by -0.80% and the forward metric by 9.7e-4, so they are left out. Tests: the exact-set test expects nine keys, and _MAXTEXT_DEFAULTS has the two keys' MaxText defaults.

This reverts commit 0968476, reversing changes made to 5f1e9ad.

…_HEAD) compute_chunked_logps scans the sequence under GSPMD. On the MLPerf Qwen3.5-397B trainer cell (fsdp32 x cp4 x ep2 on 256 devices, 64 x 65,536-token packed micro-batches, compute_logps_chunk_size 1024) the compiled program all-gathers the 2 GB vocab table in every scan step (forward and rematerialized backward) and all-reduces a 2 GB table gradient in every backward step: 129 table all-gathers and 65 table all-reduces per micro-batch. Because the scan runs over a context-sharded sequence, every context rank also computes every chunk, so the vocab matmuls are 4x redundant. With TUNIX_LOGPS_SHARDED_HEAD=local (or 1, true, on, yes), packed rows go through a shard_map'ed head that runs the same per-chunk math: - each head parameter is all-gathered once per pass; - each device scans only the tokens it holds (dynamic slices, no reshape of the sequence); - the table gradient is accumulated per device in float32 and reduce-scattered once; - the one-token shifts (targets left, outputs right) are explicit collective-permutes inside the shard_map, so they do not depend on how XLA partitions a shift of a sequence-sharded array. The head falls back to compute_chunked_logps, with a log line, when there is no mesh in context, when some mesh axis shards neither the batch nor the sequence, or when a head parameter does not divide over the mesh. Sequence axes are used in mesh order, the order in which jax.lax.ppermute numbers shards, so a logical-axis rule that lists them in another order reshards the input instead of shifting tokens across the wrong shards. Unset (the default), nothing changes. k8s_launcher.sh passes TUNIX_LOGPS_SHARDED_HEAD to the trainer process when it is set. TRAINER_EXTRA_ENV cannot carry it from the 397B v7x recipes: a value set there replaces mlperf_397b_256_v7x.sh's default TRAINER_EXTRA_ENV, including the trainer's LIBTPU_INIT_ARGS. Measured on the RL trainer bench (AI-Hypercomputer/maxtext#5701, 128 v7x chips), same pod, only the switch differs: 122.00 -> 107.59 s per optimizer step. Real-gradient A/B on TPU (reference, head, reference): loss bit-identical; grad_norm +2.8e-4 relative to the references' mean, whose own difference is 5.3e-5. Tests: tests/rl/sharded_logps_head_test.py on an 8-device CPU mesh (under pytest it runs in a child process unless the process already has 8 CPU devices): parity with compute_chunked_logps with the head engaged, token layouts, fallbacks, the legacy `with mesh:` fallback, the switch's scope and parsing, sequence axes listed out of mesh order, non-finite hidden states at masked positions, and GRPO's pattern of differentiating the log-probs only (and the entropy only). Launcher: unset, the variable appears nowhere in the render; set, it reaches the trainer JobSet only. The file measured on TPU is #2873 at a29991d. This version adds the mesh-order line, docstrings and tests; for in-order layouts the traced value_and_grad program is unchanged (same jaxpr hash, head on and off).

…AGE_NEXT_MICROBATCH) Between optimizer steps the trainer chips sat idle while the orchestrator closed step k (metrics, logging, checkpoint, weight sync) and then packed and shipped step k+1's first micro-batch: ~7.7 s per step on the 397B v7x trainer. - --pipeline_step_boundary (PIPELINE_STEP_BOUNDARY): pack the next step's first micro-batch while this step's last ones train, send it as soon as the checkpoint, weight sync (its staging, with --async_weight_sync) and commit are done, and run step k's metrics and logging once that call has returned, while the device computes it. The trainer gets the same calls in the same order. - --stage_next_microbatch (STAGE_NEXT_MICROBATCH, requires the above): send the next step's first micro-batch to the trainer host while the device still computes this step's last one (TrainerWorker.stage_fwd_bwd), so its train call carries no data. - --stage_before_update (default on): stage between the step's last micro-batch and its update, which the trainer defers (fwd_bwd(apply_optimizer=True, defer_update=True), then apply_deferred_update()). Under MaxText's retry_when_tokens_dropped, update() waits for the last micro-batch's overflow flag, so a stage after the update's dispatch would find the device idle and latch staging off. Staging needs MaxText's prepare_fwd_bwd_input and wait_for_device (AI-Hypercomputer/maxtext#5718). Without them, or when a stage measures no overlap, staging latches off and the step runs as before. While an update is deferred the worker refuses whatever needs the step closed (another micro-batch, update, get_metrics, a checkpoint, weight sync, actor log-probs, eval). A stage that fails on its way to the trainer falls back: the micro-batch is sent whole in the next step's first call and staging stays off for the rest of the run. compile and restore_checkpoint drop a stale stage. Measured on the RL trainer bench (AI-Hypercomputer/maxtext#5701) through the equivalent bench-driver switch, same pod, fix off vs on: 105.97 -> 98.31 s per optimizer step on 128 v7x chips. This orchestrator code also ran on TPU in a reduced rig: the real StandardRLProgram and TrainerWorker over gRPC, driving the team cell's MaxText trainer on Pathways with synthetic groups instead of rollouts, 3 micro-batches per step. The exact-order checks were clean, and the step-boundary idle time fell from 12.46 s to 0.02-0.03 s. Tests: rl_program_pipelined_boundary_test.py, rl_program_pipelined_boundary_robustness_test.py, rl_program_stage_before_update_test.py, distributed_rl_engine_staging_test.py, distributed_rl_engine_deferred_update_test.py, trainer_worker_staging_test.py, trainer_worker_deferred_update_test.py, trajectory_queue_manager_test.py, and the submission-launcher test (both flags off by default; STAGE_NEXT_MICROBATCH without PIPELINE_STEP_BOUNDARY is refused).

…e staged step boundary mlperf_gbs1024_submission.sh now sets, unless already set: - TUNIX_LOGPS_SHARDED_HEAD=local: the trainer computes log-probs with the sharded head (k8s_launcher.sh passes it to the trainer); 0 turns it off. - PIPELINE_STEP_BOUNDARY=1 and STAGE_NEXT_MICROBATCH=1: the orchestrator sends the next step's first micro-batch ahead of this step's bookkeeping and stages it on the trainer before this step's deferred update. Staging needs AI-Hypercomputer/maxtext#5718. STAGE_NEXT_MICROBATCH follows PIPELINE_STEP_BOUNDARY unless set, so PIPELINE_STEP_BOUNDARY=false alone turns both off. The other 397B v7x recipes keep all three off. The recipe comments name the code instead of pull request numbers. With these and the trainer settings of the first two commits, the trainer of the RL trainer bench (AI-Hypercomputer/maxtext#5701, 128 v7x chips) measured 98.31 s per optimizer step, against 130.80 s for the team recipe. Tests: the head reaches the trainer JobSet only and 0 turns it off; the staged boundary is on by default here and PIPELINE_STEP_BOUNDARY=false turns both flags off; the refusal test now sets PIPELINE_STEP_BOUNDARY=false; the other 397B v7x recipes keep the head and the boundary off.

…ext with #5759 TRAINER_TUNED_FLAGS (on in mlperf_gbs1024_submission.sh) now also sets, each measured on the 128-chip v7x RL trainer bench (AI-Hypercomputer/maxtext#5701 harness): - --xla_tpu_enable_sparse_core_collective_offload_nd_reduce_scatter=true: 97.24 -> 97.15 s/step. - --xla_jf_rematerialization_percent_shared_memory_limit=99: 10.226 -> 10.197 s per micro-batch. - moe_route=device: keeps the ragged sort's routing integers instead of replaying both argsorts in the backward (now tagged at the pin), 10.226 -> 10.184 s per micro-batch; +4.3 GiB compile temp, runs within the trainer's HBM. - wi_tile_dlhs_mlp_dim=2048 wi_tile_dlhs_embed_dim=2048: the wi input-gradient GEMM tiles the full contraction once. - REMAT_DECODER_LAYER_INPUT=device in the submission's tuned cell: the layer inputs stay in HBM, as in every bench run. MaxText pin: atwigg/mlperf plus AI-Hypercomputer/maxtext#5759 (GDN kernel, MoE routing and engine speedups) and #5741 (router and query kernels gathered once). With these flags the bench runs 9.796 s per micro-batch (89.77 s/step, 3,285 tok/s/chip) against 10.616 s (97.15 s/step) for atwigg/mlperf with #5741. Tests: mlperf_submission_launcher_test expects the new tuned keys and libtpu flags.

…ews (gdn_conv1d_dense_token_stack) The Pallas GDN forward kernel (kernels/gdn/wrapper.py fused_conv1d_gdn) convolves a chunk one token at a time and stacks the per-token [1, 1, dim] float32 rows with jnp.stack(axis=1) (kernels/gdn/compute_conv1d.py causal_conv1d). The stacked [1, chunk, 1, dim] value has a second-minor dimension of 1, and the TPU compiler of the October 2026 Pathways build (and of the libtpu 0.0.51 nightlies) lays it out in (8, 128) tiles padded from 1 to 8 sublanes. The SiLU applied next inherits that layout, so its exp and reciprocal run on 8x as many vector registers as the data fills. Older compilers keep the value dense. With gdn_conv1d_dense_token_stack=True each row is viewed as [dim / 1024, 8, 128], which is free because the row already fills dim / 1024 float32 registers; the views are concatenated along the leading axis and the result is reshaped to [1, chunk, 1, dim] (compute_conv1d.stack_conv1d_tokens). The same values in the same order; the compiler keeps the value dense on old and new builds alike. The option reaches the kernel through gdn_decoupled_conv1d's new static argument conv_dense_token_stack. Its primal and forward rule pass it to every forward-kernel call: the local call, both sequence-sharded context-parallel passes, and the forward rule's replay under remat (the backward rule never runs the forward kernel). fused_conv1d_gdn hands it to the per-sequence kernel only: its batched (decode) kernel, which convolves one token per sequence, keeps jnp.stack for those one-token rows, and the tokamax Op wrapper (kernels/gdn/pallas_mosaic_tpu.py) does not take the option. MaxText's GDN layer runs neither, since gdn_decoupled_conv1d calls fused_conv1d_gdn prefill-only. The option is left out of the kernel name and metadata, so outside the kernel body the lowered program is unchanged and profiles keep the kernel name. Numerics: every element lands where jnp.stack puts it, so the kernel's outputs do not change. On CPU, causal_conv1d and the SiLU after it give the same bits with and without the option (the kernel itself does not run on CPU). The TPU numerics A/B of the measured build below passed all six gates (nxs6gr). The validator requires a decoder with Gated DeltaNet layers (qwen3_next or qwen3_5), use_gdn_kernel=True, and a width that each forward-kernel call convolves (2 * key heads * key head dim + value heads * value head dim; one rank's heads under head-sharded GDN context parallelism) that is a multiple of 1024. Like the existing sequence-length check, it counts a context parallelism left at -1 (filled in when the mesh is built) as 1; fused_conv1d_gdn checks the per-call width again when the layer is traced, which catches such a config. Where the width is not a multiple of 1024 the measured build below silently kept jnp.stack; this change rejects the config instead. Default off. Measured on Qwen3.5-397B-A17B, 128 TPU v7x, on the RL trainer benchmark (#5701 driver, mesh fsdp 32 x context 4 x expert 2, 9 micro-batches per step, ragged_buffer_factor 2.75, staged step boundary), with the same change behind an environment switch (MAXTEXT_GDN_CONV_DENSE_STACK=1) on the stack-6 base: 85.731 s per step, regular micro-batch 9.3552 s (o275x-s6gr-l3-prof), against 86.926 s / 9.4910 s for the same-pod arm o275x-s6n2-l3-prof: -1.195 s per step. That arm's only change, a 2D GDN out_proj, is a no-op (XLA already canonicalises the 4D out_proj into the same 2D dot; TensorCore busy time -2 ms per micro-batch). The stack-6 control in another pod measured 86.970 s / 9.494 s (o275x-s6c-l3). Equivalence to the measured build, on CPU: lowered for a TPU7x target (no TPU needed), the GDN layer's forward and backward program with the key on is the measured build's program with its switch on, and with the key off is the base program (StableHLO and every Mosaic kernel body, compared without source locations). 18 configurations: no context parallelism, head-sharded context parallelism 2, sequence-sharded context parallelism 4 (also at the 397B's 16 key and 64 value heads), with and without packing and remat, and gdn_conv1d_input_in_qkv_dtype off and on. In each, the key changes only the forward-kernel bodies. Outside that layer the two differ: the measured build's switch also regrouped the batched kernel's one-token rows (at a decode tile of 1) and reached the tokamax Op wrapper. Tests (gdn_conv1d_dense_token_stack_test.py, new), on CPU: - stack_conv1d_tokens(dense=True) equals jnp.stack(axis=1) bit for bit, and dense=False is jnp.stack itself (the same jaxpr); causal_conv1d and the SiLU after it give the same bits with the option at widths 1024 and 12288, packed or not, with and without conv bias, for full and partial chunks; a token-minor regrouping is caught. - Traced and lowered for a TPU7x target on the CPU host, every forward-kernel call assembles its conv output densely exactly when the option is on: fused_conv1d_gdn; the custom VJP's primal, forward rule and remat replay; both sequence-sharded CP passes, under the forward rule and, forward only, under the primal; head-sharded CP; and model_runner reading the config key. The batched (decode) kernel is the same program with and without the option. The Mosaic lowering keeps the kernel name and changes only the kernel body. The option composes with gdn_conv1d_input_in_qkv_dtype, and the backward rule accepts every positional form. An option the conv ignores, one the forward rule drops and one the primal drops each fail every check they break. - The width check and the validator reject the option where it would do nothing, including one rank's width under head-sharded CP. The validator's copy of model_runner's head-vs-sequence CP rule agrees with model_runner at every GDN CP mode, degrees 1-8 on each context axis, nine head-count pairs and head dims 128 and 256. The kernel name and metadata do not change.

… join TRAINER_TUNED_FLAGS TRAINER_TUNED_FLAGS=true now selects the trainer cell that measured fastest on the RL trainer bench with every numerics gate passing. In mlperf_397b_256_v7x.sh the block of MaxText keys (part 2) gains sixteen keys and loses one: - wi_tile_dlhs_mlp_dim=2048, decoder_layer_input=device, moe_routing_maps=device and moe_gmm_v2_dlhs_transpose_rhs=true: the K tile of wi's input-gradient GMM matched to the fused gate and up weights, each decoder layer's input kept in HBM, the backward reusing the forward's MoE routing maps, and the input-gradient GMMs reading the expert weights through the kernel's transpose_rhs. - moe_routing_index_math=count, moe_unsplit_prefused_wi=true, moe_roe_packed_routing_gather=true, moe_router_grad_from_wo_input=true and moe_wi_gmm_fused_swiglu=true: the MoE keys of AI-Hypercomputer/maxtext#5758 in the set its commits recommend (the routing index maps by counting instead of argsorts, the fused wi kept whole up to the GEMM, one packed routing all-gather per layer, the routing-weight gradient from the down projection's input, the SwiGLU in the wi GEMM's epilogue). - gdn_conv1d_input_in_qkv_dtype=true and gdn_conv1d_dense_token_stack=true: the GDN Conv1D backward without its f32 copy of the input, and the forward kernel's Conv1D output built from dense per-token views. - wi_tile_fwd_batch_seq, wi_tile_drhs_batch_seq, wo_tile_fwd_batch_seq, wo_tile_dlhs_batch_seq and wo_tile_drhs_batch_seq=512 (from 256): row tiles of 512 for the five MoE GMMs with 1024 tiles. wi_tile_dlhs_batch_seq stays 256 next to its 2048 K tile. - sa_use_base2_exp=false leaves the block, so the trainer keeps the recipe's sa_use_base2_exp=true: base-2 exponentials measured 0.025 s faster per micro-batch (o275x-gx-b2e), and their numerics run passes (nxgx-b2e). In mlperf_gbs1024_submission.sh, which turns TRAINER_TUNED_FLAGS on, the switch also sets three defaults, each still overridable: - COMPUTE_LOGPS_CHUNK_SIZE=4096 (was 1024), and 1024 when TUNIX_LOGPS_SHARDED_HEAD turns the sharded head off (0, false, off or no, in any case): chunks of 4096 were measured with the head only, and without it the log-prob program failed to compile on TPU at chunks of 2048 and 4096. - TUNIX_LOGPS_SHARDED_HEAD=fused (was local). TRAINER_TUNED_FLAGS=false keeps local. - PACK_PAD_POSITION=start: each packed row puts its padding first. Bench (AI-Hypercomputer/maxtext#5701 driver, Qwen3.5-397B-A17B, 128 TPU v7x, mesh fsdp 32 x context 4 x expert 2, ragged_buffer_factor 2.75, 9 micro-batches per step, mean of steps 1 to 4, no token dropped in any arm): - before these keys, the fused head and pad first: 106.30 s per step (o275x-hd), 2,774 tokens/s/chip; - this cell: 82.671 and 82.670 s per step in two pods (o275x-c9pfxhv-l3), 3,567 tokens/s/chip. Part of the difference is the staged step boundary (AI-Hypercomputer/maxtext#5718), which this recipe already turns on and o275x-hd did not run; per regular micro-batch (the median of micro-batches 3 to 9), which the boundary does not change, about 10.77 -> 9.01 s; - within it, each against a control in the same pod: padding first 83.285 -> 82.043 s per step (-1.242 s) and the fused head 83.253 -> 82.708 s per step (-0.545 s). The recipe lists the measurement of each key. The measured build switched the MaxText changes with environment variables, and it also had one attention change on that neither MaxText PR carries (attention's load-balancing permute hoisted; alone 99.504 -> 99.267 s per step). The keys themselves, on the MaxText change's own code (2d82c5199) without gdn_conv1d_dense_token_stack, the fused head and pad first, measured 85.295 and 85.251 s per step (o275x-prc-l3-prof, o275x-prc-l3-r), within 0.23 s of what the environment-switch arms predicted. Both arms set moe_route=device, which nothing in #5698 or #5758 tags, so moe_route stays out of the block as before; and the stop of the frozen router's gradient, an environment switch in the measured build, is unconditional in #5758, with no key. Numerics (loss gates off, random advantages, three steps, against two reference runs whose gradient norms differ by 1.9e-5 relative): all six gates pass for this cell (nxc9pfxhv) and for the keys on the MaxText change's own code (nxprc). moe_wo_gmm_fused_dswiglu stays out: next to the fused SiLU GMM at these tiles it computed a wrong SwiGLU gradient, and its fix is still being verified on TPU. moe_routing_maps and moe_gmm_v2_dlhs_transpose_rhs exist in AI-Hypercomputer/maxtext#5698, and the seven other new MaxText keys in the MaxText change of AI-Hypercomputer/maxtext#5758 (on top of #5698); none of the nine is defined at this tree's pin (e526ea02), and MaxText refuses a key it does not define when the trainer starts. TRAINER_TUNED_FLAGS=true, the default of mlperf_gbs1024_submission.sh, therefore needs requirements/maxtext_requirements.txt moved to a MaxText with #5758 before an image built from this tree runs it. TRAINER_TUNED_FLAGS=false and the other 397B v7x recipes pass none of the nine. Tests (mlperf_submission_launcher_test.py): the exact-set test expects 23 keys, _MAXTEXT_DEFAULTS has the nine new keys' defaults, and the switch-off and other-recipe tests check that none of the nine is passed; the default launch has base-2 exponentials on, chunks of 4096, the fused head and --pack_pad_position=start; TRAINER_TUNED_FLAGS=false keeps the local head and the padding at the end; an explicit pad position, head or chunk size wins; TUNIX_LOGPS_SHARDED_HEAD=0, false, OFF or No gives chunks of 1024; the other 397B v7x recipes pass no pad position. Fifteen deliberately wrong copies of the recipes each fail at least one of these tests: a key dropped from the block or given the wrong value, one tile left out, wi_tile_dlhs_batch_seq raised to 512, sa_use_base2_exp=false kept, moe_wo_gmm_fused_dswiglu or moe_route=device added, a new key passed outside the block at its default, the local head under the switch, the chunk fallback removed or made case-sensitive, chunks of 1024 by default, pad first outside the switch, pad_align passed, and both recipes as they were before this change.

…ition) The packer writes a row's segments from position 0 and leaves its padding at the end. Under load-balanced context parallelism (MaxText's DUAL_CHUNK_SWAP: context rank r holds sequence chunks r and 2 * cp - 1 - r) the last chunk of a row is rank 0's, so trailing padding makes only rank 0 lighter, while ranks 1..cp-1 still attend every real query to every earlier position. tokamax's Splash (0.0.16.dev20261008 or later) skips the tiles whose query and key segment-id ranges do not overlap, so once the padding comes first every real query skips the padding's key tiles; an older tokamax masks segment ids element by element, and there "start" gains nothing. On the RL trainer benchmark's 60,948-token rows at context 4, counted with that rule and the benchmark's blocks (forward 1024 x 4096, dkv 2048 x 2048), the forward tiles per rank go from 80/136/136/136 to 125/120/120/120 and the dkv tiles from 74/132/132/132 to 122/116/116/116: the busiest rank, which sets the step time, does about 8% fewer. pad_position="start" (pack_chunk, pack_bin, pack_core) lays a row out exactly as "end" does, then writes it further right by its trailing padding, counted from its last token and rounded down to a multiple of pad_align (default 4096, the largest Splash block of the Qwen3.5-397B-A17B recipe), so every segment start keeps its place relative to the kernel's tiles and its segment alignment. Every per-token array moves with its token: ids, masks, advantages, segment ids, in-segment positions, per-token values and routed experts, so a model whose attention, convolutions and recurrences respect segment ids computes the same values. A row with no tokens, or with less trailing padding than pad_align, does not move, so the gain depends on how much padding the rows carry. "end" stays the default and packs byte for byte as before; pad_align is not read with "end". The option is configuration, threaded through the orchestrator the way segment_align_multiple is: BatchConfig.pack_pad_position / pack_pad_align -> create_batch_assembler -> SequencePackedBatchAssembler -> pack_chunk, so the program and its trainer warm-up payload pack alike. run_deepswe_dist.py takes --pack_pad_position / --pack_pad_align (main() builds its BatchConfig with _batch_config, so a test can check that the flags reach it), and k8s_launcher.sh passes them from PACK_PAD_POSITION / PACK_PAD_ALIGN only when they are set. Only the orchestrator is wired: pack_sequences in tunix/rl/utils.py, which rl_learner and agentic_rl_learner use, still puts the padding at the end. pack_chunk, pack_bin and pack_core raise on an unknown position and, with "start", on a pad_align that is not a positive multiple of segment_align_multiple; they leave a row that cannot move as it is. The orchestrator also refuses settings under which the option would do nothing, so a run cannot measure the default by accident: "start" with the padded (unpacked) assembler, a pad_align of at least the row length, or a non-default pad_align with "end". The overlay that the measurements below ran with read the switch from the environment inside packing.py (TUNIX_PACK_PAD_POSITION / TUNIX_PACK_PAD_ALIGN). Nothing reads those two variables any more: set PACK_PAD_POSITION / PACK_PAD_ALIGN, or the flags. Nothing in packing.py reads the environment and every other packer option is an argument, so this one is too. Measured on Qwen3.5-397B-A17B, 128 TPU v7x, on the RL trainer benchmark (AI-Hypercomputer/maxtext#5701 driver, mesh fsdp 32 x context 4 x expert 2, ragged_buffer_factor 2.75, 9 micro-batches per step), on stack 3. The benchmark builds its own rows and its driver's PAD_POSITION=start moves them; the overlay's switch changed only the trainer warm-up payload (same shapes and dtypes). So the numbers below measure the driver's layout, which is the layout this packer produces (checked field by field on the same rows). Every regular benchmark row carried 4,588 tokens of padding and moved by 4,096; rows packed from real rollouts may carry less padding and move less, or not at all. Regular micro-batch 10.033 -> 9.876 s and 99.50 -> 98.32 s per step, mean of steps 1-4 (o275x-s3-pf-prof against the stack 3 control o275x-gdn-ctl, run in another pod). Numerics (gates off, random advantages, against two reference runs whose gradient norms differ by 1.9e-5 relative; nxs3-pf): all six gates pass, loss 2.7e-7 relative, log-prob difference 7.0e-6, entropy 8.0e-8 relative, gradient norm 2.3e-4 relative (pass limit 2e-3). Tile boundaries move, so reductions round differently. Tests: - packing_test.py: "start" rows written out by hand (the shift, its rounding, segment alignment, routing and per-token values moving with their tokens, the default pad_align, and a shift counted from the last token when the budget is not a multiple of the segment alignment); 200 random packings in which "start" is the "end" row rotated right by its rounded trailing padding in every array (segment alignments 1 to 128, budgets that are and are not multiples of them, every per-token field, items without routing, empty items and bins); every token keeps its id, segment id, in-segment position, masks, advantage, per-token values and routing, checked against the items; "end" is unchanged by any pad_align, also at segment alignments that do not divide 4096; invalid settings raise. - batch_assembly_test.py: the assembler puts the padding first; the options reach it from BatchConfig on both packed paths and are refused on the padded one; settings that would do nothing raise. - grpo_recipe_wiring_test.py, mlperf_submission_launcher_test.py: the flags parse and reach BatchConfig; the default command line gives the previous BatchConfig field for field; rows packed from the parsed flags by the assembler put their padding first, at the default and at a given pad_align; main() hands the program the BatchConfig from _batch_config; the launcher passes the flags only when set, changing only the orchestrator's JobSet. Forty-two deliberately broken copies each fail at least one of these tests: an array left unshifted, given absolute positions or one token long; the shift unrounded, off by one, a float, blind to the segment alignment or rounded up to it; segment ids widened to int64; pad_align parsed or defaulted as a float; the option dropped, or taken from the wrong variable, at any hop from the launcher to pack_chunk, including main() building its own BatchConfig; a check removed, loosened or applied to "end"; "start" as the default or always passed.

…EAD=fused) With TUNIX_LOGPS_SHARDED_HEAD=local the sharded head scores each chunk with selective_log_softmax (and compute_entropy_from_logits), and its backward differentiates that with jax.vjp. At the 248,320-entry vocabulary of Qwen3.5-397B-A17B two parts of that cost more than the arithmetic. On TPU, the target logit's take_along_axis makes XLA write an f32 copy of the [chunk, vocab] logits for a SparseCore gather, and its transpose is a scatter-add: an f32 [chunk, vocab] cotangent, a flattened copy of it, a sort of the indices, a SparseCore scatter and a copy back to bf16, per chunk. And the logsumexp's max and exp-sum are computed again for the entropy in the forward, with one more exp pass, and once more in the backward. TUNIX_LOGPS_SHARDED_HEAD=fused runs the same sharded head with fused scoring: - The forward makes one pass over each chunk's logits (_head_logps_stats): the target logit by a one-hot select-sum, the max and exp-sum once, the entropy from the same exponentials, and the max and exp-sum kept for the backward (8 bytes per token). - When only the log-probs have a cotangent, as in GRPO, whose loss stops the entropy's gradient, the backward forms the logits' cotangent in one elementwise pass over the recomputed logits (_head_logps_cotangent): no f32 [chunk, vocab] intermediate, no scatter and no reductions. An optimization barrier makes both gradient matmuls read that one cotangent. - A loss that differentiates the entropy takes the "local" backward. Unset, 0 and local trace the same program as before. The option is a value of the switch that turns the head on, because that switch is the head's only setting: the head has no config field or argument, and k8s_launcher.sh already passes TUNIX_LOGPS_SHARDED_HEAD to the trainer as it is set. Numerics: the log-probs and the logits' cotangent are the f32 operations of "local" in the same order, but the exp-sum is reduced in another fusion, so results can differ in their last bits. On CPU (8 devices; a head with a 248,320-entry vocabulary called directly, and the test head through compute_per_token_logps; 12 cases), "fused" against "local": log-probs equal or 1 ulp apart, entropy 3 to 7 ulps apart, and gradients apart by at most 2.6e-8 of their largest entry. Measured on Qwen3.5-397B-A17B, 128 TPU v7x, on the RL trainer benchmark (AI-Hypercomputer/maxtext#5701 driver, staged step boundary, ragged buffer factor 2.75, log-prob chunks of 4,096, 9 micro-batches per step), against a control run in the same pod (o275x-c9hv-l3 against o275x-c9ctl-l3): 83.253 -> 82.708 s per optimizer step, mean of steps 1-4, and 9.084 -> 9.035 s per regular micro-batch. The trainer logged the head engaged with 16,384 local tokens per device in 4 steps of 4,096. Numerics on TPU, measured in the same benchmark together with two MaxText changes (deferred small all-reduces, and the shared expert applied first; nxc9pfxall against two reference runs whose gradient norms differ by 1.9e-5 relative): all six gates pass, loss 2.7e-7 relative, log-prob difference 3.1e-7 (the references' own: 3.6e-7), entropy 5.1e-7 relative, gradient norm 1.6e-4 relative (pass limit 2e-3). A deviceless TPU compile of the head alone, 16,384 tokens per device: temporaries 11.87 -> 4.35 GiB at chunks of 4,096 and 4.65 -> 2.75 GiB at chunks of 1,024, transcendentals 5.09e9 -> 2.03e9 and 1.27e9 -> 0.51e9, and no gather, sort or scatter left. The build measured on TPU switched this scoring with its own environment variable (PARAMBOLE_HEAD_VJP=1) in a copy of tunix/rl/common.py based on an earlier version of the head (md5 01eaca3f). On CPU, this file with "fused" lowers to the same StableHLO and gives bit-identical log-probs, entropy and gradients as that file with the variable set, and with "local" or unset the same as the previous commit, in all 12 cases (two meshes, f32 and bf16 logits, temperatures 1 and 0.7, padded chunks, the entropy off, returned or differentiated, linear and GRPO-shaped losses, forward only). Tests (tests/rl/sharded_logps_head_test.py): - ShardedLogpsHeadTest: "fused" against "local" through compute_per_token_logps with the log-probs alone differentiated (as in GRPO), with the entropy differentiated, with chunks that do not divide the local sequence and one chunk for all of it, at temperatures 0.7 and 1, with non-finite hidden states at masked positions, in bf16, with one output, forward only, and against the unsharded head; the program has no gather, no scatter and one optimization barrier when only the log-probs are differentiated, and the autodiff backward's scatter when the entropy is; "fused" parses and near misses raise. - FusedScoringTest (one device): _target_hit by hand, and _head_logps_stats and _head_logps_cotangent against selective_log_softmax, compute_entropy_from_logits and jax.vjp on logits with inf, -inf and NaN entries and a row of no finite logit, with negative ids and ids outside the vocabulary. Fifteen deliberately broken copies of tunix/rl/common.py each fail at least one of these tests, and each new test fails on at least one of them: the cotangent's sign flipped, the backward's temperature division dropped, the log-prob cotangent read from the first chunk, the max and exp-sum swapped, the max dropped from the log-probs, the entropy's sign flipped, an entropy cotangent sent to the one-pass backward, the barrier removed, the wrap of negative ids removed, the vocabulary bound off by one, the isfinite mask dropped from the cotangent, and the option ignored, always on, parsed as "local" or matched loosely.

…ition) The packer writes a row's segments from position 0 and leaves its padding at the end. Under load-balanced context parallelism (MaxText's DUAL_CHUNK_SWAP: context rank r holds sequence chunks r and 2 * cp - 1 - r) the last chunk of a row is rank 0's, so trailing padding makes only rank 0 lighter, while ranks 1..cp-1 still attend every real query to every earlier position. tokamax's Splash (0.0.16.dev20261008 or later) skips the tiles whose query and key segment-id ranges do not overlap, so once the padding comes first every real query skips the padding's key tiles; an older tokamax masks segment ids element by element, and there "start" gains nothing. On the RL trainer benchmark's 60,948-token rows at context 4, counted with that rule and the benchmark's blocks (forward 1024 x 4096, dkv 2048 x 2048), the forward tiles per rank go from 80/136/136/136 to 125/120/120/120 and the dkv tiles from 74/132/132/132 to 122/116/116/116: the busiest rank, which sets the step time, does about 8% fewer. pad_position="start" (pack_chunk, pack_bin, pack_core) lays a row out exactly as "end" does, then writes it further right by its trailing padding, counted from its last token and rounded down to a multiple of pad_align (default 4096, the largest Splash block of the Qwen3.5-397B-A17B recipe), so every segment start keeps its place relative to the kernel's tiles and its segment alignment. Every per-token array moves with its token: ids, masks, advantages, segment ids, in-segment positions, per-token values and routed experts, so a model whose attention, convolutions and recurrences respect segment ids computes the same values. A row with no tokens, or with less trailing padding than pad_align, does not move, so the gain depends on how much padding the rows carry. "end" stays the default and packs byte for byte as before; pad_align is not read with "end". The option is configuration, threaded through the orchestrator the way segment_align_multiple is: BatchConfig.pack_pad_position / pack_pad_align -> create_batch_assembler -> SequencePackedBatchAssembler -> pack_chunk, so the program and its trainer warm-up payload pack alike. run_deepswe_dist.py takes --pack_pad_position / --pack_pad_align (main() builds its BatchConfig with _batch_config, so a test can check that the flags reach it), and k8s_launcher.sh passes them from PACK_PAD_POSITION / PACK_PAD_ALIGN only when they are set. Only the orchestrator is wired: pack_sequences in tunix/rl/utils.py, which rl_learner and agentic_rl_learner use, still puts the padding at the end. pack_chunk, pack_bin and pack_core raise on an unknown position and, with "start", on a pad_align that is not a positive multiple of segment_align_multiple; they leave a row that cannot move as it is. The orchestrator also refuses settings under which the option would do nothing, so a run cannot measure the default by accident: "start" with the padded (unpacked) assembler, a pad_align of at least the row length, or a non-default pad_align with "end". The overlay that the measurements below ran with read the switch from the environment inside packing.py (TUNIX_PACK_PAD_POSITION / TUNIX_PACK_PAD_ALIGN). Nothing reads those two variables any more: set PACK_PAD_POSITION / PACK_PAD_ALIGN, or the flags. Nothing in packing.py reads the environment and every other packer option is an argument, so this one is too. Measured on Qwen3.5-397B-A17B, 128 TPU v7x, on the RL trainer benchmark (AI-Hypercomputer/maxtext#5701 driver, mesh fsdp 32 x context 4 x expert 2, ragged_buffer_factor 2.75, 9 micro-batches per step), on stack 3. The benchmark builds its own rows and its driver's PAD_POSITION=start moves them; the overlay's switch changed only the trainer warm-up payload (same shapes and dtypes). So the numbers below measure the driver's layout, which is the layout this packer produces (checked field by field on the same rows). Every regular benchmark row carried 4,588 tokens of padding and moved by 4,096; rows packed from real rollouts may carry less padding and move less, or not at all. Regular micro-batch 10.033 -> 9.876 s and 99.50 -> 98.32 s per step, mean of steps 1-4 (o275x-s3-pf-prof against the stack 3 control o275x-gdn-ctl, run in another pod). Numerics (gates off, random advantages, against two reference runs whose gradient norms differ by 1.9e-5 relative; nxs3-pf): all six gates pass, loss 2.7e-7 relative, log-prob difference 7.0e-6, entropy 8.0e-8 relative, gradient norm 2.3e-4 relative (pass limit 2e-3). Tile boundaries move, so reductions round differently. Tests: - packing_test.py: "start" rows written out by hand (the shift, its rounding, segment alignment, routing and per-token values moving with their tokens, the default pad_align, and a shift counted from the last token when the budget is not a multiple of the segment alignment); 200 random packings in which "start" is the "end" row rotated right by its rounded trailing padding in every array (segment alignments 1 to 128, budgets that are and are not multiples of them, every per-token field, items without routing, empty items and bins); every token keeps its id, segment id, in-segment position, masks, advantage, per-token values and routing, checked against the items; "end" is unchanged by any pad_align, also at segment alignments that do not divide 4096; invalid settings raise. - batch_assembly_test.py: the assembler puts the padding first; the options reach it from BatchConfig on both packed paths and are refused on the padded one; settings that would do nothing raise. - grpo_recipe_wiring_test.py, mlperf_submission_launcher_test.py: the flags parse and reach BatchConfig; the default command line gives the previous BatchConfig field for field; rows packed from the parsed flags by the assembler put their padding first, at the default and at a given pad_align; main() hands the program the BatchConfig from _batch_config; the launcher passes the flags only when set, changing only the orchestrator's JobSet. Forty-two deliberately broken copies each fail at least one of these tests: an array left unshifted, given absolute positions or one token long; the shift unrounded, off by one, a float, blind to the segment alignment or rounded up to it; segment ids widened to int64; pad_align parsed or defaulted as a float; the option dropped, or taken from the wrong variable, at any hop from the launcher to pack_chunk, including main() building its own BatchConfig; a check removed, loosened or applied to "end"; "start" as the default or always passed.

…EAD=fused) With TUNIX_LOGPS_SHARDED_HEAD=local the sharded head scores each chunk with selective_log_softmax (and compute_entropy_from_logits), and its backward differentiates that with jax.vjp. At the 248,320-entry vocabulary of Qwen3.5-397B-A17B two parts of that cost more than the arithmetic. On TPU, the target logit's take_along_axis makes XLA write an f32 copy of the [chunk, vocab] logits for a SparseCore gather, and its transpose is a scatter-add: an f32 [chunk, vocab] cotangent, a flattened copy of it, a sort of the indices, a SparseCore scatter and a copy back to bf16, per chunk. And the logsumexp's max and exp-sum are computed again for the entropy in the forward, with one more exp pass, and once more in the backward. TUNIX_LOGPS_SHARDED_HEAD=fused runs the same sharded head with fused scoring: - The forward makes one pass over each chunk's logits (_head_logps_stats): the target logit by a one-hot select-sum, the max and exp-sum once, the entropy from the same exponentials, and the max and exp-sum kept for the backward (8 bytes per token). - When only the log-probs have a cotangent, as in GRPO, whose loss stops the entropy's gradient, the backward forms the logits' cotangent in one elementwise pass over the recomputed logits (_head_logps_cotangent): no f32 [chunk, vocab] intermediate, no scatter and no reductions. An optimization barrier makes both gradient matmuls read that one cotangent. - A loss that differentiates the entropy takes the "local" backward. Unset, 0 and local trace the same program as before. The option is a value of the switch that turns the head on, because that switch is the head's only setting: the head has no config field or argument, and k8s_launcher.sh already passes TUNIX_LOGPS_SHARDED_HEAD to the trainer as it is set. Numerics: the log-probs and the logits' cotangent are the f32 operations of "local" in the same order, but the exp-sum is reduced in another fusion, so results can differ in their last bits. On CPU (8 devices; a head with a 248,320-entry vocabulary called directly, and the test head through compute_per_token_logps; 12 cases), "fused" against "local": log-probs equal or 1 ulp apart, entropy 3 to 7 ulps apart, and gradients apart by at most 2.6e-8 of their largest entry. Measured on Qwen3.5-397B-A17B, 128 TPU v7x, on the RL trainer benchmark (AI-Hypercomputer/maxtext#5701 driver, staged step boundary, ragged buffer factor 2.75, log-prob chunks of 4,096, 9 micro-batches per step), against a control run in the same pod (o275x-c9hv-l3 against o275x-c9ctl-l3): 83.253 -> 82.708 s per optimizer step, mean of steps 1-4, and 9.084 -> 9.035 s per regular micro-batch. The trainer logged the head engaged with 16,384 local tokens per device in 4 steps of 4,096. Numerics on TPU, measured in the same benchmark together with two MaxText changes (deferred small all-reduces, and the shared expert applied first; nxc9pfxall against two reference runs whose gradient norms differ by 1.9e-5 relative): all six gates pass, loss 2.7e-7 relative, log-prob difference 3.1e-7 (the references' own: 3.6e-7), entropy 5.1e-7 relative, gradient norm 1.6e-4 relative (pass limit 2e-3). A deviceless TPU compile of the head alone, 16,384 tokens per device: temporaries 11.87 -> 4.35 GiB at chunks of 4,096 and 4.65 -> 2.75 GiB at chunks of 1,024, transcendentals 5.09e9 -> 2.03e9 and 1.27e9 -> 0.51e9, and no gather, sort or scatter left. The build measured on TPU switched this scoring with its own environment variable (PARAMBOLE_HEAD_VJP=1) in a copy of tunix/rl/common.py based on an earlier version of the head (md5 01eaca3f). On CPU, this file with "fused" lowers to the same StableHLO and gives bit-identical log-probs, entropy and gradients as that file with the variable set, and with "local" or unset the same as the previous commit, in all 12 cases (two meshes, f32 and bf16 logits, temperatures 1 and 0.7, padded chunks, the entropy off, returned or differentiated, linear and GRPO-shaped losses, forward only). Tests (tests/rl/sharded_logps_head_test.py): - ShardedLogpsHeadTest: "fused" against "local" through compute_per_token_logps with the log-probs alone differentiated (as in GRPO), with the entropy differentiated, with chunks that do not divide the local sequence and one chunk for all of it, at temperatures 0.7 and 1, with non-finite hidden states at masked positions, in bf16, with one output, forward only, and against the unsharded head; the program has no gather, no scatter and one optimization barrier when only the log-probs are differentiated, and the autodiff backward's scatter when the entropy is; "fused" parses and near misses raise. - FusedScoringTest (one device): _target_hit by hand, and _head_logps_stats and _head_logps_cotangent against selective_log_softmax, compute_entropy_from_logits and jax.vjp on logits with inf, -inf and NaN entries and a row of no finite logit, with negative ids and ids outside the vocabulary. Fifteen deliberately broken copies of tunix/rl/common.py each fail at least one of these tests, and each new test fails on at least one of them: the cotangent's sign flipped, the backward's temperature division dropped, the log-prob cotangent read from the first chunk, the max and exp-sum swapped, the max dropped from the log-probs, the entropy's sign flipped, an entropy cotangent sent to the one-pass backward, the barrier removed, the wrap of negative ids removed, the vocabulary bound off by one, the isfinite mask dropped from the cotangent, and the option ignored, always on, parsed as "local" or matched loosely.

… TRAINER_STEP_TIMING, fix lint - bench_driver: finish_step() calls update() only when the last fwd_bwd did not already apply the optimizer (google/tunix#2719 reports `updated` in the response metadata); log_step_timing() logs the worker's get_step_timing() in the orchestrator's TRAINER_STEP_TIMING format. Both have CPU tests in tests/step_test.py. - README: describe both. - The DRY_RUN fixture interleaves launcher log lines with YAML documents, so it is now launch_dryrun_fixture.txt (yamllint no longer parses it as YAML). - pre-commit clean (codespell, pylint, pyink): docstrings, explicit encodings, dict literals, f-strings, unused imports; the deliberate function-level imports are documented with a module-level disable.

github-actions · 2026-10-30T01:10:07Z

🤖 Hi @ultrons, I'm running an automatic MaxText review of this pull request. You can track my progress in the logs. To request another review later, add the gemini-review label or comment @gemini-cli /review.

gemini-code-assist · 2026-10-30T01:10:36Z

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 a standalone benchmarking harness for the RL trainer program used in Tunix GRPO runs. By isolating the trainer on a single slice with a synthetic, calibrated workload, this harness enables efficient A/B testing and profiling of MaxText and libtpu configurations without the overhead of the full RL system. The PR includes comprehensive tooling for rendering benchmark JobSets, validating numerics, and performing static memory checks.

Highlights

New RL Trainer Benchmark Harness: Introduced a standalone benchmarking harness under experimental/rl_trainer_bench that allows for isolated profiling of the RL trainer program on a single slice.
Benchmark Configuration Tools: Added scripts for rendering benchmark JobSets from e2e dry runs, including support for A/B testing MaxText and libtpu flags.
Validation and Diagnostics: Included utilities for ahead-of-time (AOT) memory checking, numerics validation, and collective overlap analysis to ensure benchmark accuracy.

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.

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. ↩

github-actions

## 📋 Review Summary

This pull request introduces src/maxtext/experimental/rl_trainer_bench/, a standalone benchmarking and verification harness for the RL trainer component in Tunix GRPO pipelines without requiring live rollout engines or orchestrators. The implementation is comprehensive, well-structured, thoroughly documented in the accompanying README, and accompanied by CPU-runnable unit tests.

🔍 General Feedback

Architecture & Design: The harness cleanly isolates the RL training step, using realistic synthetic payloads calibrated to production MoE routing distributions while maintaining parity with e2e runtime configurations.
Robustness: A few edge-case refinements are suggested inline to prevent unexpected failures when handling empty/missing CLI flags, heredoc generation, or timing serialization.
Previous Review Follow-up: Note that the existing comment on bench_driver.py regarding np.arange(64, real_len, 64) raising a ValueError for small sequences/populations when replace=False remains relevant for smaller test shapes or custom batch configs.

github-actions · 2026-10-30T01:19:38Z

+  if not hasattr(worker, "get_step_timing"):
+    return
+  try:
+    timing = json.dumps(worker.get_step_timing(), separators=(",", ":"))


🟢 If `worker.get_step_timing()` returns `None`, `json.dumps(None)` serializes to the literal string `"null"`, producing a spurious `timing=null` log line. Checking if the returned timing is `None` before serialization prevents emitting empty timing lines.

Suggested change

timing = json.dumps(worker.get_step_timing(), separators=(",", ":"))

raw_timing = worker.get_step_timing()

if raw_timing is None:

return

timing = json.dumps(raw_timing, separators=(",", ":"))

github-actions · 2026-10-30T01:19:38Z

+  trainer_flags = set_flag(trainer_flags, "maxtext_output_directory", f"{a.output_prefix}/{a.name}")
+  trainer_flags = set_flag(trainer_flags, "worker_id", a.name)
+  if a.compute_logps_chunk_size:
+    trainer_flags = set_in_place(trainer_flags, [f"compute_logps_chunk_size={int(a.compute_logps_chunk_size)}"])


🟡 If the input launch render did not explicitly specify `--compute_logps_chunk_size`, `set_in_place` will raise a `ValueError` because the flag does not exist in `trainer_flags`. Using `upsert_flag` (consistent with how `profiler_steps` and `profiler_skip` are handled below) allows setting this flag even when it wasn't present in the baseline argv.

Suggested change

trainer_flags = set_in_place(trainer_flags, [f"compute_logps_chunk_size={int(a.compute_logps_chunk_size)}"])

trainer_flags = upsert_flag(trainer_flags, "compute_logps_chunk_size", str(int(a.compute_logps_chunk_size)))

github-actions · 2026-10-30T01:19:38Z

+  costs = [ast.literal_eval(m) for m in re.findall(r"Cost analysis: (\{.*\})", s)]
+  mems = re.findall(r"Memory analysis: CompiledMemoryStats\(([^)]*)\)", s)
+  for name, c, m in zip(("fwd_bwd", "accumulate", "update"), costs, mems):
+    d = dict(kv.split("=") for kv in m.split(", "))


🟢 Using `split("=")` will raise a `ValueError: too many values to unpack` if any property value in `CompiledMemoryStats` contains an equals sign. Using `split("=", 1)` safely limits splitting to the key-value separator.

Suggested change

d = dict(kv.split("=") for kv in m.split(", "))

d = dict(kv.split("=", 1) for kv in m.split(", "))

github-actions · 2026-10-30T01:19:39Z

+  return f"""# PID 1 wrapper: forwards SIGTERM to the benchmark so the grace period shuts it down cleanly.
+echo "Program Start: $(date)"
+cat > /tmp/bench_driver.py <<'BENCH_DRIVER_EOF'
+{driver_src}BENCH_DRIVER_EOF


🟢 If `driver_src` lacks a trailing newline, `BENCH_DRIVER_EOF` will concatenate directly onto the last line of the script instead of appearing on its own line, resulting in an unclosed heredoc syntax error in bash. Splitting it with an explicit newline ensures the delimiter is always on a clean line.

Suggested change

{driver_src}BENCH_DRIVER_EOF

{driver_src.rstrip('\n')}

BENCH_DRIVER_EOF

github-actions · 2026-10-30T01:19:39Z

+      "compile_engine_loss=grpo",
+      f"compile_engine_grpo_config={grpo}",
+      # Match the trainer's --compute_logps_chunk_size.
+      f"compile_engine_logps_chunk_size={int(os.environ.get('AOT_LOGPS_CHUNK_SIZE', '512'))}",


🟢 If `AOT_LOGPS_CHUNK_SIZE` is exported as an empty string in the environment, `int("")` raises a `ValueError`. Using `or '512'` provides a fallback when the variable is set but empty.

Suggested change

f"compile_engine_logps_chunk_size={int(os.environ.get('AOT_LOGPS_CHUNK_SIZE', '512'))}",

f"compile_engine_logps_chunk_size={int(os.environ.get('AOT_LOGPS_CHUNK_SIZE') or '512')}",

codecov · 2026-10-30T01:26:49Z

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

ultrons added 2 commits October 13, 2026 18:37

rl_trainer_bench: label example numbers, fail fast in aot_e2e.sh, not…

d6db22f

…e DRY_RUN context switch

gemini-code-assist Bot reviewed Oct 13, 2026

View reviewed changes

parambole added a commit that referenced this pull request Oct 14, 2026

Revert "Integration: merge #5701 sivaibhav/rl-trainer-bench (70d6444)"

d9848b5

This reverts commit 0968476, reversing changes made to 5f1e9ad.

parambole mentioned this pull request Oct 15, 2026

Qwen3.5-397B RL trainer on 128 v7x: MoE and GDN speed options as config keys #5758

Open

khatwanimohit mentioned this pull request Oct 15, 2026

397B v7x RL trainer: GDN kernel, MoE routing and engine speedups (10.62 -> 9.80 s per micro-batch on the 128-chip bench) #5759

Closed

parambole mentioned this pull request Oct 15, 2026

Qwen3.5-397B RL trainer on 128 v7x: pad-first sequence packing, fused sharded log-prob head VJP, and the tuned trainer recipe google/tunix#2913

Open

ultrons marked this pull request as ready for review October 30, 2026 01:09

ultrons requested review from RissyRan, abhinavclemson, bvandermoon, gagika, gobbleturk, khatwanimohit, richjames0, shralex and vipannalla as code owners October 30, 2026 01:10

ultrons requested review from A9isha, NuojCheng, SurbhiJainUSC, aireenmei, darisoy, dipannita08, hengtaoguo, huytransformer, igorts-git, jiangjy1982, shuningjin and xibinliu as code owners October 30, 2026 01:10

github-actions Bot reviewed Oct 30, 2026

View reviewed changes

igorts-git approved these changes Oct 30, 2026

View reviewed changes

A9isha approved these changes Oct 30, 2026

View reviewed changes

A9isha merged commit c847919 into AI-Hypercomputer:atwigg/mlperf Oct 30, 2026
46 of 75 checks passed

-    cuts = np.sort(rng.choice(np.arange(64, real_len, 64), size=n_seg - 1, replace=False)) if n_seg > 1 else []
+    pop = np.arange(64, real_len, 64)
+    if len(pop) >= n_seg - 1:
+      cuts = np.sort(rng.choice(pop, size=n_seg - 1, replace=False)) if n_seg > 1 else []
+    else:
+      cuts = np.sort(rng.choice(np.arange(1, real_len), size=n_seg - 1, replace=False)) if n_seg > 1 and real_len > n_seg else np.linspace(0, real_len, n_seg + 1, dtype=int)[1:-1]

	with open(path) as f:
	with open(path, encoding="utf-8") as f:

	with open(a.driver) as f:
	with open(a.driver, encoding="utf-8") as f:

	with open(a.out, "w") as f:
	with open(a.out, "w", encoding="utf-8") as f:



		def summarize(path, cap_gib):
		with open(path, errors="replace") as f:

	with open(path, errors="replace") as f:
	with open(path, encoding="utf-8", errors="replace") as f:

-    timing = json.dumps(worker.get_step_timing(), separators=(",", ":"))
+    raw_timing = worker.get_step_timing()
+    if raw_timing is None:
+      return
+    timing = json.dumps(raw_timing, separators=(",", ":"))

	trainer_flags = set_in_place(trainer_flags, [f"compute_logps_chunk_size={int(a.compute_logps_chunk_size)}"])
	trainer_flags = upsert_flag(trainer_flags, "compute_logps_chunk_size", str(int(a.compute_logps_chunk_size)))

	d = dict(kv.split("=") for kv in m.split(", "))
	d = dict(kv.split("=", 1) for kv in m.split(", "))

	{driver_src}BENCH_DRIVER_EOF
	{driver_src.rstrip('\n')}
	BENCH_DRIVER_EOF

	f"compile_engine_logps_chunk_size={int(os.environ.get('AOT_LOGPS_CHUNK_SIZE', '512'))}",
	f"compile_engine_logps_chunk_size={int(os.environ.get('AOT_LOGPS_CHUNK_SIZE') or '512')}",

Conversation

ultrons commented Oct 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

How it works

Placement

Caveats

Tests

Checklist

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Oct 13, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Oct 13, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Oct 13, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Oct 13, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Oct 13, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Oct 13, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Oct 13, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Oct 13, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Oct 13, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented Oct 30, 2026

Uh oh!

gemini-code-assist Bot commented Oct 30, 2026

Summary of Changes

Highlights

Footnotes

Uh oh!

github-actions Bot left a comment

Choose a reason for hiding this comment

🔍 General Feedback

Uh oh!

github-actions Bot Oct 30, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Oct 30, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Oct 30, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Oct 30, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Oct 30, 2026

Choose a reason for hiding this comment

Uh oh!

codecov Bot commented Oct 30, 2026

Codecov Report

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

ultrons commented Oct 13, 2026 •

edited

Loading