Repository navigation
Conversation
…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.
… 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.
parambole
changed the title
parambole
marked this pull request as ready for review
October 15, 2026 17:42
parambole
requested review from
abheesht17,
hgao327,
jiangyangmu,
lc5211,
s-noghabi,
sizhit2,
tianshub and
wang2yn84
as code owners
October 15, 2026 17:42
Contributor
Summary of ChangesHello, I'm Gemini Code Assist1! 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 PR introduces performance optimizations for the Qwen3.5-397B RL trainer on TPU v7x infrastructure. By implementing pad-first sequence packing and a fused sharded log-prob head, the trainer achieves significant speed improvements while maintaining numerical stability. A new tuned trainer recipe is provided to easily enable these optimizations. Highlights
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 AssistThe 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
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 Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request introduces a fused sharded log-prob head (TUNIX_LOGPS_SHARDED_HEAD=fused) to optimize chunk scoring and backward passes, alongside a new "start" padding position option in sequence packing to balance attention workloads under context parallelism. The changes are accompanied by extensive unit tests and recipe updates. The reviewer feedback suggests adding strict type annotations to the shape parameter of _target_hit and a return type annotation to _head_logps_stats in tunix/rl/common.py to comply with the repository's strict typing guidelines.
| # and differs in its last bits. | ||
|
|
||
|
|
||
| def _target_hit(ids: jax.Array, shape) -> tuple[jax.Array, jax.Array]: |
Contributor
There was a problem hiding this comment.
Please add strict type annotations to the shape parameter of _target_hit to adhere to the repository style guide's typing requirements.
Suggested change
| def _target_hit(ids: jax.Array, shape) -> tuple[jax.Array, jax.Array]: | |
| def _target_hit(ids: jax.Array, shape: tuple[int, ...]) -> tuple[jax.Array, jax.Array]: |
References
- Always annotate function signatures, attributes, and data structures with concrete, specific types rather than Any or untyped dict/object. (link)
| return hit, (ids >= 0) & (ids < vocab) | ||
|
|
||
|
|
||
| def _head_logps_stats(logits: jax.Array, ids: jax.Array, with_entropy: bool): |
Contributor
There was a problem hiding this comment.
Please add a strict return type annotation to _head_logps_stats to adhere to the repository style guide's typing requirements.
Suggested change
| def _head_logps_stats(logits: jax.Array, ids: jax.Array, with_entropy: bool): | |
| def _head_logps_stats(logits: jax.Array, ids: jax.Array, with_entropy: bool) -> tuple[jax.Array, jax.Array | None, jax.Array, jax.Array]: |
References
- Always annotate function signatures, attributes, and data structures with concrete, specific types rather than Any or untyped dict/object. (link)
s-noghabi
reviewed
Oct 15, 2026
Collaborator
There was a problem hiding this comment.
have we verified the correctness impact of this? I am a little concerned it might introduce convergence regression, but nothing pops out right away. ideally we would do a real run to compare. especially less confident on the bwd pass.
s-noghabi
reviewed
Oct 15, 2026
Collaborator
There was a problem hiding this comment.
also some findings from my AI agent:
4 subtle edge cases / hazards identified:
- Next-token shift boundary bleed on completion-only (len(prompt_ids) == 0) items: Shifting segment 1 to offset = shift > 0 causes its first token (t = shift, where completion_mask == 1 when p == 0) to be scored from hidden_states[:, shift - 1] (a leading segment_ids == 0 padding token) instead of 0.0, corrupting log-probs and leaking gradients into padding.
- Catastrophic MoE token dropping if pack_pad_position="start" is used without moe_exclude_padding_from_routing=true: All shift leading padding tokens (0..shift-1) have identical embeddings, route to the same top_k experts, and occupy the first buffer slots in stable ragged sort—evicting real tokens when expert capacity is finite.
- Split segment_ids == 0 across start, 64-token alignment gaps, and end in Splash Attention: Because Splash Attention masks via q_ids == kv_ids, trailing pad queries ([shift + used, budget)) attend to leading pad keys ([0, shift)) and inter-segment gap keys; moreover, if tokamax's [min, max] tile-skipping does not ignore 0, any 64-token alignment gap (segment_ids == 0) inside a Q tile lowers q_min to 0 and defeats leading-pad KV tile skipping.
- Default-on TRAINER_TUNED_FLAGS=true in
mlperf_gbs1024_submission.sh
passes 9 MaxText config keys that do not exist at the MaxText commit (60425bfbd) currently pinned in
requirements/maxtext_requirements.txt
.
s-noghabi
approved these changes
Oct 15, 2026
Rohan-Bierneni
added a commit
to Rohan-Bierneni/tunix
that referenced
this pull request
Oct 15, 2026
…oundary masking, and add 256-chip 64x1x8 recipe
Rohan-Bierneni
added a commit
to Rohan-Bierneni/tunix
that referenced
this pull request
Oct 15, 2026
…oundary masking, and add 256-chip 64x1x8 recipe
Rohan-Bierneni
added a commit
to Rohan-Bierneni/tunix
that referenced
this pull request
Oct 15, 2026
…oundary masking, and add 256-chip 64x1x8 recipe
…shes 32x2x8 and 64x1x8 With TRAINER_TPU_SLICE=tpu7x:4x8x8 and TRAINER_MESH_FSDP/EXPERT/CONTEXT 32/2/8 or 64/1/8, TRAINER_TUNED_FLAGS=true adds query_proj, key_proj, value_proj, out_proj and qkv_proj=device, the settings of Rohan Bierneni's 256-chip trainer runs. At 64x1x8, NUM_MOE_TOKEN_CHUNKS defaults to 1; an explicit value wins. One MoE token chunk now turns moe_chunk_pipeline off, and with it the tuned moe_roe_packed_routing_gather: MaxText refuses the chunk pipeline with one chunk, and the packed routing gather without the pipeline. Every existing render is unchanged.
Rohan-Bierneni
pushed a commit
to Rohan-Bierneni/tunix
that referenced
this pull request
Oct 16, 2026
…shes 32x2x8 and 64x1x8 Cherry-picked from google#2913 (8c75ae3) onto this PR's 256-chip recipe. With TRAINER_TPU_SLICE=tpu7x:4x8x8 and TRAINER_MESH_FSDP/EXPERT/CONTEXT 32/2/8 or 64/1/8, TRAINER_TUNED_FLAGS=true adds query_proj, key_proj, value_proj, out_proj and qkv_proj=device. At 64x1x8 (the default 4x8x8 mesh with tuned flags), NUM_MOE_TOKEN_CHUNKS defaults to 1; an explicit value wins. One MoE token chunk turns moe_chunk_pipeline off, and with it the tuned moe_roe_packed_routing_gather: MaxText refuses the chunk pipeline with one chunk, and the packed routing gather without the pipeline. This replaces the previous 4x8x8 block, which forced moe_chunk_pipeline=false at every 256-chip mesh while leaving moe_roe_packed_routing_gather=true (rejected by MaxText config validation) and only added the projection keys at EP=1. Measured on the RL trainer bench (9 micro-batches of 64 rows, steps 1-4 average): 32x2x8 60.62 s/step, 64x1x8 55.97 s/step, vs 82.17 s/step for 32x2x4 on 128 chips.
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Companion to AI-Hypercomputer/maxtext#5758. Stacked on
atwigg/mlperf(affb1c643, after #2878).Two opt-in speed options for the Qwen3.5-397B-A17B RL trainer, plus the recipe that turns on the fastest trainer cell measured without the follow-up key
moe_wo_gmm_fused_dswiglu(see Not in this PR), with every numerics gate passing. On the RL trainer bench (AI-Hypercomputer/maxtext#5701, 128 TPU v7x), this recipe on both PR heads (#575820c877aa1and this PR'sbee63532f), with no environment switches, takes 82.18 s per step instead of 106.30 s: 2,774 -> 3,589 tok/s/chip (two runs, 82.200 and 82.158 s; row 4 of Results). The six-gate numerics A/B passes for an overlay build of the same cell (82.67 s = 3,567 tok/s/chip, row 3) and for #5758's pre-review code (row 2); it has not yet run on the exact head build. The 106.30 s cell lacked the already-merged staged step boundary (AI-Hypercomputer/maxtext#5718); per regular micro-batch, which the boundary does not change, the time goes from 10.77 to 8.95 s (row 4). The 82.67 s cell ran the MaxText changes as an overlay build (environment switches) on the bench image's MaxText, without the merged AI-Hypercomputer/maxtext#5568, with one attention change that follows in a separate PR, and withmoe_route=device, which the recipe leaves at its default (see Results). Both options are off by default in the library: with them unset, the packer and the sharded log-prob head run the same program as before. The recipe commit turns both on inmlperf_gbs1024_submission.sh, whoseTRAINER_TUNED_FLAGSdefaults to true.What changes
Blue: changed by this PR, labelled with its option. Gray: unchanged.
flowchart TB classDef pr fill:#dbeafe,stroke:#1d4ed8,stroke-width:2px,color:#0b1f4d classDef same fill:#f3f4f6,stroke:#9ca3af,color:#6b7280 REC["launch recipe<br/>TRAINER_TUNED_FLAGS=true"]:::pr subgraph STEP["trainer step"] direction LR PACK["sequence packer<br/>pack_pad_position=start"]:::pr --> MT["MaxText model<br/>keys set by the recipe"]:::same --> HEAD["sharded log-prob head<br/>TUNIX_LOGPS_SHARDED_HEAD=fused"]:::pr --> LOSS["GRPO loss"]:::same end REC -. sets .-> STEP style STEP fill:none,stroke:#9ca3afChanges
pack_pad_position=startpack_sequencesintunix/rl/utils.pystill pads at the end), aligned to 4096 tokens, so every load-balanced context rank skips the padding's key tiles in Splash Attention (tokamax 0.0.16.dev20261008 or later; an older tokamax masks segment ids element by element and gains nothing)nxs3-pf)TUNIX_LOGPS_SHARDED_HEAD=fusednxc9pfxhv,nxc9pfxall)TRAINER_TUNED_FLAGS=true(recipe)nxc9pfxhv)* Stack 9 is the recipe's MaxText side on the overlay build, without pad first and the fused head, plus
moe_wo_gmm_fused_dswiglu, which is in neither PR because it computed a wrong gradient next to the fused SiLU, so these two arms are speed measurements only. Stack 3 is an earlier bench stack, the pad-first control. On the recipe stack, across pods, the fused head measured 83.316 -> 82.671 s per step.Results
RL trainer bench, 128 TPU v7x. Seconds per optimizer step, mean of steps 1-4.
o275x-hd)o275x-prc-l3-profand its repeato275x-prc-l3-r, one pod)nxprc)o275x-c9pfxhv-l3, two pods)nxc9pfxhv)o275x-prf-l3-profand its repeato275x-prf-l3-r, one pod)gdn_conv1d_dense_token_stack, pad first and the fused head.o275x-c9ypfxhv-l3), took 82.945 s per step (3,556 tok/s/chip; regular micro-batch 9.04 s). That pod's overlay build also contained the follow-up dSwiGLU fix, with its key off.moe_route=device(see Not in this PR).20c877aa1, this PR'sbee63532f) installed in the bench image and run with the recipe's settings as config, with no environment switches: the recipe's 16 MaxText keys, the fused log-prob head at chunks of 4096, padding first (the bench driver's equivalent ofpack_pad_position=start),ragged_buffer_factor2.75 with 0 dropped tokens, and the same Pathways server image as row 3. It does not have the attention hoist that row 3 had, and it does have the merged perf(training_engine): buffer microbatch metrics in Python lists and gate checkpoint metrics on intra-step state AI-Hypercomputer/maxtext#5568, which row 3 lacked. The six-gate numerics A/B has not been run on this build; rows 2 and 3 pass it.How to turn it on
Needs a MaxText with #5758. The recipe passes nine MaxText keys that this tree's pin (
e526ea02) does not define:moe_routing_mapsandmoe_gmm_v2_dlhs_transpose_rhsfrom AI-Hypercomputer/maxtext#5698, and seven of #5758's nine keys (its recommended set). MaxText refuses an undefined key when the trainer starts, sorequirements/maxtext_requirements.txtmust move to a MaxText with #5758 before an image built from this branch runsTRAINER_TUNED_FLAGS=true.TRAINER_TUNED_FLAGS=falseand the other 397B v7x recipes pass none of the nine.Testing
tests/rl/packing_test.py,tests/experimental/orchestrator/batch_assembly_test.py,tests/experimental/orchestrator/grpo_recipe_wiring_test.py,tests/experimental/examples/deepswe_dist/mlperf_submission_launcher_test.py.tests/rl/sharded_logps_head_test.py(ShardedLogpsHeadTest,FusedScoringTest), shown to fail on 15 deliberately wrong copies. Its sharded cases run in an 8-device child process: 40 tests OK on base, 62 on this branch.mlperf_submission_launcher_test.py41 of 41 pass (38 of 38 onaffb1c643); 15 deliberately wrong copies of the recipes each fail at least one test.o275x-c9pfxhv-l3ran on 24 of 25 MaxText settings: 17 of the 18 keys that arm passed (the 18th ismoe_route, left out on purpose) and the seven #5758 keys, which that arm set through the overlay build's environment switches. The chunk size 4096, the fused head (in that armlocalplus the overlay's head switch), pad first (in that arm the bench driver'sPAD_POSITION=start, the same layout as--pack_pad_position=start, checked field by field) and the SparseCore copy flag also match.nxc9pfxhv), #5758's keys on #5758's pre-review code (nxprc), pad first on stack 3 (nxs3-pf) and the recipe plus two levers in neither PR (nxc9pfxall). These gates compare scalars of the first of three steps (the later two are logged, not gated), so they cannot see a wrong gradient direction: they passed the wrongmoe_wo_gmm_fused_dswiglugradient, which was found by reading the source.Not in this PR (follow-up PRs)
moe_wo_gmm_fused_dswiglu: not in the recipe.nxc9ypfckhv, PASS 6/6). These gates also passed the unfixed version, so the kernel check is the evidence for the fix.moe_route=device: the measured arms set it, but no code in #5698 or #5758 tags a tensor with that name, so it does nothing. The recipe leaves it at its default.The recipe commit: every setting it changes
TRAINER_TUNED_FLAGS=falseand the other 397B v7x recipes render exactly as before.mlperf_397b_256_v7x.sh(MaxText keys)wi_tile_dlhs_mlp_dimdecoder_layer_inputmoe_routing_maps(#5698)moe_gmm_v2_dlhs_transpose_rhs(#5698)moe_routing_index_math(#5758)moe_unsplit_prefused_wi(#5758)moe_roe_packed_routing_gather(#5758)moe_router_grad_from_wo_input(#5758)moe_wi_gmm_fused_swiglu(#5758)gdn_conv1d_input_in_qkv_dtype(#5758)gdn_conv1d_dense_token_stack(#5758)wi_tile_{fwd,drhs}_batch_seq,wo_tile_{fwd,dlhs,drhs}_batch_seqwi_tile_dlhs_batch_seqstays 256 next to its 2048 K tile)sa_use_base2_expmlperf_gbs1024_submission.shCOMPUTE_LOGPS_CHUNK_SIZETUNIX_LOGPS_SHARDED_HEADturns the head offTUNIX_LOGPS_SHARDED_HEADTRAINER_TUNED_FLAGS=falsekeeps local)PACK_PAD_POSITIONThe recipe's comment lists the measurement of every key.
Why padding first is faster
Each packed row of 65,536 tokens is split over 4 context ranks with load balancing (dual chunk swap): 8 chunks, and rank r holds chunks r and 7-r, so causal-attention work is even across ranks.
pack_pad_position=start): every real query would attend to the padding keys before it, but they belong to another segment, so Splash Attention (tokamax 0.0.16.dev20261008 or later) skips those tiles. By tile count, ranks 1-3, which set the pace, drop from 136 to 120 forward tiles; rank 0 rises from 80 to 125, so the busiest rank does about 8% fewer (136 -> 125).The padding moves by a multiple of
pack_pad_align(default 4096).How the fused log-prob head works
TUNIX_LOGPS_SHARDED_HEAD=fusedruns the same sharded head aslocal, with fused scoring of each chunk of hidden states:take_along_axisgather; each token's log-prob is target - (log exp-sum + max), with no combine across devices.localand1keep the previous head unchanged.Measurement, logs and profiles
ragged_buffer_factor2.75, with 0 dropped tokens in every arm.gdn_conv1d_dense_token_stack+ pad first, without the fused head (o275x-c9pfxctl-l3): 83.316 s per step, regular micro-batch 9.085 s, 3,540 tok/s/chip; numerics not run alone, but its supersetsnxc9pfxhv(the recipe) andnxc9pfxallpass 6/6.o275x-s9p-l3-profvso275x-s9pf-l3-prof), fused head 83.253 -> 82.708 s per step (o275x-c9ctl-l3vso275x-c9hv-l3).gs://parambole-397b-replay-nap/win-tpu/sched/<booking>/<job>/<arm>.log, for example:r-5f2bf/parambole-bvx2-5f2bf/o275x-c9pfxhv-l3.log, replicater-5f2bf/parambole-bvx4-5f2bf/o275x-c9pfxhv-l3-r.log, third podr-2d4c2/parambole-bn3s-2d4c2/o275x-c9ypfxhv-l3.log; numericsr-5f2bf/parambole-bvx3-5f2bf/nxc9pfxhv.logr-5f2bf/parambole-bpr-5f2bf/o275x-prc-l3-prof.logando275x-prc-l3-r.log; numericsr-5f2bf/parambole-bnpr-5f2bf/nxprc.logr-4e5a9/parambole-bxa-4e5a9/o275x-hd.logo275x-prc-l3-prof.gs://parambole-397b-replay-nap/win-tpu/sched/r-5f2bf/parambole-bvx-5f2bf/o275x-c9pfxall-l3-prof/parambole-bvx-5f2bf-o275x-c9pfxall-l3-prof/tensorboard/plugins/profile/2026_10_15_14_43_36/server-gke-tpu-174635e9-sj65.0.0.xplane.pb.Tests: details of the recipe test and the regression subset
All run in the bench image on CPU (
JAX_PLATFORMS=cpu,--network none), comparing per-test outcomes from junit XML.The recipe test,
mlperf_submission_launcher_test.py:--pack_pad_position=start.TRAINER_TUNED_FLAGS=falsekeeps the local head and end padding, and explicit values win.TUNIX_LOGPS_SHARDED_HEAD=0|false|OFF|Nogives chunks of 1024.wi_tile_dlhs_batch_seqraised to 512;sa_use_base2_exp=falsekept;moe_wo_gmm_fused_dswigluormoe_route=deviceadded; 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;pack_pad_alignpassed; both recipes as they were before.Regression subset (22 files: the 20 the pad-first review used, plus
tests/rl/common_test.pyandtests/rl/sharded_logps_head_test.py), baseaffb1c643against this branch:rl_program_stage_before_update_test.pythat failed on base and passed here, under CPU load.grpo_recipe_wiring_test.py(PaddedBatchAssemblerTest); 6 incpu_distributed_rl_integration_test.py; 2 inyaml_generator_test.py; 1 each incluster_reaper_test.py,run_trainer_node_test.py,deepswe_mllog_utils_test.pyandrl_program_pipelined_boundary_test.py.router_replay_end_to_end_test.pydoes not finish without network, on either tree.256-chip trainer meshes (
8c75ae30d)With
TRAINER_TPU_SLICE=tpu7x:4x8x8andTRAINER_MESH_FSDP/EXPERT/CONTEXTset to 32/2/8 or 64/1/8,TRAINER_TUNED_FLAGS=trueadds the settings of Rohan Bierneni's 256-chip trainer runs:query_proj,key_proj,value_proj,out_projandqkv_proj=deviceNUM_MOE_TOKEN_CHUNKSdefaults to 1; an explicit value wins), and somoe_chunk_pipeline=falseandmoe_roe_packed_routing_gather=falsemoe_chunk_pipelineoff at any mesh, and the tunedmoe_roe_packed_routing_gatherwith it. MaxText at #5758 refuses the chunk pipeline with one chunk, and the packed routing gather without the pipeline, so 64x1x8 as listed (one chunk, pipeline off) on top of the tuned keys would not start.qkv_projalso names the GDN layers' fused qkvz projection in MaxText, soqkv_proj=devicekeeps those outputs too.bee63532f. That covers 10 launcher renders: the submission launcher at its default, at the bench cell, with tuned flags off, with buffer factor -1 and with the capacity-factor trainer, plus the 256/512/1024-device and GBS-512 recipes. Only the two new meshes and a one-chunk setting change.20c877aa1) was run on the trainer argument list of a bench run, with each rendered cell's keys and mesh. It accepts 32x2x4, 32x2x8, 64x1x8, and 64x1x8 with an explicit 2 chunks.moe_roe_packed_routing_gather=True requires: moe_chunk_pipeline=True), and one chunk with the pipeline on (moe_chunk_pipeline=True requires num_moe_token_chunks > 1).