Repository navigation
[Qwen3.5-397B MLPerf] Fused sharded logps head, start-padded packing, and 128c/256c trainer recipe defaults - #2884
Open
Rohan-Bierneni wants to merge 9 commits into
Conversation
Rohan-Bierneni
requested review from
abheesht17,
hgao327,
jiangyangmu,
lc5211,
s-noghabi,
sizhit2,
tianshub and
wang2yn84
as code owners
October 14, 2026 17:37
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 pull request implements significant performance optimizations for the Qwen3.5-397B-A17B MLPerf DeepSWE recipe on TPU v7x. By introducing a 'cp_local' mode for chunked log-probability computation and tuning various training flags, the changes aim to reduce steady-state time per micro-batch and increase overall throughput. 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 CP-local chunked LM-head execution mode (cp_local) in tunix/rl/common.py to optimize log-probabilities and entropy computation, along with corresponding tests and configuration updates. The code review identified several critical issues and style guide violations: the use of the non-existent jax.typeof function which will cause runtime errors, an unsupported preferred_element_type argument in jnp.einsum, broad exception handling that violates the repository's error-handling guidelines, and deeply nested speculative getattr calls that breach the 'One well-lit path' design principle.
Comment on lines
+1028
to
+1029
| hs_sharding = getattr(jax.typeof(hidden_states), "sharding", None) | ||
| w_sharding = getattr(jax.typeof(w), "sharding", None) |
Contributor
There was a problem hiding this comment.
Using jax.typeof is incorrect as JAX does not have a typeof function. This will raise an AttributeError at runtime. Since hidden_states and w are JAX arrays, you can directly access their sharding attribute or use getattr on the instances directly.
Suggested change
| hs_sharding = getattr(jax.typeof(hidden_states), "sharding", None) | |
| w_sharding = getattr(jax.typeof(w), "sharding", None) | |
| hs_sharding = getattr(hidden_states, "sharding", None) | |
| w_sharding = getattr(w, "sharding", None) |
| mesh, | ||
| getattr(base, "mesh", None) if base is not None else None, | ||
| getattr(model, "mesh", None) if model is not None else None, | ||
| getattr(getattr(jax.typeof(x), "sharding", None), "mesh", None) |
Contributor
There was a problem hiding this comment.
Using jax.typeof(x) is incorrect as JAX does not have a typeof function. This will raise an AttributeError at runtime. You should use x directly to check its sharding.
Suggested change
| getattr(getattr(jax.typeof(x), "sharding", None), "mesh", None) | |
| getattr(getattr(x, "sharding", None), "mesh", None) |
Comment on lines
+1320
to
+1331
| dx_c = jnp.einsum( | ||
| "nv,hv->nh", | ||
| s_c, | ||
| w_bwd, | ||
| preferred_element_type=jnp.float32, | ||
| ).astype(x_c.dtype) | ||
| dw_step = jnp.einsum( | ||
| "nh,nv->hv", | ||
| x_c.astype(w_bwd.dtype), | ||
| s_c, | ||
| preferred_element_type=jnp.float32, | ||
| ) |
Contributor
There was a problem hiding this comment.
jax.numpy.einsum does not support the preferred_element_type keyword argument. Passing it will result in a TypeError at runtime. To perform the computation in float32, cast the inputs to jnp.float32 before calling einsum.
dx_c = jnp.einsum(
"nv,hv->nh",
s_c.astype(jnp.float32),
w_bwd.astype(jnp.float32),
).astype(x_c.dtype)
dw_step = jnp.einsum(
"nh,nv->hv",
x_c.astype(jnp.float32),
s_c.astype(jnp.float32),
)
Comment on lines
+1124
to
+1128
| try: | ||
| if pltpu.is_tpu_device(): | ||
| return True | ||
| except Exception: # pylint: disable=broad-exception-caught | ||
| pass |
Contributor
There was a problem hiding this comment.
Catching a broad Exception to swallow errors violates the repository style guide. Please catch more specific exceptions, such as AttributeError or RuntimeError.
Suggested change
| try: | |
| if pltpu.is_tpu_device(): | |
| return True | |
| except Exception: # pylint: disable=broad-exception-caught | |
| pass | |
| try: | |
| if pltpu.is_tpu_device(): | |
| return True | |
| except (AttributeError, RuntimeError): | |
| pass |
References
- Avoid catching or raising overly broad exceptions like bare Exception to swallow errors. (link)
Comment on lines
+1147
to
+1148
| except Exception: # pylint: disable=broad-exception-caught | ||
| continue |
Contributor
There was a problem hiding this comment.
Catching a broad Exception to swallow errors violates the repository style guide. Please catch a more specific exception, such as AttributeError.
Suggested change
| except Exception: # pylint: disable=broad-exception-caught | |
| continue | |
| except AttributeError: | |
| continue |
References
- Avoid catching or raising overly broad exceptions like bare Exception to swallow errors. (link)
Comment on lines
+916
to
+931
| def resolve_logps_head_mode(model: Any, head_mode: str | None = None) -> str: | ||
| """Resolves the chunked LM-head execution mode, defaulting to 'off'.""" | ||
| raw = head_mode | ||
| if not raw: | ||
| raw = getattr(model, "logps_head_mode", None) | ||
| if not raw: | ||
| raw = getattr(getattr(model, "config", None), "logps_head_mode", None) | ||
| if not raw: | ||
| base = getattr(model, "base", None) | ||
| raw = getattr(getattr(base, "config", None), "logps_head_mode", None) | ||
| if not raw: | ||
| raw = ( | ||
| os.environ.get("TUNIX_LOGPS_HEAD_MODE") | ||
| or os.environ.get("LOGPS_HEAD_MODE") | ||
| or "off" | ||
| ) |
Contributor
There was a problem hiding this comment.
The deeply nested speculative getattr calls with multiple fallback layers violate the 'One well-lit path (No defensive fallbacks)' rule in the repository style guide. Consider simplifying the configuration resolution or using direct attribute access on a strictly typed contract.
References
- Never use speculative getattr(obj, "field", default) ... Access attributes and keys directly (obj.field) on a single, strictly typed contract. (link)
Rohan-Bierneni
force-pushed
the
rbierneni/mlperf-397b-cp-local-logps
branch
from
October 15, 2026 07:48
406c312 to
a4b6be1
Compare
Rohan-Bierneni
changed the title
Rohan-Bierneni
force-pushed
the
rbierneni/mlperf-397b-cp-local-logps
branch
from
October 15, 2026 08:30
a4b6be1 to
063cb17
Compare
Rohan-Bierneni
changed the title
…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.
Rohan-Bierneni
changed the title
…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.
Rohan-Bierneni
force-pushed
the
rbierneni/mlperf-397b-cp-local-logps
branch
from
October 15, 2026 17:27
4e72309 to
aeb9467
Compare
Rohan-Bierneni
changed the title
Rohan-Bierneni
force-pushed
the
rbierneni/mlperf-397b-cp-local-logps
branch
from
October 15, 2026 18:46
aeb9467 to
04719d0
Compare
s-noghabi
reviewed
Oct 15, 2026
s-noghabi
left a comment
s-noghabi
left a comment
Collaborator
There was a problem hiding this comment.
Two issues in the 256-chip (TRAINER_TPU_SLICE=tpu7x:4x8x8) branch of the tuned trainer flags. Details are in the inline comments.
| export MAXTEXT_EXTRA_FLAGS="${MAXTEXT_EXTRA_FLAGS} context_parallel_strategy=all_gather sa_fuse_reciprocal=true sa_use_base2_exp=false gdn_cp_mode=seq gdn_cp_state=device gdn_cp_matmul_precision=high cast_grads_after_all_reduce=true" | ||
| export MAXTEXT_EXTRA_FLAGS="${MAXTEXT_EXTRA_FLAGS} context_parallel_strategy=all_gather sa_fuse_reciprocal=true gdn_cp_mode=seq gdn_cp_state=device gdn_cp_matmul_precision=high cast_grads_after_all_reduce=true wi_tile_dlhs_mlp_dim=2048 decoder_layer_input=device moe_route=device prepare_moe_weights_once_per_step=true moe_roe_packed_routing_gather=true gdn_conv1d_input_in_qkv_dtype=true wi_tile_fwd_batch_seq=512 wi_tile_drhs_batch_seq=512 wo_tile_fwd_batch_seq=512 wo_tile_dlhs_batch_seq=512 wo_tile_drhs_batch_seq=512" | ||
| if [[ "${TRAINER_TPU_SLICE}" == "tpu7x:4x8x8" ]]; then | ||
| export MAXTEXT_EXTRA_FLAGS="${MAXTEXT_EXTRA_FLAGS} moe_chunk_pipeline=false" |
Collaborator
There was a problem hiding this comment.
Blocker: this crashes trainer startup on every 256-chip tuned run (both 64x1x8 and 64x2x4).
Line 263 already adds moe_roe_packed_routing_gather=true, and this line then adds moe_chunk_pipeline=false. At the MaxText this PR pins (e93800779, maxtext#5721), MoEGeneral.validate_moe_roe_packed_routing_gather in src/maxtext/configs/types.py requires moe_chunk_pipeline=True whenever the packed gather is on:
ValueError: moe_roe_packed_routing_gather=True requires: moe_chunk_pipeline=True (only the pipelined chunk loop gathers per chunk).
I can see why moe_chunk_pipeline=false is here: with 64x1x8, NUM_MOE_TOKEN_CHUNKS=1, and validate_moe_chunks rejects moe_chunk_pipeline=True when num_moe_token_chunks <= 1. But the packed gather has to be turned off along with it. test_256_chip_trainer_slice_uses_64x1x8_tuned_layout doesn't catch this because it only checks the rendered YAML strings and never builds a MaxTextConfig.
Suggested fix:
export MAXTEXT_EXTRA_FLAGS="${MAXTEXT_EXTRA_FLAGS} moe_chunk_pipeline=false moe_roe_packed_routing_gather=false"A test that runs the rendered MAXTEXT_EXTRA_FLAGS through pyconfig/MaxTextConfig validation would also catch this kind of conflict in the future.
| if [[ "${TRAINER_TPU_SLICE}" == "tpu7x:4x8x8" ]]; then | ||
| export MAXTEXT_EXTRA_FLAGS="${MAXTEXT_EXTRA_FLAGS} moe_chunk_pipeline=false" | ||
| if [[ "${TRAINER_MESH_EXPERT}" == "1" ]]; then | ||
| export MAXTEXT_EXTRA_FLAGS="${MAXTEXT_EXTRA_FLAGS} query_proj=device key_proj=device value_proj=device out_proj=device qkv_proj=device" |
Collaborator
There was a problem hiding this comment.
These conditions check the slice name, not the settings that actually require them.
moe_chunk_pipeline=false (L265) and query_proj/key_proj/value_proj/out_proj/qkv_proj=device (here) only depend on TRAINER_TPU_SLICE == tpu7x:4x8x8 (and TRAINER_MESH_EXPERT == 1). The settings that matter are NUM_MOE_TOKEN_CHUNKS and REMAT_CONTEXT:
- 64x2x4 (
TRAINER_MESH_EXPERT=2, TRAINER_MESH_CONTEXT=4):NUM_MOE_TOKEN_CHUNKS=2, so the 2-chunk pipeline is valid, but it still gets turned off just because of the slice name. That loses the chunk overlap (and, together with the issue above, also crashes). - Fully dropless / replay overrides on 4x8x8 (
RAGGED_BUFFER_FACTOR=-1,RETRY_DROPLESS_FINAL_MICRO_BATCH=true, orFIRST_PHASE_RAGGED_BUFFER_FACTORset):mlperf_gbs1024_submission.sh(L122-123) deliberately falls back toREMAT_CONTEXT=offloadandNUM_MOE_TOKEN_CHUNKS=4so the larger buffers fit. This block still keeps every attention projection on device, which works against that memory fallback and risks an HBM OOM.
Suggestion: condition on the settings themselves, e.g.
if [[ "${NUM_MOE_TOKEN_CHUNKS}" == "1" ]]; then
export MAXTEXT_EXTRA_FLAGS="${MAXTEXT_EXTRA_FLAGS} moe_chunk_pipeline=false moe_roe_packed_routing_gather=false"
fi
if [[ "${TRAINER_TPU_SLICE}" == "tpu7x:4x8x8" && "${TRAINER_MESH_EXPERT}" == "1" && "${REMAT_CONTEXT}" != "offload" ]]; then
export MAXTEXT_EXTRA_FLAGS="${MAXTEXT_EXTRA_FLAGS} query_proj=device key_proj=device value_proj=device out_proj=device qkv_proj=device"
fiAlso, mlperf_gbs1024_submission.sh L115 reads ${TRAINER_MESH_EXPERT:-1} before mlperf_397b_256_v7x.sh has set that variable, so the topology logic is duplicated in two scripts. Leaving NUM_MOE_TOKEN_CHUNKS unset there and letting L179-183 here resolve it would keep that logic in one place.
Rohan-Bierneni
force-pushed
the
rbierneni/mlperf-397b-cp-local-logps
branch
from
October 15, 2026 22:01
5037bde to
6539b6e
Compare
4 tasks done
Rohan-Bierneni
force-pushed
the
rbierneni/mlperf-397b-cp-local-logps
branch
from
October 15, 2026 22:55
6539b6e to
e5eeeab
Compare
…oundary masking, and add 256-chip 64x1x8 recipe
Rohan-Bierneni
force-pushed
the
rbierneni/mlperf-397b-cp-local-logps
branch
from
October 15, 2026 23:10
e5eeeab to
f63569b
Compare
…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.
…UNED_FLAGS Cherry-picked from google#2940 (2cfa895), companion to AI-Hypercomputer/maxtext#5783 (20c93dc4f), which AI-Hypercomputer/maxtext#5721 now carries in place of its earlier 2-line cp-as-ep.yml edit. TRAINER_TUNED_FLAGS=true now also passes custom_mesh_and_rule=cp-as-ep-dense-no-cp after the recipe's default custom_mesh_and_rule=cp-as-ep (MaxText keeps the last value of a repeated key): the dense weights (attention and shared-expert projections, the MoE router and the shared-expert gate) are sharded over fsdp x expert and replicated over context, their gradients reduce-scattered over every device and then all-gathered over context (linears.route_dense_weight_grad); the GDN projections (embed_gdn), routed experts, embeddings and activations stay as under cp-as-ep. TRAINER_TUNED_FLAGS=false renders as before. 128 TPU v7x (maxtext#5701 bench, from google#2940): 82.185 -> 81.253 s/step (3,588 -> 3,630 tok/s/chip), numerics PASS 6/6. Moves requirements/maxtext_requirements.txt to AI-Hypercomputer/maxtext#5721 at d66a82096.
…ompleted #5698 port (22fa0c5598993549ce3e80125c54300b14c7310c) d66a82096 does not start the trainer: it ports only part of #5698, so moe.py's CompoundLogicalAxis, max_utils.is_eval, models.weaver, DecoderBlockType.WEAVER and the tgmm kernel's GmmConfigs(transpose_rhs) are missing. 22fa0c559 completes the port.
…-cp-local-logps Conflict in mlperf_gbs1024_submission.sh: atwigg/mlperf (0d4d96a) removed the recipe's comment blocks that this branch had edited. Resolved by taking the removal and keeping this branch's code (the fused log-prob head default under TRAINER_TUNED_FLAGS=true, log-prob chunks of 4096, pad-first packing, the 64x1x8 token-chunk default and the PACK_PAD_POSITION=start capacity check).
…mputer/maxtext#5721 (bbd0f6f10) AI-Hypercomputer/maxtext#5721 was squash-merged into atwigg/mlperf as bbd0f6f10b17bb8038e926b0461405f2e829fc41, whose tree is identical to the PR head 22fa0c559 pinned before. Pinning the commit on atwigg/mlperf keeps the image build reproducible if the PR branch is deleted. Updates the recipe and test comments that still described the old e526ea02 pin.
andytwigg
approved these changes
Oct 17, 2026
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
5 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.
Summary
Consolidates the highest-ROI, cleanest-to-review Tunix trainer optimizations on top of
atwigg/mlperf(a928c1d), paired with companion MaxText PR AI-Hypercomputer/maxtext#5721:32×2×4)pack_pad_position="start"(--pack_pad_position=start): Places each packed row's padding at the start of the row (rounded to a multiple of64tokens, with--pack_pad_align=4096in the recipe so undercp=4the leading padding lands in CP rank 0's causal prefix where attention blocks are skipped viasa_block_q=1024, sa_block_kv=4096)1435de0)897e406TUNIX_LOGPS_SHARDED_HEAD=fused,COMPUTE_LOGPS_CHUNK_SIZE=4096): Replaces the two-pass logits projection +log_softmaxin_sharded_per_token_logps_from_hiddenwith a single-pass online-softmax + target-gather inside the vocab-tile scan loop ([B, C, V_tile]temporaries are never materialized across tiles), allowingCOMPUTE_LOGPS_CHUNK_SIZEto scale from1024→4096(16 chunks → 4 chunks)86e17b5)03b7745mlperf_gbs1024_submission.shandmlperf_397b_256_v7x.sh) + pinrequirements/maxtext_requirements.txtto AI-Hypercomputer/maxtext#5721 (e93800779)aeb9467Out-of-the-Box Recipe Defaults (
TRAINER_TUNED_FLAGS=true, default inmlperf_gbs1024_submission.sh)1. 128-chip trainer (
tpu7x:4x4x8,32×2×4)TUNIX_LOGPS_SHARDED_HEAD=fused,COMPUTE_LOGPS_CHUNK_SIZE=4096,PACK_PAD_POSITION=start.MAXTEXT_EXTRA_FLAGS): the sixteen tuned keys from Qwen3.5-397B RL trainer on 128 v7x: pad-first sequence packing, fused sharded log-prob head VJP, and the tuned trainer recipe #2913 (context_parallel_strategy=all_gather,gdn_cp_mode=seq,gdn_cp_matmul_precision=high,decoder_layer_input=device,moe_routing_maps=device,moe_routing_index_math=count,moe_unsplit_prefused_wi,moe_roe_packed_routing_gather,moe_router_grad_from_wo_input,moe_wi_gmm_fused_swiglu,gdn_conv1d_input_in_qkv_dtype,gdn_conv1d_dense_token_stack, GMMbatch_seq=512tiles, ...),num_moe_token_chunks=2,moe_chunk_pipeline=true.custom_mesh_and_rule=cp-as-ep-dense-no-cp(Add mesh rules that replicate dense weights across context (cp-as-ep-dense-no-cp, cp-as-ep-dense-fsdp) AI-Hypercomputer/maxtext#5783, carried in #5721; cherry-picked from recipes(397b v7x): the cp-as-ep-dense-no-cp mesh rule joins TRAINER_TUNED_FLAGS #29402cfa8950a): dense attention / shared-expert / router weights sharded overfsdp × expertand replicated overcontext, withroute_dense_weight_gradreduce-scatter + context all-gather for their gradients; GDN projections stay onfsdp × context × expert(embed_gdn). 128c: 82.185 → 81.253 s/step (recipes(397b v7x): the cp-as-ep-dense-no-cp mesh rule joins TRAINER_TUNED_FLAGS #2940).DiffState,xor_dncGDN triangular inverse, GDN prefill zero-fill skip, unpacked-row segment skip).requirements/maxtext_requirements.txtpinsatwigg/mlperfatbbd0f6f10, the squash-merge of [Qwen3.5-397B MLPerf] Consolidated top-ROI trainer & GDN/MoE optimizations on atwigg/mlperf AI-Hypercomputer/maxtext#5721 (same tree as its final head22fa0c559, which completes the #5698 port).2. 256-chip trainer (
TRAINER_TPU_SLICE=tpu7x:4x8x8, 512 devices)Logic cherry-picked from #2913 (
8c75ae30d) so both PRs render the same 256-chip flags:fsdp×ep×cp)64×1×8(default with tuned flags)TRAINER_TPU_SLICE=tpu7x:4x8x8num_moe_token_chunks=1,moe_chunk_pipeline=false,moe_roe_packed_routing_gather=falsequery/key/value/out/qkv_proj=device32×2×8+ TRAINER_MESH_FSDP=32 TRAINER_MESH_EXPERT=2num_moe_token_chunks=2,moe_chunk_pipeline=true,moe_roe_packed_routing_gather=truequery/key/value/out/qkv_proj=devicemoe_chunk_pipelinefollows the chunk count andmoe_roe_packed_routing_gatherfollows the pipeline (MaxText refuses the pipeline with one chunk, and the packed gather without the pipeline). An explicitNUM_MOE_TOKEN_CHUNKSwins.TRAIN_MICRO_BATCH_SIZEdefaults tofsdp × ep(64 for both meshes).Results (RL trainer bench, same 9 micro-batches × 64 rows per step, steps 1–4 average)
32×2×432×2×864×1×8Re-run with
custom_mesh_and_rule=cp-as-ep-dense-no-cp(#5721d66a82096code + #2884, same bench)cp-as-eps/stepcp-as-ep-dense-no-cps/step32×2×432×2×864×1×8*The reservation ended during the 32×2×8 run; only steps 0–1 were captured. Its steady microbatches (6.538–6.539 s) were flat, as in every other run.
Best 256-chip config:
64×1×8+cp-as-ep-dense-no-cp= 55.63 s/step, 1.477× the 128ccp-as-epstep (82.17 s), 2,651 tok/s/chip.Launcher tests:
tests/experimental/examples/deepswe_dist/mlperf_submission_launcher_test.py48/48 pass (incl. the 5 new 256-chip mesh tests).