Repository navigation
[Qwen3.5-397B MLPerf] Consolidated top-ROI trainer & GDN/MoE optimizations on atwigg/mlperf - #5721
Merged
Rohan-Bierneni merged 16 commits intoOct 16, 2026
Merged
Conversation
Rohan-Bierneni
requested review from
A9isha,
NuojCheng,
RissyRan,
SurbhiJainUSC,
abhinavclemson,
aireenmei,
bvandermoon,
darisoy,
dipannita08,
gagika,
gobbleturk,
hengtaoguo,
huytransformer,
igorts-git,
jiangjy1982,
khatwanimohit,
michelle-yooh,
parambole,
richjames0,
shralex,
shuningjin,
vipannalla,
xibinliu and
zxhe-sean
as code owners
October 14, 2026 17:36
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces several performance optimizations for MoE models and GDN kernels, including deferring small all-reduces, hoisting MoE weight preparation, and adding fast paths for convolution operations. It also includes corresponding configuration updates and unit tests. The review identified a missing import in maxtext_engine.py and an incorrect string comparison for identifying the RMSNorm scale parameter in nnx_decoders.py, both of which require correction.
| if getattr(self._config, "prepare_moe_weights_once_per_step", False) and not getattr( | ||
| self._config, "parameter_memory_host_offload", False | ||
| ): | ||
| moe.prepare_moe_weights_for_step(mdl, self._config) |
There was a problem hiding this comment.
The moe module is referenced here (and in collect_moe_audit and diff_wrapper), but it is not imported in maxtext_engine.py. This will cause a NameError at runtime. Please add from maxtext.layers import moe to the imports at the top of maxtext_engine.py.
| if not norm_axes and (batch_axes or seq_axes) and set(batch_axes + seq_axes) == active_mesh_axes: | ||
|
|
||
| def _maybe_expand_norm_scale(path, leaf): | ||
| if isinstance(leaf, nnx.Param) and path and str(path[-1]) == "scale": |
There was a problem hiding this comment.
Using str(path[-1]) == "scale" to identify the scale parameter will always evaluate to False because JAX tree path elements (like DictKey or GetAttrKey) do not have a string representation equal to the raw key string (e.g., DictKey(key='scale') stringifies to "['scale']" and GetAttrKey stringifies to ".scale"). This silently disables the defer_small_all_reduces optimization for RMSNorm scale.\n\nInstead, check the key or name attribute of the path element directly.
Suggested change
| if isinstance(leaf, nnx.Param) and path and str(path[-1]) == "scale": | |
| if isinstance(leaf, nnx.Param) and path and (getattr(path[-1], "key", None) == "scale" or getattr(path[-1], "name", None) == "scale"): |
Rohan-Bierneni
force-pushed
the
rbierneni/mlperf-397b-trainer-perf
branch
from
October 15, 2026 07:46
0f32015 to
23c7490
Compare
Rohan-Bierneni
changed the title
Rohan-Bierneni
added a commit
to Rohan-Bierneni/tunix
that referenced
this pull request
Oct 15, 2026
Rohan-Bierneni
force-pushed
the
rbierneni/mlperf-397b-trainer-perf
branch
from
October 15, 2026 08:23
23c7490 to
1fcdee5
Compare
Rohan-Bierneni
requested review from
Lumosis,
jshin1394,
liudangyi,
patemotter and
xuefgu
as code owners
October 15, 2026 21:58
4 tasks done
…o atwigg/mlperf
fwd_bwd differentiated every parameter, and update() then zeroed the gradients of the ones trainable_parameters_mask freezes before the norm, clipping and the optimizer saw them. Every micro-batch computed, reduce-scattered and accumulated those gradients only to throw them away. diff_wrapper now applies jax.lax.stop_gradient to the frozen leaves, inside the differentiated function: JAX emits no weight-gradient op for them and fwd_bwd returns a constant zero in their place. The gradient through a frozen weight into the activations still comes from the same ops, so the trainable gradients, the norm and the update are the same computation. Unconditional whenever a mask is set, with no config key. Nothing in the engine reads a frozen leaf's gradient: the norm, clipping, the NaN and spike skips and the optimizer all come after update() zeroes it, and no API returns gradients. update() still zeroes them, for an accumulator restored from an intra-step checkpoint written by an older engine or under another mask. The leaves keep their slots in the gradient tree, as zeros, so kernel signatures, the accumulator and checkpoints are unchanged. pre_train/train.py is untouched: it clips over the whole gradient tree, frozen leaves included, so the same stop there would change its updates (README, "Differences from pre_train/train.py"). Both kernels read the mask through one helper, _freeze_mask_for, so the leaves fwd_bwd stops are exactly the leaves update() zeroes. It reads the compile-time cache once and, when it is empty, computes the mask from the tree, as _update_kernel already did: on the eager path, and for a dropless or final kernel traced on a background thread after _disable_pure_state emptied the cache, where checking the cache alone would silently skip the stop. The MLPerf RL recipes freeze the MoE router (^(?!.*routed_experts/gate/kernel).*). 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), with the same stop behind an environment switch: - AOT with the TPU libtpu: the router's f32[4096,512] weight-gradient fusions and f32[16,512] reduce-scatters are gone from fwd_bwd (8 -> 0 each in the program, one per MoE layer when run). - Regular micro-batch 10.189 -> 10.131 s, both with moe_routing_maps=device (o275x-stk1a-a vs o275x-stk2, two jobs on one booking): -58 ms against a predicted -45 ms (-27 to -65). - Numerics A/B (nxstk2 vs nxref, nxref2): 5 of 5 gates pass; loss 4.6e-8 relative, grad_norm 4.8e-5 relative against a replicate spread of 1.9e-5. Tests (maxtext_engine_frozen_grad_test.py, new), against the same engine with the stop removed: on a tiny decoder with one MLP kernel frozen, trained by AdamW with clipping on four CPU devices, fwd_bwd lowers to one dot_general fewer, shaped like the kernel, and none added; it returns exactly zero for the kernel; and every other gradient, the running sums, the gradient norms and the trained weights are bitwise equal, under fsdp, under the deferred all-reduce and on the eager path. On the one-layer Qwen3.5 MoE rig, with the recipes' router mask and Tunix's router replay payload and loss, the same holds. A stop that also hits one trainable kernel is caught, and the failure names it. README: the config key, the test and the difference from pre_train/train.py.
…e_routing_index_math) The ring-of-experts ragged sort derives its routing index maps from the top-k expert ids with sorts: a stable argsort for the expert permutation, a gather of a repeated iota for the sorted token indices, a second argsort for the inverse permutation, a stable argsort in the SparseCore gather-reduce's validity preprocessing, and, for tokens without forced ids (every token under moe_exclude_padding_from_routing), jax.lax.top_k on the router probabilities. XLA lowers each stable argsort in the forward pass with an extra, duplicated iota operand. moe_routing_index_math="count" (default "sort", unchanged) computes the same maps with less work: - the permutation by one single-operand sort of the packed uint32 key (expert << pos_bits | slot), whose values are distinct, so the unstable sort has exactly the stable order (packed_stable_sort); - the sorted token indices as slot // topk, without a gather; - the inverse permutation by counting, on this shard's slots only, -1 elsewhere (ragged_sort._local_slot_positions): a slot's rank is the slots of lower local experts, plus the slots of its expert in earlier tokens (a blocked prefix count as a triangular bf16 matmul, exact), plus earlier slots of the same token with the same expert. Every reader of the inverse masks it to this shard's window, so nothing reads the -1s, and the unsort backward gets the forward permutation instead of argsorting the partial inverse back (it raises without it); - the router-weight gradient back to slot order by a masked gather instead of a TensorCore scatter (the table is padded to the slot count so XLA offloads the gather to the SparseCore); - the gather-reduce validity partition as one packed uint32 sort (_preprocess(packed=True)); - the top-k indices of non-forced tokens by k argmax passes (topk_indices_by_argmax), which break ties the way top_k does. That these are cheaper than top_k's sort is expected, not measured: they were timed only as part of this key. Exact: the same permutation everywhere, so the outputs and gradients are the same bits. The validator rejects "count" where no ring-of-experts ragged sort runs (it needs sparse_matmul, te_moe_block=False, use_ring_of_experts and use_ragged_sort). It starts the validator form the later keys of this series share: _reject_unless names every unmet requirement, and the fields of the other config mixins are read directly rather than through getattr with a default. With the key off the lowered TPU program of the layer is the same as before the change (checked on the forward-plus-backward lowering of a shrunk Qwen3.5 MoE layer for v7x). Measured on Qwen3.5-397B-A17B, 128 TPU v7x, on the RL trainer benchmark (mesh fsdp 32 x context 4 x expert 2, 9 micro-batches per step), with this change switched by an environment variable in a build that a TPU lowering check shows compiling to the same program as this commit: regular micro-batch 10.022 -> 9.875 s (-0.147 s; arm o275x-mf-cnt-prof against o275x-mf-ctl, the flags-off control in the same pod), 99.25 -> 98.05 s per step, 0 dropped tokens in both. The same control measured in different pods spreads by +/-0.02-0.03 s per micro-batch. The timing arms run with the benchmark's log-prob gates on (loss and gradient norm read 0), so they say nothing about numerics. TPU numerics: a 3-step A/B with the gates off, against a reference run and a replicate of it, passes all six checks with this key in both measured stacks: with moe_unsplit_prefused_wi, moe_prefused_swiglu_vjp, moe_roe_packed_routing_gather and moe_router_grad_from_wo_input (run nxs6c: loss relative difference 4.6e-8, N1 2.9e-6, gradient-norm difference 7.0e-6 against a replicate spread of 1.9e-5), and with the fused SiLU GMM in place of the fused activation (run nxs6silu: gradient-norm difference 2.5e-6). Caveats: that is the stack, not this key alone; and it ran on the environment-switch build, not on this branch's binary. Tests (moe_routing_index_math_test.py, new), on CPU: - _local_slot_positions equals argsort(argsort(ids)) minus the window start on every shard's slots and -1 elsewhere, across partial blocks, the padding id and repeated experts; topk_indices_by_argmax equals top_k's indices with ties; _preprocess(packed=True) equals the argsort version. The SparseCore gather-reduce is TPU-only (CPU takes its reference fallback), so its preprocessing is tested directly. - On the shrunk Qwen3.5 routed MoE on 8 CPU devices (cp-as-ep, local routing, two chunks), "count" is bit-identical to "sort" in the output, the loss and every gradient: plain and pipelined chunk loop, padding excluded from routing, forced ids with a buffer that drops tokens, dropless, saved routing maps, bf16 router, and the benchmark's MoE settings (expert-weight prefetch, row-major reduce-scatter, routing maps on device, ragged_buffer_factor 2.75, fp8 fake-quant). - Structure, two-sided: "count" traces fewer sorts than "sort", none of them stable or with more than one operand (so a partial fallback to jnp.argsort fails), and no top_k under padding exclusion; every ragged_gather_reduce call gets packed_preprocess=True under "count" and False under "sort" (a spy, since only the TPU kernel reads it). A counted inverse shifted by one is caught, and "count" without the forward permutation raises. - The validator rejects "count" for each unmet requirement, naming it.
…ort_unique_pair_argsort, moe_prefused_swiglu_vjp) Three rewrites of MoE glue ops that leave every value unchanged, each behind its own key (default off, program unchanged): moe_unsplit_prefused_wi. With prefuse_moe_weights, the layer sliced the fused wi = [w0 | w1] into its halves at entry, so the fp8 fake-quant ran twice, the FSDP all-gather and the expert-weight prefetch barrier carried two operands, and gmm_up concatenated the gathered halves back before the up-projection. The key keeps wi whole up to the GEMM; w1_kernel=None marks the layout in sparse_matmul. The fake-quant scales are per output channel, so rounding the whole tensor gives the halves' values. moe_ragged_sort_unique_pair_argsort. With moe_routing_index_math="sort", the ring-of-experts ragged sort's two argsorts and the SparseCore gather-reduce's validity sort run as unstable sorts of the unique (key, index) pairs (argsort_unique_pairs). A stable jnp.argsort lowers in the forward pass with a third, duplicated iota operand; the pair sort has two operands and exactly one result, the stable permutation (on v7x, an s32[655360] argsort takes 0.61-0.63 ms with three operands and 0.42 ms with two). moe_prefused_swiglu_vjp. With prefuse_moe_weights, the activation on the whole [gate | up] up-projection output runs as one custom-VJP op whose backward takes jax.vjp of the same activation and returns d[gate | up] as one array. Split, the two [rows, mlp] cotangents leave their fusion separately and XLA joins them (add of two pads) in an extra pass over the expert buffer. Any gated activation works (SwiGLU, GPT-OSS's clipped form). Same values; the gradient bits can differ only in the sign of an exact zero (the add of zero pads gives -0.0 + 0.0 = +0.0, the concatenate keeps -0.0; measured on XLA:CPU). Which keys to use. With the rest of this series the recommended set is moe_routing_index_math="count", moe_unsplit_prefused_wi, moe_roe_packed_routing_gather, moe_router_grad_from_wo_input and moe_wi_gmm_fused_swiglu. Of this commit's keys only the whole wi is in it: "count" removes the sorts the pair argsort rewrites, and the fused SiLU GMM never forms the [gate | up] output the fused activation reads, so the validators reject the pair argsort with "count" and the fused activation with the fused SiLU. The two stay as standalone alternatives for configurations the recommended keys do not cover: the pair argsort for runs that keep moe_routing_index_math="sort" (it keeps the maps as sorts, the smaller change), and the fused activation where the fused SiLU GMM does not apply (it needs tokamax gmm_v2 and plain silu + linear; the fused activation works with any GMM backend and gated activation). The validators, in the form moe_routing_index_math started, reject each key for every unmet requirement: the pair argsort off the ring-of-experts ragged sort or with "count"; the whole wi off prefused sparse-matmul training, with weight quantization, non-default expert-weight sharding, embedding chunks, chunk weight-gradient accumulation, or an expert MLP dimension that tensor, tensor_sequence or autoregressive parallelism shards (a device would then hold a slice of [w0 | w1] instead of matching slices of w0 and w1); the fused activation off prefused sparse matmul, with MLP biases or embedding chunks, or under a remat policy that saves the routed experts' gate or up half (_remat_keeps_moe_tensor_requirement, shared with the later keys of this series: the minimal* presets, or custom with mlpwi_0, moe_mlpwi_0 or moe_mlpwi_1 saved or offloaded). The fused activation never forms those halves, so such a policy would silently stop saving the MoE tensor (the shared expert's dense MLP, tagged mlpwi_0 too, would still be saved). An MLP dimension sharded under custom logical axis rules, weight quantization objects and sparsity, which only the built layer knows, raise at trace time instead of falling back. Measured on Qwen3.5-397B-A17B, 128 TPU v7x, on the RL trainer benchmark (mesh fsdp 32 x context 4 x expert 2, 9 micro-batches per step), all three together, switched by an environment variable in a build that a TPU lowering check shows compiling to the same program as this commit: regular micro-batch 10.022 -> 9.852 s (-0.170 s; arm o275x-mf-gb-prof against its in-pod control o275x-mf-ctl), 99.25 -> 97.70 s per step, 0 dropped tokens in both. The same control measured in different pods spreads by +/-0.02-0.03 s per micro-batch. The parts were not timed separately, and moe_routing_index_math was "sort" in that arm. The timing arms run with the benchmark's log-prob gates on (loss and gradient norm read 0), so they say nothing about numerics. TPU numerics: a 3-step A/B with the gates off, against a reference run and a replicate of it, passes all six checks for moe_unsplit_prefused_wi and moe_prefused_swiglu_vjp stacked with moe_routing_index_math="count", moe_roe_packed_routing_gather and moe_router_grad_from_wo_input (run nxs6c: loss relative difference 4.6e-8, N1 2.9e-6, gradient-norm difference 7.0e-6 against a replicate spread of 1.9e-5), and for the whole wi in the recommended set above (run nxs6silu: gradient-norm difference 2.5e-6). Caveats: stack level, not per key; the environment-switch build, not this branch's binary; and the pair argsort is in neither run (it is inert next to "count"; a permutation, exact by construction, and bit-identical on CPU). Tests (moe_glue_rewrites_test.py, new), on CPU: - argsort_unique_pairs equals the stable argsort (ties, booleans, every axis, index dtype); _preprocess gives the same partition with it (the SparseCore gather-reduce is TPU-only, so its preprocessing is tested directly); _prefused_ffn_activation gives the split value bit for bit and the split VJP value for value, with bit differences only at zeros (SwiGLU and clipped GPT-OSS, bf16 and f32, a -0.0 cotangent row). - On the shrunk Qwen3.5 routed MoE on 8 CPU devices with prefuse_moe_weights (cp-as-ep, local routing, two chunks), each key alone and all three are bit-identical to the plain program in the output, the loss and every gradient: plain and pipelined chunk loop, fp8 fake-quant with padding excluded, forced ids with drops and dropless, with "count", and with the benchmark's MoE settings (expert-weight prefetch, row-major reduce-scatter, routing maps on device, ragged_buffer_factor 2.75). - Structure: fewer concatenates (and fake-quant barriers) with the whole wi; exactly two two-key sorts per chunk with the pair argsort, and unique_pair_argsort=True on every gather-reduce call (a spy, since only the TPU kernel reads it); one more custom_vjp per chunk with the fused activation. A fused backward that swaps d(gate) and d(up), and a pair argsort with reversed ties under a dropping buffer, are caught. - On an expert 4 x tensor 2 mesh the layer runs without the key and raises at trace time when the key reaches it past the validator. - Each validator rejects its key for each unmet requirement, naming it.
… experts (moe_roe_packed_routing_gather) With ring_of_experts_local_routing and moe_chunk_pipeline, every token chunk routes its own tokens and all-gathers its top-k weights and expert ids over the expert-parallel axis: two all-gathers, two SparseCore offloads, per chunk, and in the backward one weight-cotangent reduce-scatter per chunk. moe_roe_packed_routing_gather routes all of the layer's local tokens once and all-gathers weights and ids as ONE float32 [batch, seq, 2 * top_k] array (pack_routing_for_all_gather: ids are < num_experts or -1, exact in float32), which each chunk slices: one gather per layer and one reduce-scatter in the backward. Routing is per token, so each slice holds what the chunk would have computed: bit-identical with float32 routing weights (float32_weight_sum, the default). With bfloat16 weights the values round-trip exactly and their cotangent is reduce-scattered in float32 instead of bfloat16. The all-gather is the ring of experts' own (_ep_all_gather, now defined once for roe_ag_and_route and this gather), so a SparseCore pin applies to both alike. The per-chunk load-balance loss is computed as without the key, from this shard's own ids of the chunk, so the loss and the output are bit-identical; the router gradient then agrees to float32 reduction order (the top-k and the loss terms reach the logits in another order; 1.2e-7 relative on CPU). The gathered array is a routing map: tagged moe_routing_maps, so moe_routing_maps=device saves it and the rematerialized forward skips the gather; untagged under remat. The validator (the shared form) requires sparse_matmul, te_moe_block=False, ring_of_experts_local_routing and moe_chunk_pipeline (the packed gather replaces the pipelined loop's per-chunk gathers only), and rejects use_random_routing (random routing draws per call, so one draw for the layer differs from one per chunk). If XLA strips optimization barriers the pipeline is skipped with a warning, and the packed gather too. Measured on Qwen3.5-397B-A17B, 128 TPU v7x, on the RL trainer benchmark (mesh fsdp 32 x context 4 x expert 2, 9 micro-batches per step, load_balance_loss_weight 0, moe_routing_maps=device), switched by an environment variable in a build that a TPU lowering check shows compiling to the same program as this commit: regular micro-batch 9.938 s (arm o275x-mf-pra-prof), 98.68 s per step, 0 dropped tokens. That pod ran no flags-off control; against the flags-off control of another pod (o275x-gdn-ctl, 10.033 s, 99.50 s per step) that is -0.095 s per micro-batch, with a spread of the same control across pods of +/-0.02-0.03 s per micro-batch. The environment-switch build dropped the load-balance loss (it returned None for it); this change computes it, which the measured run, at weight 0, does not exercise. The timing arms run with the benchmark's log-prob gates on (loss and gradient norm read 0), so they say nothing about numerics. TPU numerics: a 3-step A/B with the gates off, against a reference run and a replicate of it, passes all six checks with this key in both measured stacks: with moe_routing_index_math="count", moe_unsplit_prefused_wi, moe_prefused_swiglu_vjp and moe_router_grad_from_wo_input (run nxs6c: loss relative difference 4.6e-8, N1 2.9e-6, gradient-norm difference 7.0e-6 against a replicate spread of 1.9e-5), and with the fused SiLU GMM in place of the fused activation (run nxs6silu: gradient-norm difference 2.5e-6). Caveats: stack level, not this key alone; the environment-switch build, not this branch's binary. Tests (moe_roe_packed_routing_gather_test.py, new): - pack_routing_for_all_gather / unpack_routing round-trip bit for bit with ids up to 4095 and -1 (bfloat16 holds every integer only up to 256), float32 and bfloat16 weights, and the packed array is float32. - On the shrunk Qwen3.5 routed MoE on 8 CPU devices (cp-as-ep, local routing, pipelined loop, two chunks), bit-identical output, loss and every gradient against the per-chunk gathers: real routing, padding excluded from routing, forced ids with a dropping buffer, saved routing maps, with moe_routing_index_math=count, and with the benchmark's MoE settings (expert-weight prefetch, row-major reduce-scatter, ragged_buffer_factor 2.75, prefused weights with fp8 fake-quant). With the load-balance loss on: output and loss bit-identical, gradients within 1e-5 of max (measured 1.2e-7). bf16 routing weights: within relative L2 1e-2 (bit-identical on CPU). - Structure: the plain loop all-gathers two [..., top_k] arrays per chunk and the packed one exactly one [..., 2 * top_k] array, tagged moe_routing_maps only under moe_routing_maps=device; nothing is tagged moe_route. - A packed gather shifted along the sequence, and the packed path with its load-balance loss dropped, are caught. - The validator rejects the key for each unmet requirement, naming it, and accepts the load-balance loss. The SparseCore-pinned all-gather is a TPU-only placement and is not reached on CPU.
The up projection wrote the [buffer, 2 * mlp] GMM output and a separate
elementwise pass read it back to apply the SwiGLU, over every buffer row,
the ragged buffer's padding included.
moe_wi_gmm_fused_swiglu calls megablox ops.gmm_silu instead: the tokamax
gmm_v2 kernel with fuse_act="silu" applies the activation to the float32
accumulator of the rows it computes and writes the [buffer, mlp]
activation. When differentiated it also writes the gate and up halves
([buffer, 3 * mlp] in all), rounded to the output dtype exactly as the
unfused GMM rounds its output (the new return_pre_act option of gmm_v2:
both halves go through the store path of the 2D output, with its
zero-init DMAs and partial-row carry), and the custom VJP is the unfused
backward: the SwiGLU VJP on those bf16 halves, then gmm's own DLHS / TGMM
on [d gate | d up]. ops._gmm_silu_bwd calls gmm's backward with keyword
arguments, as gmm's custom VJP would for a plain bf16 gmm_v2 GEMM. Works
with split weights and with the whole [w0 | w1] of moe_unsplit_prefused_wi.
The layer hands the fused GMM the routing operands explicitly
(get_gmm_operands_for_local_experts, which get_gmm_for_local_experts now
wraps; the key-off program is unchanged).
Numerics change: silu reads the f32 accumulator instead of its bf16
rounding (the more precise value), so the activation, and through it
everything downstream, differs by rounding.
The validator (the shared form) requires prefused sparse-matmul training,
use_tokamax_gmm and use_gmm_v2, silu + linear activations without
clipping (and not GPT-OSS), no weight quantization, MLP bias, embedding
chunks, flat FSDP weights, 3D TC GMM, token all-gather quantization or
chunk weight-gradient accumulation, and a remat policy that saves neither
half of the routed experts' up projection (the rule
moe_prefused_swiglu_vjp uses: the fused GMM never forms the
[buffer, 2 * mlp] output, so such a policy would silently stop saving the
MoE tensor; remat_policy=full and the shared expert's own names are
fine). moe_prefused_swiglu_vjp is rejected next to it, as it would do
nothing. Weight quantization objects and an explicit weight all-gather,
which only the built layer knows, raise at trace time.
Measured on Qwen3.5-397B-A17B, 128 TPU v7x, on the RL trainer benchmark
(mesh fsdp 32 x context 4 x expert 2, 9 micro-batches per step), switched
by an environment variable in a build whose kernel files differ from these
only in docstrings, comments and line breaks (AST-identical) and whose
program a TPU lowering check matches to this commit. Same pod, stacked on
moe_routing_index_math=count, moe_unsplit_prefused_wi,
moe_prefused_swiglu_vjp, moe_roe_packed_routing_gather and
moe_router_grad_from_wo_input (the fused activation has nothing to act on
next to this key, and the validator rejects the two together): regular
micro-batch 9.494 -> 9.418 s (-0.0755 s), 86.970 -> 86.253 s per step
(-0.716 s), and the step's second micro-batch, a proxy for the step
boundary, 9.450 -> 9.375 s, so the boundary pays nothing extra (arms
o275x-s6silu-l3-prof against o275x-s6c-l3-prof, one pod, on a build with
the staged step boundary, which is why its steps are shorter than the
other arms of this series). The same control measured in different pods
spreads by +/-0.02-0.03 s per micro-batch. The timing arms run with the
benchmark's log-prob gates on (loss and gradient norm read 0), so they say
nothing about numerics.
TPU numerics: a 3-step A/B with the gates off, against a reference run and
a replicate of it, passes all six checks for this key stacked with
moe_routing_index_math=count, moe_unsplit_prefused_wi,
moe_roe_packed_routing_gather and moe_router_grad_from_wo_input (run
nxs6silu: N1 1.4e-6, gradient-norm difference 2.5e-6, within the
replicate spread of 1.9e-5 measured with the same reference). Caveats:
stack level, not this key alone; the environment-switch build, not this
branch's binary. megablox_gmm_silu_test.py, the kernel's own TPU test,
has not been run yet, so that stack run is the kernel's only TPU
evidence.
Tests:
- megablox_gmm_silu_test.py (new, TPU only, skipped on CPU): gmm_v2's
gate and up equal the unfused output bit for bit, rows outside the
groups are zero, act is silu of the f32 accumulator, the act-only call
matches, with an EP shard (group_offset); gmm_silu's VJP equals the
unfused path's value for value.
- moe_wi_gmm_fused_swiglu_test.py (new, CPU):
- ops.gmm_silu with stand-ins for the Pallas kernel and gmm's TPU
backward only: the forward tile is tiling[2] halved, act alone is
written undifferentiated and act + gate + up differentiated, the
backward hands gmm's backward [d gate | d up] of the SwiGLU on the
bf16 halves bit for bit, with this op's operands as residuals and its
DLHS options on the right keywords (each option set alone, both
group offsets); an odd forward tile is rejected.
- The shrunk Qwen3.5 MoE on 8 CPU devices, with jnp stand-ins for gmm
and gmm_silu (gmm_silu as the unfused composition of the same
stand-in): the layer with the key is bit-identical to the layer
without it (output, loss, every gradient): plain and pipelined loop,
dropless (non-zero group offsets), padding excluded, the whole
prefused wi, with the other MoE keys of this series, and with those
and the benchmark's MoE settings. One gmm_silu call per chunk, with
the plain path's wi tiling and the config's DLHS options
(moe_gmm_v2_dlhs_transpose_rhs True in the benchmark case). Stand-ins
that swap the gate and up halves, or ignore the group offset, are
caught.
- Key off: the dropless ring-of-experts ragged sort matches the ring of
experts without it (use_ragged_sort=False), and fails to once its GMM
ignores the group offset, so the routing operands that both paths now
take from get_gmm_operands_for_local_experts are checked against a
reference this change does not touch.
- The validator rejects the key for each unmet requirement, naming it.
…ward that reads the GEMM input (moe_router_grad_from_wo_input)
On the ring-of-experts ragged path the down projection y = h @ wo (per
expert, over the sorted rows) feeds ring_ragged_unsort, whose backward
needs y for the routing-weight gradient <g, y>. Under remat the layer's
backward therefore recomputes the down-projection GEMM only to form y,
and the combine backward makes a pass over two [rows, hidden] tensors that
also writes dy = w * g.
moe_router_grad_from_wo_input runs the activation, the down projection and
the unsort as one custom-VJP op (ragged_sort.ring_ragged_unsort_down_projection).
Its forward is the GEMM followed by ring_ragged_unsort's forward, so the
output is bit-identical. Its backward never reads y: row j of the GEMM is
linear in h_j and the weight is a scalar per row, so
<g, h wo> = <g wo^T, h>, (w g) wo^T = w (g wo^T) and h^T (w g) = (w h)^T g.
It takes the GEMM's own VJP at the input w * h with the unweighted
cotangent g (u = g wo^T, dwo = (w h)^T g), then dh = w u and the routing
gradient <u, h>. The GEMMs keep their shapes, the recomputed forward GEMM
has no reader and XLA drops it, and the weighting and the dot run over
[rows, mlp] tensors. In the pipelined chunk loop, tie 2 moves into the op
as a forward-only barrier on the GEMM output. The cotangent rows outside
this shard's groups are left unmasked: the GEMM backward never mixes them
into group rows and every output is masked.
The op shares its code with ring_ragged_unsort: the unsort's forward and
the fan-out half of its backward are now helpers (_ring_unsort_window,
_ring_unsort_combine, _ring_unsort_fan_out, with _unsort_kwargs for their
kernel options) that both call, so the fused forward is the unsort's by
construction. The helpers keep the unsort's ops in their order: lowered
for TPU v7x (forward plus backward of a shrunk Qwen3.5 MoE layer, the
benchmark's MoE settings, also the dropless retry program), the key-off
program is the same as the parent commit's and every key-on program the
same as before this refactor.
Numerics change: each gradient keeps one rounding to the activation dtype,
at a different point (w * h instead of w * g; dh after the product with w
instead of before the GEMM; the routing gradient from g wo^T instead of
y), so the gradients agree with the two ops apart to bf16 rounding, not bit
for bit. On the CPU layer below they differ by 1e-3 to 4.2e-3 relative L2.
On the op alone, against the same function in float32, the fused
gradients are as far as the two ops' are: the ratio of the two errors is
0.99-1.01 for dh and dwo and 0.94-1.07 for the routing gradient, over 6
seeds with a full and a dropping buffer.
The validator (the shared form) requires the ring-of-experts ragged sort
(sparse_matmul, te_moe_block=False, use_ring_of_experts and
use_ragged_sort) and a down projection that is the bare GEMM of the
sorted rows: no tensor parallelism (also checked on the mesh at trace
time), MLP bias, weight quantization, TC ragged sort, chunk weight-gradient
accumulation, pre-quantized unsort backward, token all-gather quantization
or embedding chunks; and a remat policy that saves no name of the routed
experts' down-projection output (mlpwo, moe_mlpwo: the minimal* and
save_dot_except_mlpwi presets, or custom with either saved or offloaded),
the rule of the series: the op never forms that output, so such a policy
would silently stop saving the MoE tensor. It composes with
moe_routing_index_math, the prefused-weight keys and
moe_wi_gmm_fused_swiglu. The layer raises if it is asked to return the
weights for chunk weight-gradient accumulation; _moe_body's three exits
now share one construction of its outputs.
Measured on Qwen3.5-397B-A17B, 128 TPU v7x, on the RL trainer benchmark
(mesh fsdp 32 x context 4 x expert 2, 9 micro-batches per step), switched
by an environment variable in a build that a TPU lowering check shows
compiling to the same program as this commit: regular micro-batch 10.022
-> 9.765 s (-0.257 s; arm o275x-mf-a14-prof against its in-pod control
o275x-mf-ctl), 99.25 -> 97.05 s per step, 0 dropped tokens in both. All
five MoE keys of this series together (moe_routing_index_math=count,
moe_unsplit_prefused_wi, moe_roe_packed_routing_gather, this key and
moe_wi_gmm_fused_swiglu), in one pod on a build with the staged step
boundary: 10.0445 -> 9.510 s per micro-batch (-0.535 s), 91.87 -> 87.04 s
per step (o275x-s5moe5-l3-prof against o275x-s5c-l3-prof). The same
control measured in different pods spreads by +/-0.02-0.03 s per
micro-batch. The timing arms run with the benchmark's log-prob gates on
(loss and gradient norm read 0), so they say nothing about numerics.
TPU numerics: a 3-step A/B with the gates off, against a reference run and
a replicate of it, passes all six checks with this key in both measured
stacks: with moe_routing_index_math=count, moe_unsplit_prefused_wi,
moe_prefused_swiglu_vjp and moe_roe_packed_routing_gather (run nxs6c: loss
relative difference 4.6e-8, N1 2.9e-6, gradient-norm difference 7.0e-6
against a replicate spread of 1.9e-5), and with the fused SiLU GMM in
place of the fused activation (run nxs6silu: gradient-norm difference
2.5e-6). Caveats: stack level, not this key alone; the environment-switch
build, not this branch's binary.
Tests (moe_router_grad_from_wo_input_test.py, new), on 8 CPU devices:
- The op on a 4-device EP mesh with a jnp grouped matmul as the down
projection, the full buffer and a packed buffer that drops tokens,
"sort" and "count": output bit-identical to the two ops apart, every
gradient within relative L2 1e-2 of theirs; a forward tie comes back
unchanged with its cotangent passed through and is tied to the GEMM
output by a barrier; "count" without the forward permutation raises.
- Precision of the op alone (its input h given, so only the backward
rounds), against the same function written densely in float32 (the
ragged gather-reduce returns bf16, so the two ops cannot be the
reference): each fused gradient at most 1.10x as far as the two ops',
on 6 seeds x {full buffer, dropping buffer}, with routing weights where
bf16 rounds them worst. A backward that rounds the routing weight to
bf16 measures 1.31-1.34 and one that shrinks the routing gradient by
0.5% 3.0-3.3; both are caught.
- The shrunk Qwen3.5 routed MoE: output and loss bit-identical to the key
off, gradients within relative L2 1e-2: plain and pipelined loop,
dropless, "count" with padding excluded, saved routing maps, prefused
weights with the whole wi and the fused SwiGLU VJP, and the benchmark's
MoE settings with every compatible key of the series on against all
off. The compiled rematerialized backward has exactly
num_moe_token_chunks GEMM loops fewer (19 -> 17 for 2 chunks), the
pipelined forward keeps its four barriers, and the next chunk's routed
tokens are read from the op's tie output.
- On an expert 4 x tensor 2 mesh the layer runs without the key and
raises at trace time when the key reaches it past the validator.
- A fused backward that drops the routing weight is caught at the op and
in the layer; the validator rejects the key for each unmet requirement,
naming it.
…kv_dtype) The GDN kernel path's Conv1D + SiLU (kernels/gdn/gdn_bwd/compute_conv1d_bwd.py) builds its halo-extended input [conv_state | qkv] in float32 and slices the gdn_conv_kernel_dim taps out of it. On the Qwen3.5-397B RL trainer that is an f32[1, 16387, 12288] = 805.5 MB copy per GDN layer in the backward pass, written once and read once per tap. With gdn_conv1d_input_in_qkv_dtype=True the input stays in qkv's dtype and each tap is cast to float32 where it is read, so XLA fuses the cast into the consumer and the float32 copy is never written. No optimization barrier. The option reaches the Conv1D through gdn_decoupled_conv1d's new static argument conv_input_in_qkv_dtype, which scopes it around the op's primal, forward and backward rules (the backward rule is traced after the model's forward returns, so only the custom_vjp's static arguments reach it). Numerics: the same float32 values in the same order, so op by op every output is bit-identical. Compiled, the forward outputs and the qkv, conv-state and bias gradients stay bit-identical; the conv weight gradient (a float32 reduction over batch x sequence) can be summed in a different order once XLA fuses the cast into it (CPU: 11 of 64 elements, <= 5.1e-6 relative, in 2 of 16 configurations). Default off. Requires use_gdn_kernel=True (validated).
Follow-up to the six commits of this PR (moe_routing_index_math, moe_unsplit_prefused_wi / moe_ragged_sort_unique_pair_argsort / moe_prefused_swiglu_vjp, moe_roe_packed_routing_gather, moe_wi_gmm_fused_swiglu, moe_router_grad_from_wo_input, gdn_conv1d_input_in_qkv_dtype). No new keys and no behaviour change. Refactors: - gmm_v2 kernel: inner_kernel's 2D store is now the helper _store_group_rows, which the activation output and the return_pre_act gate and up halves all call, instead of the halves using a copy of the store; the 3D stores share its row mask and partial-row carry (_mask_group_rows, _carry_partial_rows). - GDN api.py: gdn_decoupled_conv1d and its forward and backward rules keep their bodies at the original indentation. A small decorator enters conv_input_in_qkv_dtype_scope from the static argument for the primal and forward rule; the backward rule runs its argument shim, then its body (_gdn_decoupled_conv1d_bwd_impl) inside the scope. The api.py diff against the parent shrinks from 848 lines to the real changes. - moe.py: _use_prefused_swiglu_vjp inlined (it only returned the config flag); two pointer comments dropped. - types.py: the MoE validators' requirement helpers come before the validators. Validators and guards: - Both chunk loops now refuse moe_router_grad_from_wo_input with chunk weight-gradient accumulation at trace time, before any work: the pipelined loop had no such check, and the plain loop's came after the up projection, which can raise a different error first (the config validator rejects the pair in both cases). - gdn_conv1d_input_in_qkv_dtype is read directly instead of through getattr with a default, in its validator and in model_runner (gdn_cp_head_sharded_test.py's SimpleNamespace config gains the field). Tests: - moe_router_grad_from_wo_input_test.py: chunk_wgrad_guard, both chunk loops raise with moe_accumulate_chunk_wgrad past the validator. The op precision check also bounds both paths' output and gradients against the float32 reference (relative L2 1e-2, measured at most 2.4e-3): the two paths share the unsort helpers, so a defect there leaves the fused/two-ops ratio at 1, and the forward on the dropping buffer was unchecked. - moe_wi_gmm_fused_swiglu_test.py: GmmSiluVjpTest also with heuristic tiling (use_gmm_v2_heuristic_tiling must reach gmm's backward). - moe_glue_rewrites_test.py and moe_routing_index_math_test.py: validator cases for tensor_sequence and autoregressive parallelism, for every minimal* remat preset, and an error that names two unmet requirements. - Test docstrings no longer cite history or TPU runs that did not happen (megablox_gmm_silu_test.py is TPU-only and not yet run on TPU); each _BENCH_FLAGS is commented as the subset of the benchmark's MoE settings that file needs; one-line docstring summaries. - gdn_conv1d_input_in_qkv_dtype_test.py: dict literals, a docstring for _conv_inputs, the repo's license header. Docs and comments: - base.yml, types.py and the compute_conv1d_bwd.py comment describe gdn_conv1d_input_in_qkv_dtype without model-specific numbers, and say the compiled bit identity was measured on CPU. - nnx_decoders.py and types.py point at each other for the MLP checkpoint names of the remat presets (_REMAT_PRESET_MLP_NAMES), so the two stay in sync. Lint (pylint 3.3.8, pyink 24.10.1, codespell 2.4.1, yamllint 1.35.0, repo config): the changed lines add no messages against the base; the 3 pyink hunks and 5 pylint messages of the GDN test are gone. Program identity against the previous head: unchanged, by normalized TPU v7x lowering, except where intentionally changed (no change was meant to alter a program).
…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.
_run_local_gdn_decoupled_fwd called fused_conv1d_gdn with the default zero_initialize_out=True, which aliases a zero buffer into the kernel output. XLA hoists the zeros out of the layer scan into a loop-carried buffer and copies it before every kernel call, and the GDN recompute in the backward broadcasts it again. The prefill kernel never reads that buffer and writes every row (seq_len is a multiple of chunk_size, which both functions already check), so the fill is dead work. On the 397B v7x bench (16,384 tokens and 64 value heads per device) this removes six 256 MiB copies per forward layer-block iteration (about 1 ms per iteration) and three zero broadcasts per backward iteration, and frees the 256 MiB loop-carried buffer. Outputs are bit-identical.
#5395 removed enable_gdn_sequence_packing so that packing=True and use_gdn_kernel=True alone select the segment-aware GDN kernel. The kernel, though, keys on whether segment IDs are present, and unpacked data still supplies them (all ones, or a 0/1 padding mask), so unpacked training now runs the segment path too. At the Qwen3.5-397B v7x training cell the segment path cost ~2 s of a 165 s step; with unpacked rows it adds work and changes nothing, because the padding mask applied before the kernel already makes padded positions a no-op (outputs agree to fp32 rounding). Drop the segment IDs before the kernel in training when packing=False, after the padding mask. Prefill and decode keep them: there the segment path also ends the carried conv state at the last valid token. packing=False now tells the GDN kernel that every training row holds one sequence, so reject the inputs that pack rows without packing=True: - RL with max_seq_token_per_tpu > 0 (Tunix sequence packing); - grain mmap/mmap_npy with reset_attention_mask, which splits rows into documents at EOD without reading packing; - a MaxTextTrainingEngine payload carrying segment_ids, which Tunix sets only on sequence-packed rows; - a distillation teacher whose packing differs from the student's, since the teacher consumes the student's batches.
… attention query kernels
Rohan-Bierneni
force-pushed
the
rbierneni/mlperf-397b-trainer-perf
branch
from
October 15, 2026 23:10
94c69b8 to
ff76561
Compare
cp-as-ep shards the dense weights (attention and shared-expert projections,
the MoE router and the shared-expert gate) over fsdp x context x expert, so
every all-gather of them spans all devices and synchronizes the context ranks,
which do unequal attention work. The two new rules keep cp-as-ep's
activations, vocabulary weights and routed experts and take those weights off
context: cp-as-ep-dense-no-cp shards them over fsdp x expert,
cp-as-ep-dense-fsdp over fsdp only. Their all-gathers then run in independent
groups of 64 or 32 devices on the 256-device mesh, at the cost of replicating
the weights and their gradients over the axes they leave.
- RoutedMoE keeps ("context", "expert") as its expert-parallel axes under both
rules, as it does for cp-as-ep.
- Left to GSPMD, the gradient of a weight replicated over context is
all-reduced at full size over every device and then sliced. Under a rule set
that maps the new logical axis embed_grad, linears.route_dense_weight_grad
reduce-scatters it over every token axis instead, as under cp-as-ep, and
all-gathers the shard over the replicated axes. An optimization barrier
between the two keeps XLA:TPU from folding them back into the all-reduce.
DenseGeneral and GateLogit apply it to kernels with an embed or embed_attn
dimension.
- The GatedDeltaNet projections take their embed axis from embed_gdn when the
rules map it. Both new rules keep them on fsdp x context x expert: their
all-gathers are already overlapped, and replicating them would cost far more
HBM than the rest of the dense weights.
Rule sets that map neither embed_grad nor embed_gdn compile to the same
program as before.
The CPU test runs a tiny hybrid Qwen3.5 on 16 devices (fsdp 2, context 4,
expert 2) under cp-as-ep and both rules: exactly the dense weights move off
context, the loss is bit-identical, and the gradients agree within bfloat16
reduction order (1e-5 with float32 activations). A rule without the RoutedMoE
change and a perturbed weight both fail the check.
Rohan-Bierneni
force-pushed
the
rbierneni/mlperf-397b-trainer-perf
branch
from
October 16, 2026 00:50
ff76561 to
d66a820
Compare
Rohan-Bierneni
pushed a commit
to Rohan-Bierneni/tunix
that referenced
this pull request
Oct 16, 2026
…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.
At d66a820, the MoE, GDN and model files of this branch come from #5758, which is built on all of #5698. But this branch ports only part of #5698 onto atwigg/mlperf. Several files it needs are still the atwigg/mlperf versions, so the trainer does not start: - utils/sharding.py lacks CompoundLogicalAxis, so `import maxtext.layers.moe` raises ImportError; - utils/max_utils.py lacks is_eval (moe.py); - models/weaver.py is missing, and common_types.py lacks DecoderBlockType.WEAVER, which nnx_decoders.py maps for every decoder; - kernels/megablox/pallas_mosaic_tpu_v2_tgmm_kernel.py builds GmmConfigs without transpose_rhs, which the gmm_v2 kernel requires: TypeError when the MoE backward compiles. This brings src/maxtext to the tree it was developed and measured on: #5758 (20c877a, on #5698 + #5718) plus this branch's own commits (the GDN segment-path skip, the GDN forward zero-fill skip, gather_scanned_kernels_once and the cp-as-ep-dense rules), which apply to it cleanly.
parambole
added a commit
to Rohan-Bierneni/tunix
that referenced
this pull request
Oct 16, 2026
…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.
Rohan-Bierneni
added a commit
to Rohan-Bierneni/tunix
that referenced
this pull request
Oct 16, 2026
…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.
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.
Summary
Consolidates the 10 highest-ROI, easiest-to-review trainer optimizations across the competing Qwen3.5-397B-A17B MLPerf RL optimization PRs (#5667, #5721, #5741 / #5740, #5758, and #5759) cleanly on top of
atwigg/mlperf(e526ea02d).To keep code review fast, low-risk, and surgical, all invasive / low-ROI refactors were culled (see What Was Removed / Culled below). Every remaining change is self-contained, preserves bit-for-bit or statistically identical numerics (
0MoE dropped tokens), and is either a zero-flag-needed surgical fix or gated behind a clean boolean config flag (default: Falsewhere applicable).Kept Optimizations (Ordered by ROI & Review Simplicity)
32×2×4)nnx.DiffStatefilter inmaxtext_engine.pywhentrainable_parameters_maskis set, skipping VJP through frozenrouted_experts/gate/kernel)c650199a3)01f2d3502+6 / -1trainable_parameters_maskis setquerykernel once per scanned layer stack (_StackPreparedParam+ custom VJP beforennx.scaninmoe.py,nnx_decoders.py)55dab682b..4f41ede3c+84netxor_dnc+gdn_chunk_size=128support (pltpu.repeat+ elementwise XOR masks incompute_gdn.py/gdn_bwd, plus fullC=128Pallas backward tiling)268213d68xor_dnc=Trueby default (gdn_chunk_size=128optional)2 × num_moe_token_chunksseparate[B, S/C, K]routing all-gathers with 1 packed[B, S, K, 2]bf16 all-gather per layer, tagged undermoe_routeremat policy)c88674bbc)e93800779+116moe_roe_packed_routing_gather=True(defaultFalse)_prepare_prefused_weights+quantize_routed_expert_weightsbeforennx.scaninstead of 2× per layer in fwd + remat bwd)38240301f)a6bdc6615+117prepare_moe_weights_once_per_step=True(defaultFalse)qkv.dtype(bf16) (avoids f32 widening before theppermutehalo exchange and Conv1D kernel)cb709c93c)8e749332d+40gdn_conv1d_input_in_qkv_dtype=True(defaultFalse)o_hbm/s_hbmis unconditionally written by the Pallas kernel)d5b1be0df)eeac81487+3 / -0gate/upweight gradients (jnp.concatenatein_SplitGateUpFromFusedcustom VJP instead of two full-sized zero-pads + elementwise add across 60 MoE layers)0b7cdb7f0)8ebbe9327+41 / -1prefuse_moe_weights=True[512, 4096, 2048]zero-pad + add ops in bwdtop_konly on fallback tokens (routed_experts < 0) under forced routing (1D gather of fallback rows beforelax.top_kinstead of full[N, 512]top_kwhen>90%of tokens have rollout-replayed experts)464767398)69db483fe+36 / -2lax.top_kon>90%of tokens per layerring_ragged_unsortbackward (custom VJP reuses savedsort_indicesviaring_ragged_sortinstead of recomputingjnp.argsort(inv_sort_indices))bb836e7e6)0cd7c4d95+23 / -5argsortkernels per MB in bwdWhat Was Removed / Culled (Compared to Earlier Drafts & Competing PRs)
To minimize review burden and merge risk, the following heavier or lower-ROI code changes were excluded from this PR:
#5721:gdn_cp_num_passes=2): Added multi-file state propagation plumbing incp_gdn.py/wrapper.py, but in the 128-chip trainer (cp=4,S_local=16384) the single-passB×4flattening trick already runs at268 chunks >= 256, yielding only~10–15 ms/MBgain.[B, S, H, D]layout across GDN forward/backward (use_natural_layout): Large multi-file Pallas kernel refactor (compute_gdn.py,api.py,pallas_mosaic_tpu_bwd.py) for only~35 ms/MBnet gain aftergdn_conv1d_input_in_qkv_dtype=true.maxtext_engine.py(defer_grad_norm_and_loss_all_reduce,defer_moe_load_metrics_all_reduce): Culled (~350lines inmaxtext_engine.py) becausemaxtext_engine.pyonatwigg/mlperfalready providesprepare_fwd_bwd_inputandwait_for_devicefor host-staging overlap (STAGE_BEFORE_UPDATE=on), and most of the step-boundary saving comes from overlapping next-step H2D staging rather than custom metric deferral.remat_Saveableplumbing (gdn_fwd_bwd_Input,gdn_conv_Output,moe_pkg_Input,attention_Input): Culled; only the 1-linead_checkpoint.checkpoint_name(..., "moe_route")tag needed bymoe_roe_packed_routing_gatheris kept.#5758/#5759:moe_roe_Async/ customall_gather_AsyncPallas barrier-semaphore primitive (#5758a744dc661, ~785 lines of Pallas/Mosaic semaphore code).cp_gdn.pyfused Conv1D + L2-norm + recurrence rewrite (#5759e068d8017).Measured Results (with google/tunix#2884, RL trainer bench, 9 micro-batches × 64 rows per step, steps 1–4 average)
fsdp×ep×cp)4x4x8)32×2×44x8x8)32×2×84x8x8)64×1×8moe_roe_packed_routing_gather=false)Measured with
custom_mesh_and_rule=cp-as-ep, on the code without the top commit (the earlier 2-linecp-as-ep.ymledit was never applied in these runs). The top commit is now #578320c93dc4f(cp-as-ep-dense-no-cp/cp-as-ep-dense-fsdprules +route_dense_weight_grad+embed_gdn), replacing that edit;cp-as-ep.ymlis unchanged. See the re-run table below forcp-as-ep-dense-no-cp. 256-chip runs addquery_proj=device key_proj=device value_proj=device out_proj=device qkv_proj=device. Loss matches across all three runs to ~1e-8 (0.03044535), with 0 MoE drops and 0 dropless replays.Note:
moe_roe_packed_routing_gather=truerequiresmoe_chunk_pipeline=true; withnum_moe_token_chunks=1(64×1×8) it must befalse(the tunix recipe sets this automatically).Re-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.Companion Tunix PR
paired with google/tunix#2884, which adds:
pack_pad_position=start(tunix#29131435de0,-1.24 s/step)TUNIX_LOGPS_SHARDED_HEAD=fused+COMPUTE_LOGPS_CHUNK_SIZE=4096(tunix#291386e17b5,-0.55 s/step)mlperf_gbs1024_submission.sh) and 256-chip (mlperf_397b_256_v7x.sh) recipe defaults.