Repository navigation
Conversation
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).
![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 advanced performance optimizations and configuration keys for Mixture of Experts (MoE) and GatedDeltaNet (GDN) layers in MaxText. Key MoE enhancements include keeping prefused weights unsplit during quantization and FSDP all-gather (moe_unsplit_prefused_wi), fusing gated activations (moe_prefused_swiglu_vjp and moe_wi_gmm_fused_swiglu), packing routing gathers (moe_roe_packed_routing_gather), and optimizing backward passes to avoid down-projection GEMM recomputation (moe_router_grad_from_wo_input). For GDN, it adds support for a 128 chunk size, a new triangular-inverse algorithm (xor_dnc), and an option to keep halo-extended inputs in their native activation dtype (gdn_conv1d_input_in_qkv_dtype) to save HBM. Additionally, it integrates support for freezing parameters via trainable_parameters_mask in the training engine. Extensive unit tests are added to validate these exact mathematical rewrites and optimizations. I have no feedback to provide as there are no review comments.
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.
parambole
added a commit
to google/tunix
that referenced
this pull request
Oct 15, 2026
… join TRAINER_TUNED_FLAGS TRAINER_TUNED_FLAGS=true now selects the trainer cell that measured fastest on the RL trainer bench with every numerics gate passing. In mlperf_397b_256_v7x.sh the block of MaxText keys (part 2) gains sixteen keys and loses one: - wi_tile_dlhs_mlp_dim=2048, decoder_layer_input=device, moe_routing_maps=device and moe_gmm_v2_dlhs_transpose_rhs=true: the K tile of wi's input-gradient GMM matched to the fused gate and up weights, each decoder layer's input kept in HBM, the backward reusing the forward's MoE routing maps, and the input-gradient GMMs reading the expert weights through the kernel's transpose_rhs. - moe_routing_index_math=count, moe_unsplit_prefused_wi=true, moe_roe_packed_routing_gather=true, moe_router_grad_from_wo_input=true and moe_wi_gmm_fused_swiglu=true: the MoE keys of AI-Hypercomputer/maxtext#5758 in the set its commits recommend (the routing index maps by counting instead of argsorts, the fused wi kept whole up to the GEMM, one packed routing all-gather per layer, the routing-weight gradient from the down projection's input, the SwiGLU in the wi GEMM's epilogue). - gdn_conv1d_input_in_qkv_dtype=true and gdn_conv1d_dense_token_stack=true: the GDN Conv1D backward without its f32 copy of the input, and the forward kernel's Conv1D output built from dense per-token views. - wi_tile_fwd_batch_seq, wi_tile_drhs_batch_seq, wo_tile_fwd_batch_seq, wo_tile_dlhs_batch_seq and wo_tile_drhs_batch_seq=512 (from 256): row tiles of 512 for the five MoE GMMs with 1024 tiles. wi_tile_dlhs_batch_seq stays 256 next to its 2048 K tile. - sa_use_base2_exp=false leaves the block, so the trainer keeps the recipe's sa_use_base2_exp=true: base-2 exponentials measured 0.025 s faster per micro-batch (o275x-gx-b2e), and their numerics run passes (nxgx-b2e). In mlperf_gbs1024_submission.sh, which turns TRAINER_TUNED_FLAGS on, the switch also sets three defaults, each still overridable: - COMPUTE_LOGPS_CHUNK_SIZE=4096 (was 1024), and 1024 when TUNIX_LOGPS_SHARDED_HEAD turns the sharded head off (0, false, off or no, in any case): chunks of 4096 were measured with the head only, and without it the log-prob program failed to compile on TPU at chunks of 2048 and 4096. - TUNIX_LOGPS_SHARDED_HEAD=fused (was local). TRAINER_TUNED_FLAGS=false keeps local. - PACK_PAD_POSITION=start: each packed row puts its padding first. Bench (AI-Hypercomputer/maxtext#5701 driver, Qwen3.5-397B-A17B, 128 TPU v7x, mesh fsdp 32 x context 4 x expert 2, ragged_buffer_factor 2.75, 9 micro-batches per step, mean of steps 1 to 4, no token dropped in any arm): - before these keys, the fused head and pad first: 106.30 s per step (o275x-hd), 2,774 tokens/s/chip; - this cell: 82.671 and 82.670 s per step in two pods (o275x-c9pfxhv-l3), 3,567 tokens/s/chip. Part of the difference is the staged step boundary (AI-Hypercomputer/maxtext#5718), which this recipe already turns on and o275x-hd did not run; per regular micro-batch (the median of micro-batches 3 to 9), which the boundary does not change, about 10.77 -> 9.01 s; - within it, each against a control in the same pod: padding first 83.285 -> 82.043 s per step (-1.242 s) and the fused head 83.253 -> 82.708 s per step (-0.545 s). The recipe lists the measurement of each key. The measured build switched the MaxText changes with environment variables, and it also had one attention change on that neither MaxText PR carries (attention's load-balancing permute hoisted; alone 99.504 -> 99.267 s per step). The keys themselves, on the MaxText change's own code (2d82c5199) without gdn_conv1d_dense_token_stack, the fused head and pad first, measured 85.295 and 85.251 s per step (o275x-prc-l3-prof, o275x-prc-l3-r), within 0.23 s of what the environment-switch arms predicted. Both arms set moe_route=device, which nothing in #5698 or #5758 tags, so moe_route stays out of the block as before; and the stop of the frozen router's gradient, an environment switch in the measured build, is unconditional in #5758, with no key. Numerics (loss gates off, random advantages, three steps, against two reference runs whose gradient norms differ by 1.9e-5 relative): all six gates pass for this cell (nxc9pfxhv) and for the keys on the MaxText change's own code (nxprc). moe_wo_gmm_fused_dswiglu stays out: next to the fused SiLU GMM at these tiles it computed a wrong SwiGLU gradient, and its fix is still being verified on TPU. moe_routing_maps and moe_gmm_v2_dlhs_transpose_rhs exist in AI-Hypercomputer/maxtext#5698, and the seven other new MaxText keys in the MaxText change of AI-Hypercomputer/maxtext#5758 (on top of #5698); none of the nine is defined at this tree's pin (e526ea02), and MaxText refuses a key it does not define when the trainer starts. TRAINER_TUNED_FLAGS=true, the default of mlperf_gbs1024_submission.sh, therefore needs requirements/maxtext_requirements.txt moved to a MaxText with #5758 before an image built from this tree runs it. TRAINER_TUNED_FLAGS=false and the other 397B v7x recipes pass none of the nine. Tests (mlperf_submission_launcher_test.py): the exact-set test expects 23 keys, _MAXTEXT_DEFAULTS has the nine new keys' defaults, and the switch-off and other-recipe tests check that none of the nine is passed; the default launch has base-2 exponentials on, chunks of 4096, the fused head and --pack_pad_position=start; TRAINER_TUNED_FLAGS=false keeps the local head and the padding at the end; an explicit pad position, head or chunk size wins; TUNIX_LOGPS_SHARDED_HEAD=0, false, OFF or No gives chunks of 1024; the other 397B v7x recipes pass no pad position. Fifteen deliberately wrong copies of the recipes each fail at least one of these tests: a key dropped from the block or given the wrong value, one tile left out, wi_tile_dlhs_batch_seq raised to 512, sa_use_base2_exp=false kept, moe_wo_gmm_fused_dswiglu or moe_route=device added, a new key passed outside the block at its default, the local head under the switch, the chunk fallback removed or made case-sensitive, chunks of 1024 by default, pad first outside the switch, pad_align passed, and both recipes as they were before this change.
parambole
marked this pull request as ready for review
October 15, 2026 17:42
parambole
requested review from
NuojCheng,
RissyRan,
aireenmei,
bvandermoon,
darisoy,
gagika,
gobbleturk,
huytransformer,
igorts-git,
jiangjy1982,
michelle-yooh,
richjames0,
shralex,
shuningjin,
xibinliu and
zxhe-sean
as code owners
October 15, 2026 17:42
parambole
requested review from
A9isha,
SurbhiJainUSC,
abhinavclemson,
dipannita08,
hengtaoguo,
khatwanimohit and
vipannalla
as code owners
October 15, 2026 17:42
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
2 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.
Stacked on
parambole/rl-trainer-4600-base(#5698's head3792fdef8, the MaxText the RL trainer bench image runs, merged withatwigg/mlperfe526ea0), so the diff shows only this work. Retarget toatwigg/mlperfonce #5698 lands. Companion tunix PR: google/tunix#2913.Nine opt-in config keys, all off by default (seven recommended, plus two alternatives in
048d61701), that speed up the Qwen3.5-397B-A17B RL trainer on 128 TPU v7x. On the RL trainer bench (#5701), this PR's code with the companion recipe's MaxText keys exceptgdn_conv1d_dense_token_stacktakes 85.30 s per step, and the full companion tunix recipe 82.67 s, against 106.30 s before this work: 2,774 -> 3,567 tok/s/chip, with all six numerics gates passing. The 106.30 s cell also lacked the already-merged staged step boundary (#5718); per regular micro-batch, which the boundary does not change, the time goes from 10.77 to 9.01 s. The 82.67 s cell ran on an overlay build of the bench image's MaxText that also had one attention change, which follows in a separate PR, and lacked the merged #5568 (see Results).Where each change acts
One decoder layer: full attention in 15 of the 60 layers, GDN in the other 45, and the MoE block in every layer. Blue: changed by this PR, labelled with the key that turns it on. Gray: unchanged.
flowchart TB classDef pr fill:#dbeafe,stroke:#1d4ed8,stroke-width:2px,color:#0b1f4d classDef same fill:#f3f4f6,stroke:#9ca3af,color:#6b7280 subgraph MIX["token mixer: attention or GDN"] ATT["full attention"]:::same CONV["GDN Conv1D + SiLU<br/>gdn_conv1d_dense_token_stack<br/>gdn_conv1d_input_in_qkv_dtype"]:::pr DR["GDN gated delta rule<br/>xor_dnc, always on"]:::pr CONV --> DR end subgraph MOE["MoE block"] RT["routing<br/>moe_routing_index_math=count<br/>moe_roe_packed_routing_gather"]:::pr WI["wi GMM + SwiGLU<br/>moe_unsplit_prefused_wi<br/>moe_wi_gmm_fused_swiglu"]:::pr WO["wo GMM + unsort<br/>moe_router_grad_from_wo_input"]:::pr RT --> WI --> WO end ATT ---> RT DR --> RT style MIX fill:none,stroke:#9ca3af style MOE fill:none,stroke:#9ca3afChanges
moe_routing_index_math=countmoe_roe_packed_routing_gathero275x-mf-ctl; -0.096 againsto275x-gdn-ctl)moe_unsplit_prefused_wimoe_wi_gmm_fused_swiglufuse_act="silu")moe_prefused_swiglu_vjpin its place)moe_router_grad_from_wo_inputgdn_conv1d_input_in_qkv_dtypegdn_conv1d_dense_token_stackjnp.stackrows, which the October 2026 TPU compiler pads from 1 to 8 sublanes (older compilers keep them dense)jnp.stackGains are against a control in the same pod unless marked (for
gdn_conv1d_dense_token_stack, a same-pod arm whose only change, a 2D GDN out_proj, is a no-op); mb = micro-batch, 9 per step. Also in the branch, with no key:dede90a9f: the GDN triangular inverse with xor_dnc, Rohan's Optimize GDN triangular inverse with xor_dnc and support gdn_chunk_size=128 #5667 cherry-picked as a dependency (drop it when Optimize GDN triangular inverse with xor_dnc and support gdn_chunk_size=128 #5667 merges). -0.191 s/mb (cross-pod), TPU A/B PASS 6/6.7f5dcc7ce:stop_gradienton frozen parameters insidefwd_bwd; the RL recipe freezes the router. -0.058 s/mb, both arms withmoe_routing_maps=device(cross-pod), TPU A/B PASS 6/6.8cc954eb2: review feedback. It lowers the MoE layer, the gmm_v2 kernel and the GDN op to the same TPU programs as the commit before it (see Testing).Results
RL trainer bench (#5701), 128 TPU v7x. Seconds per optimizer step, mean of steps 1-4.
o275x-hd)o275x-prc-l3-profand its repeato275x-prc-l3-r, one pod)nxprc)o275x-c9pfxhv-l3, two pods)nxc9pfxhv)2d82c5199(same MoE-layer program as the published commits, GDN and model sources unchanged; see Testing), with the keys set as config: the companion recipe's key set withoutgdn_conv1d_dense_token_stack, pad first and the fused log-prob head. It includes the six recommended earlier keys of this PR (048d61701's two alternative keys were off), the GMM tiles at 512 and log-prob chunk 4096.gdn_conv1d_input_in_qkv_dtypewas process-wide in that build and is scoped here, but the GDN layer program is the same (see Testing). The stop_gradient on frozen parameters has no such check. That build also hoisted attention's load-balancing permute, which is in neither PR (alone, against another pod's control: 99.504 -> 99.267 s per step, regular micro-batch 10.033 -> 10.020 s).o275x-c9ypfxhv-l3), took 82.945 s per step (3,556 tok/s/chip; regular micro-batch 9.04 s). That pod's overlay build also contained the follow-up dSwiGLU fix, with its key off.moe_route=device, which the recipe leaves at its default; no code on this branch tags a tensor with that name.How to turn it on
Each MoE key's validator refuses a config in which the key would do nothing and names every unmet requirement.
gdn_conv1d_dense_token_stackis refused without Gated DeltaNet layers, withoutuse_gdn_kernel, or when the Conv1D width each kernel call convolves is not a multiple of 1024, one requirement at a time;gdn_conv1d_input_in_qkv_dtypeis refused only foruse_gdn_kernel=Falseon a GDN decoder. The companion recipe,TRAINER_TUNED_FLAGS=truein google/tunix#2913, sets these seven together with its other MaxText keys; the 82.67 s cell measured those settings on the overlay build (see Results).Testing
8cc954eb2; the dense-stack test, and the Conv1D dtype test again, on the tip20c877aa1. The gmm SiLU kernel test (megablox_gmm_silu_test.py) needs TPU and has not run there yet.7f5dcc7ce).cb709c93clowers to the same program as the pre-review tree that ran row 2 andnxprc, in 45 of 45 cases (9 key sets in 5 variants), with the GDN and model sources unchanged between the two;8cc954eb2lowers it identically (16 of 16) and the tip changes no MoE source. Withgdn_conv1d_dense_token_stackon, the GDN layer equals the overlay's switch applied to8cc954eb2in 18 of 18 configurations. Against the measured overlay tree itself, with its switches, the tip lowers the MoE layer (the recipe's five MoE keys) and the GDN layer (both GDN keys, andgdn_conv1d_input_in_qkv_dtypealone, at the 397B head counts with sequence-sharded context parallelism 4; both keys also without context parallelism) to the same programs. Negative controls differ, as they must.8cc954eb2;gdn_cpalso on the tip): GDN backward tests 14 passed, 1 skipped;gdn_cp7 passed; in a 31-file regression subset, run oncb709c93c(before the review-feedback and dense-stack commits), every failure also fails at7f5dcc7ce(the tree before the keyed commits).nxprc) and for the full recipe (nxc9pfxhv). The gates compare scalars of the first of the run's 3 steps (the later two are logged, not gated), so they catch gross errors, not a wrong gradient direction; exactness rests on the CPU bit-identity tests, and the two rounding-level keys have finer checks (details below).alogsix times ingdn_conv1d_dense_token_stack_test.py(read as "along"); the rename follows in a separate PR. pyink reformats 4 hunks, all in lines of the cherry-picked Optimize GDN triangular inverse with xor_dnc and support gdn_chunk_size=128 #5667 (dede90a9f); the other pyink hunks in the changed files are already at base.Not in this PR (follow-up PRs)
moe_wo_gmm_fused_dswiglu(SwiGLU backward in the wo DLHS GMM epilogue) and its layout fix, in a follow-up PR.wi_tile_fwd_mlp_dim1024 the forward kernel halves its N tile to 512, so the array was laid out per tile as [g0|u0|g1|u1], and the DLHS epilogue split it in half: a wrong gradient. The key's own commit refuses that tile setting at config time; the fix commit reads the per-tile layout. This PR's fused SiLU writes gate and up as separate arrays and is not affected.nxc9ypfckhv, PASS 6/6 (these gates also passed the unfixed version, see Numerics, so the kernel check is the evidence for the fix).o275x-c9pfxctl-l3-> 82.148 ino275x-c9ypfctl-l3, two pods).How it was measured: bench, timing, logs, profiles
ragged_buffer_factor2.75, with 0 dropped tokens in every arm.r-4e5a9/parambole-bxa-4e5a9; rows 2-3:r-5f2bf/parambole-bpr-5f2bf,parambole-bvx2-5f2bfandparambole-bvx4-5f2bf; the intermediate arm:parambole-bvx-5f2bf). The same control measured in different pods spreads by at most +/-0.01 s per micro-batch:o275x-hdin four pods took 10.7735-10.7845 s per regular micro-batch and 106.04-106.30 s per step.o275x-c9pfxctl-l3, same overlay build): 83.32 s per step. Per regular micro-batch the chain is 10.77 -> 9.30 -> 9.08 -> 9.01 s. Between row 2 and this arm,gdn_conv1d_dense_token_stackand tunix pad first are added and the build changes too (overlay build: the hoist on, perf(training_engine): buffer microbatch metrics in Python lists and gate checkpoint metrics on intra-step state #5568 absent).gs://parambole-397b-replay-nap/win-tpu/sched/<booking>/<job>/<arm>.log, for example:r-4e5a9/parambole-bxa-4e5a9/o275x-hd.logr-5f2bf/parambole-bpr-5f2bf/o275x-prc-l3-prof.logando275x-prc-l3-r.log; numericsr-5f2bf/parambole-bnpr-5f2bf/nxprc.logr-5f2bf/parambole-bvx2-5f2bf/o275x-c9pfxhv-l3.log, replicater-5f2bf/parambole-bvx4-5f2bf/o275x-c9pfxhv-l3-r.log, third podr-2d4c2/parambole-bn3s-2d4c2/o275x-c9ypfxhv-l3.log; numericsr-5f2bf/parambole-bvx3-5f2bf/nxc9pfxhv.logo275x-prc-l3-prof, under.../r-5f2bf/parambole-bpr-5f2bf/o275x-prc-l3-prof/.../tensorboard/plugins/profile/.o275x-c9pfxall-l3-prof, 82.787 s per step; the recipe arm itself ran without a profiler):.../r-5f2bf/parambole-bvx-5f2bf/o275x-c9pfxall-l3-prof/parambole-bvx-5f2bf-o275x-c9pfxall-l3-prof/tensorboard/plugins/profile/2026_10_15_14_43_36/server-gke-tpu-174635e9-sj65.0.0.xplane.pb.o275x-hd-prof.o275x-*/nx*arm ID maps to a log at the path above. Arms ending in-profalso have an xplane profile under.../<job>/<arm>/<job>-<arm>/tensorboard/plugins/profile/.Numerics: the six checks, per-run results, and what they cannot see
The bench's gates-off A/B, run for every kept change: log-prob gates off, random advantages, 3 steps, against a reference run and a replicate of it. Six checks: every arm keeps its tokens; loss, relative difference; mean log-prob difference; last-micro-batch log-prob difference; entropy, relative difference; gradient norm, relative difference.
nxprc2d82c5199) with the six MoE/GDN keys of stack 8T as confignxc9pfxhvnxs6grgdn_conv1d_dense_token_stackon the environment-switch build, stacked as in its timing arm (withmoe_prefused_swiglu_vjp, not the fused SiLU); steps 0-1 of 3, cut in step 2 by a dropped proxy connectionnxs6silumoe_wi_gmm_fused_swigluon the environment-switch build, stacked as in its timing armnxgx2-c1b1gdn_conv1d_input_in_qkv_dtypealonenxt67,nxstk2stop_gradientWhat these checks cannot see. The six gates compare scalars of the first optimizer step (the run has three steps; the later two are logged, not gated), so they rule out gross errors only, not a wrong gradient direction.
moe_wo_gmm_fused_dswiglu) next to the fused SiLU at a 512-wide forward tile. That key is not in this PR.moe_router_grad_from_wo_input: an op-level precision test against a float32 reference, both paths within relative L2 1e-2, measured at most 2.4e-3.moe_wi_gmm_fused_swiglu: kernel testmegablox_gmm_silu_test.py, which needs TPU and has not yet run on TPU. Runsnxprcandnxs6siluare that kernel's TPU evidence.Tests: per-file counts, program identity, mutation tests, lint
All tests run on CPU in the bench image. Unit tests on the tree of
8cc954eb2; the dense-stack test, and the Conv1D dtype test again, on the tip20c877aa1.add4a1e74tests/unit/moe_routing_index_math_test.py048d61701tests/unit/moe_glue_rewrites_test.pyc88674bbctests/unit/moe_roe_packed_routing_gather_test.py4ab9f123dtests/unit/moe_wi_gmm_fused_swiglu_test.py;tests/unit/megablox_gmm_silu_test.py(TPU only)63cba2632tests/unit/moe_router_grad_from_wo_input_test.py, including the float32-reference precision testcb709c93ctests/unit/gdn_conv1d_input_in_qkv_dtype_test.py20c877aa1tests/unit/gdn_conv1d_dense_token_stack_test.py, including a program check in a child processRelated suites (tree of
8cc954eb2; the threegdn_cpsuites also on the tip, same results):gdn_cp: 7 passed.gdn_cp_head_shardedandgdn_cp_precision: one failing test each, the same at base and at tip. It is a pre-existingSimpleNamespaceerror atqwen3.py:897.cb709c93c(before the review-feedback and dense-stack commits): every failure there also fails at7f5dcc7ce(the tree before the keyed commits).Program identity (normalised TPU v7x StableHLO, Mosaic kernel bodies decoded, source locations stripped):
7f5dcc7ce), for the MoE layer and the GDN op on the local and context-parallel paths.8cc954eb2vs the publishedcb709c93c: MoE layer identical for 8 key sets in 2 variants (bench, retry), 16 of 16; gmm_v2 kernel identical for 2D, group offset, transpose_rhs, 3D, windowed 3D and SiLU; GDN op identical. The tip20c877aa1changes no MoE source.cb709c93cvs the pre-review tree that ran theo275x-prc-l3arms andnxprc: the MoE layer is identical for all 9 key sets in 5 variants (45 of 45), and the GDN and model sources are unchanged, so those measurements apply to this branch.gdn_conv1d_dense_token_stack: key off equals8cc954eb2, and key on equals the overlay's switch applied to8cc954eb2, in 18 of 18 configurations (162 probes). Negative controls differ, as they must.gdn_conv1d_input_in_qkv_dtypealone matches at the first. Negative controls differ, as they must.Mutation tests (tests shown to fail on deliberately broken copies):
8cc954eb2: R1, R4, R6, R8, R8p, R9, R10 and R11.gdn_conv1d_dense_token_stack: 11 of 11 mutants from the test reviewer are killed; each fails at least one test that passes on the clean tip (41 passed). They cover each context-parallel kernel call site dropping the option, the backward shim using the wrong slot, the batched kernel getting the option, a concat off the token path, the runner reading the wrong key, a bf16 round trip and a dtype guard that accepts bf16, a token shift by one, and a validator threshold.Lint (the CI pins), tip vs base: pylint 3.3.8, 0 new messages; yamllint clean; codespell 2.4.1 flags
alogsix times intests/unit/gdn_conv1d_dense_token_stack_test.py(read as "along"; the rename toa_logfollows in a separate PR); pyink 24.10.1 reformats 4 hunks, all in lines the cherry-picked #5667 adds (dede90a9f: 3 ingdn_bwd/api.py, 1 ingdn_bwd_pallas_test.py), and the other pyink hunks in the changed files are already at base. Each tool fires on a deliberately broken file.Tried and not kept (same bench)
ragged_buffer_factor2.25: not allowed; 2.75 stays, or 2.0 only with zero drops.moe_pin_sparse_core_all_gathers: +0.363.gdn_cp_mode=head: +0.294.o275x-gdn-ctl).Corrections to the published commit messages
The six key commits (
add4a1e74tocb709c93c) were published before the last review pass, so their messages cannot be changed. The follow-up commit "Address review feedback on the MoE and GDN keys" carries the code changes; the corrected wording for the messages is here.TPU numerics of this PR's own code (replaces "it ran on the environment-switch build, not on this branch's binary"). Applies to
moe_routing_index_math="count",moe_unsplit_prefused_wi,moe_roe_packed_routing_gather,moe_wi_gmm_fused_swiglu,moe_router_grad_from_wo_inputandgdn_conv1d_input_in_qkv_dtype.nxprcran this PR's code on TPU with those six keys set as config, not as environment switches. It ran the pre-review revision of these commits; a TPU v7x lowering check shows that revision lowers the MoE layer to the same program as the published commits for all 9 key sets checked, in 5 variants (45 of 45), with the GDN and model sources unchanged between the two, and the follow-up commit keeps that identity.o275x-prc-l3-prof, and 85.251 s in its repeato275x-prc-l3-r.nxprc:moe_prefused_swiglu_vjp(TPU numerics only from environment-switch-build runs:nxs6c,nxs6lc,nxs6t512andnxs6gr, each PASS 6/6) andmoe_ragged_sort_unique_pair_argsort(in no TPU run; a permutation, exact by construction, bit-identical on CPU).What the A/B can catch (add wherever the A/B is cited). These gates compare scalars of the first of 3 steps, and they have passed a run with known-wrong MoE gradients, so they rule out gross errors only. Exactness of the exact rewrites rests on the CPU bit-identity tests. The two numerics changes need more care:
moe_wi_gmm_fused_swiglu, whose kernel test (megablox_gmm_silu_test.py) is not yet run on TPU, andmoe_router_grad_from_wo_input, a rounding-level gradient change whose finer check is the op-level precision test against a float32 reference.moe_ragged_sort_unique_pair_argsortandmoe_prefused_swiglu_vjp(048d61701). Neither key is in the recommended set, and neither was timed alone. Both were timed only together with the whole wi (moe_unsplit_prefused_wi): all three keys in one arm, withmoe_routing_index_math="sort", -0.170 s per micro-batch (o275x-mf-gb-profagainst its in-pod controlo275x-mf-ctl). They are kept as alternatives for configurations the recommended set does not cover. The pair argsort is for runs that keepmoe_routing_index_math="sort". The fused activation is for runs withoutmoe_wi_gmm_fused_swiglu, which needs tokamax gmm_v2 and plain silu + linear, while the fused activation works with any GMM backend and gated activation.moe_wi_gmm_fused_swiglu(4ab9f123d).megablox_gmm_silu_test.py, the kernel's own TPU test, is not yet run on TPU. Runsnxprcandnxs6siluare the kernel's only TPU evidence._store_group_rows), not a copy of it. Withreturn_pre_actoff, the kernel lowers to the same Mosaic body as before.moe_router_grad_from_wo_input(63cba2632).Test additions from the mutant re-review.
moe_wi_gmm_fused_swiglu_test.py: the gmm_silu VJP test also runs with heuristic tiling, so the option must reach gmm's backward.moe_glue_rewrites_test.py: the whole wi is rejected under tensor_sequence and autoregressive parallelism, and the fused activation under every minimal* remat preset.moe_routing_index_math_test.py: the validator error names every unmet requirement when several are unmet.8cc954eb2: these additions kill the re-review's surviving mutants R1 (forward and backward error bound, dropping buffer), R4 (forward error 0.585 against the 1e-2 bound), R6, R9 (both parallelism cases), R10 and R11, plus the trace-time guard removed from either chunk loop. Not addressed: R5 (its killing test,moe_megatron_aux_loss_test, still needs 8 host devices to run) and R14 (needs a down-projection stand-in that writes NaN outside its groups).gdn_conv1d_input_in_qkv_dtype(cb709c93c): measurement and numerics.o275x-gx2-c1b1-prof). That pod ran no flags-off control. Against the flags-off control of another pod (o275x-gdn-ctl, 10.033 s, 99.504 s per step), that is -0.050 s per micro-batch and -0.45 s per step. The same control measured in different pods spreads by at most +/-0.01 s per micro-batch (o275x-hdin four pods: 10.7735-10.7845 s).conv1d_siluforward and backward in the process through an environment variable. This commit scopes the option togdn_decoupled_conv1d. A lowering check of one GDN layer (16 key / 64 value heads, sequence-sharded context parallelism 4, segment ids, remat) gives the same program for the build's switch and this commit's key, alone and together with the dense stack; with both on, also without context parallelism.nxgx2-c1b1, environment-switch build): all six checks pass. Loss relative difference 4.6e-8, mean log-prob difference 2.6e-6, gradient-norm relative difference 9.3e-5. That is 4.9x the replicate spread of 1.9e-5, under the 2e-3 tolerance. Runnxprcran this commit's scoped option together with the five recommended MoE keys and passes all six checks. The option is also in runnxs6c's stack.dede90a9f). Its api.py import hunk does not apply without it.