[JAX] Add an MoE Block (Layer) that compound router, permutation, groupedGEMM and communication - #2912
Merged
Merged
Conversation
tdophung
marked this pull request as ready for review
May 5, 2026 21:47
Contributor
Greptile SummaryThis PR introduces a self-contained MoE block (
Confidence Score: 5/5This PR is safe to merge. The single-VJP MoE block and its A2A helpers are logically correct, the recv_buffer_rows alignment fix is present, and the new test suite provides meaningful end-to-end coverage. The core mathematical paths (dispatch, combine, FFN bwd, routing bwd) were reviewed and verified consistent. The only new finding is two dead operations in the PURE_JAX routing-bwd that are immediately overwritten and have no effect on correctness. transformer_engine/jax/moe.py — the dead-code cleanup on lines 1414-1415 of the routing bwd is trivial but worth a pass before the file grows further. Important Files Changed
Sequence DiagramsequenceDiagram
participant Caller
participant moe_fwd as _moe_fwd_rule
participant body_fwd as _body_fwd
participant dispatch as _dispatch
participant A2A as ragged_all_to_all
participant FFN as grouped_gemm x3
participant combine as _combine
participant moe_bwd as _moe_bwd_rule
participant body_bwd as _body_bwd
Caller->>moe_fwd: x, gate_kernel, wi_0, wi_1, wo
moe_fwd->>body_fwd: shard_map or direct
body_fwd->>dispatch: inputs_2d, routing_map
dispatch->>A2A: ragged_all_to_all fwd
A2A-->>dispatch: recv tokens
dispatch-->>body_fwd: sorted_x, dispatch_state
body_fwd->>FFN: grouped_gemm
FFN-->>body_fwd: expert_outputs
body_fwd->>combine: expert_outputs, dispatch_state
combine->>A2A: ragged_all_to_all reverse
A2A-->>combine: gathered outputs
combine-->>body_fwd: output, expert_outputs_residual
body_fwd-->>moe_fwd: output, aux_loss, ctx
moe_fwd-->>Caller: output, aux_loss
Caller->>moe_bwd: d_output, ctx
moe_bwd->>body_bwd: shard_map or direct
body_bwd->>combine: _combine_bwd
combine->>A2A: ragged_all_to_all fwd params
A2A-->>combine: d_expert_outputs
combine-->>body_bwd: d_expert_outputs, d_routing_weights
body_bwd->>FFN: grouped_gemm bwds
FFN-->>body_bwd: d_sorted_x, d_wi
body_bwd->>dispatch: _dispatch_bwd
dispatch->>A2A: ragged_all_to_all reverse params
A2A-->>dispatch: d_inputs_2d
dispatch-->>body_bwd: d_inputs_2d
body_bwd-->>moe_bwd: grads
moe_bwd-->>Caller: d_x, d_gate_kernel, d_wi_0, d_wi_1, d_wo
Reviews (13): Last reviewed commit: "[JAX] CI lint: reach pylint 10.00/10 on ..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Signed-off-by: tdophung <tdophung@nvidia.com>
…ody single GPU vs. multi GPU Signed-off-by: tdophung <tdophung@nvidia.com>
Signed-off-by: tdophung <tdophung@nvidia.com>
Signed-off-by: tdophung <tdophung@nvidia.com>
…e and single device initial params in the MoEBlock. Tests should pass now Signed-off-by: tdophung <tdophung@nvidia.com>
Signed-off-by: tdophung <tdophung@nvidia.com>
for more information, see https://pre-commit.ci
Signed-off-by: tdophung <tdophung@nvidia.com>
…int in C++ files, make FP8 works. Tested with current scaling Signed-off-by: JAX Toolbox <jax@nvidia.com>
for more information, see https://pre-commit.ci
Contributor
|
Want your agent to iterate on Greptile's feedback? Try greploops. |
tdophung
commented
May 7, 2026
tdophung
commented
May 7, 2026
… grad tol to 5e-2, move arch/align_size docs into MoEBlock class Signed-off-by: tdophung <tdophung@nvidia.com>
Comment on lines
+909
to
+914
| batch_divisor = num_ep * dp_size | ||
| if global_batch_size % batch_divisor != 0: | ||
| raise ValueError( | ||
| f"batch={global_batch_size} not divisible by prod(data_parallelism_axes)={dp_size}" | ||
| ) | ||
| recv_buffer_rows = (global_batch_size // dp_size) * sequence_length * topk |
Contributor
There was a problem hiding this comment.
align_size > 0 + EP are combined
recv_buffer_rows is computed assuming unpadded token counts, but when align_size > 0 the per-expert group_sizes are the aligned counts, so the send_sizes in compute_ragged_all_to_all_params include padding tokens. The worst-case receive per shard is num_ep * ((B/(num_ep*dp_size))*S*K + num_experts_per_shard*(align_size-1)), which exceeds the current recv_buffer_rows = (B/dp_size)*S*K by up to num_experts*(align_size-1) rows. ragged_all_to_all writing beyond the buffer produces incorrect results or a crash. The correct worst-case size is:
recv_buffer_rows = (global_batch_size // dp_size) * sequence_length * topk + num_experts * (self.align_size - 1 if self.align_size > 0 else 0)
This combination (EP + align_size > 0) is not exercised by the current distributed test (which defaults to align_size=0), so the bug is latent.
phu0ngng
requested changes
May 12, 2026
phu0ngng
left a comment
phu0ngng
left a comment
Collaborator
There was a problem hiding this comment.
I think we should go with exposing GroupMLP VJP first before the MoE module to enable future possible fusions.
…ing None as group_topk, align_size rename, Signed-off-by: tdophung <tdophung@nvidia.com>
Signed-off-by: tdophung <tdophung@nvidia.com>
Signed-off-by: tdophung <tdophung@nvidia.com>
tdophung
marked this pull request as draft
May 15, 2026 16:46
Comment on lines
+568
to
+700
| # Per-shard compile-time-constant shape info (Python ints / int tuples). | ||
| # See ``_compute_static_shape_info`` and the note in ``_dispatch`` | ||
| # for why these are kwargs rather than state-dict entries. | ||
| num_real_tokens: int, | ||
| padding_size: int, | ||
| post_a2a_buffer_shape: Optional[Tuple[int, int]], | ||
| # EP-only: | ||
| ep_axis: Optional[str], | ||
| shard_id: Optional[jnp.ndarray] = None, | ||
| num_ep: int = 1, | ||
| ) -> Tuple[jnp.ndarray, Optional[jnp.ndarray]]: | ||
| """Inverse of :func:`_combine` on the cotangent. | ||
|
|
||
| Returns ``(d_expert_outputs, d_routing_weights_or_merging_probs)``. | ||
|
|
||
| ``expert_outputs`` is the *forward* output of the FFN (same value the | ||
| fwd handed to :func:`_combine`). It's required by the TRITON | ||
| combine_bwd kernel; for PURE_JAX we don't need it but accept it for | ||
| a symmetric signature. | ||
| """ | ||
| # Step 3 inverse: global combine bwd. | ||
| d_output_2d = d_output.reshape(-1, d_output.shape[-1]) | ||
| if backend is PermutationBackend.PURE_JAX: | ||
| # The pure-jax combine is: | ||
| # unsort = _sort_activations(expert_outputs, argsort(sorted_indices)) | ||
| # if pad: unsort = unsort[:num_real] | ||
| # reshape -> einsum BKE,BK -> BE -> reshape to BSE | ||
| # Hand-derive the bwd in plain JAX (no custom_vjp involved): | ||
| unsort_indices = jnp.argsort(state["sorted_indices"]) | ||
| topk = num_experts_per_tok | ||
| num_real = num_real_tokens | ||
| padding = padding_size | ||
| # Recover the unsorted intermediate that the fwd produced (we | ||
| # need it for the d_routing_weights pullback). Apply the same | ||
| # gather the fwd did. | ||
| unsort_intermediate = expert_outputs[unsort_indices] | ||
| if padding > 0: | ||
| unsort_intermediate = unsort_intermediate[:num_real] | ||
| # Bwd of einsum/reshape: | ||
| # output[B, E] = sum_K intermediate[B, K, E] * weights[B, K] | ||
| # d_intermediate[B, K, E] = d_output[B, E] * weights[B, K] | ||
| # d_weights[B, K] = sum_E d_output[B, E] * intermediate[B, K, E] | ||
| rw = state["routing_weights"].reshape(-1, topk) | ||
| intermediate_3d = unsort_intermediate.reshape(rw.shape[0], topk, -1) | ||
| rw_cast = rw.astype(intermediate_3d.dtype) | ||
| d_intermediate_3d = jnp.einsum("BE,BK -> BKE", d_output_2d, rw_cast) | ||
| d_routing_weights = jnp.einsum("BE,BKE -> BK", d_output_2d, intermediate_3d).astype( | ||
| state["routing_weights"].dtype | ||
| ) | ||
| d_routing_weights = d_routing_weights.reshape(state["routing_weights"].shape) | ||
| d_unsort_intermediate = d_intermediate_3d.reshape(num_real, -1) | ||
| # Pad back with zeros if the fwd stripped padding. | ||
| if padding > 0: | ||
| d_unsort_intermediate = jnp.concatenate( | ||
| [ | ||
| d_unsort_intermediate, | ||
| jnp.zeros( | ||
| (padding, d_unsort_intermediate.shape[-1]), | ||
| dtype=d_unsort_intermediate.dtype, | ||
| ), | ||
| ], | ||
| axis=0, | ||
| ) | ||
| # Bwd of the gather is gather-by-original-indices: | ||
| # sorted = unsort[argsort(sorted_indices)] | ||
| # d_sorted = scatter d_unsort via argsort(sorted_indices) | ||
| # = d_unsort[sorted_indices] (gather by original sorted_indices, | ||
| # which is the inverse of argsort(sorted_indices)). | ||
| d_expert_outputs_global = d_unsort_intermediate[state["sorted_indices"]] | ||
| else: | ||
| # TRITON combine bwd: requires fwd_input (expert_outputs). | ||
| num_tokens = state["row_id_map"].shape[0] | ||
| n_experts = (state["row_id_map"].shape[1] - 1) // 2 | ||
| hidden = d_output_2d.shape[-1] | ||
| num_out_tokens = expert_outputs.shape[0] | ||
| if state["pad_offsets"] is not None: | ||
| d_expert_outputs_global, d_merging_probs = unpermute_bwd_with_merging_probs_and_unpad( | ||
| d_output_2d, | ||
| state["row_id_map"], | ||
| expert_outputs, | ||
| state["merging_probs"], | ||
| state["pad_offsets"], | ||
| num_tokens, | ||
| n_experts, | ||
| num_out_tokens, | ||
| hidden, | ||
| ) | ||
| # The kernel only writes positions tokens map to; padded | ||
| # positions may contain NaN. Replace with zeros (matches | ||
| # ``_token_combine_bwd_rule``). | ||
| d_expert_outputs_global = jnp.where( | ||
| jnp.isnan(d_expert_outputs_global), 0.0, d_expert_outputs_global | ||
| ) | ||
| else: | ||
| d_expert_outputs_global, d_merging_probs = unpermute_bwd_with_merging_probs( | ||
| d_output_2d, | ||
| state["row_id_map"], | ||
| expert_outputs, | ||
| state["merging_probs"], | ||
| num_tokens, | ||
| n_experts, | ||
| num_out_tokens, | ||
| hidden, | ||
| ) | ||
| d_routing_weights = d_merging_probs | ||
|
|
||
| if not ep_active: | ||
| return d_expert_outputs_global, d_routing_weights | ||
|
|
||
| # Step 2 (EP) inverse: bwd of reverse ragged_all_to_all is a forward | ||
| # ragged_all_to_all using the SAME forward parameters (sender / | ||
| # receiver roles swap from the reverse direction back to forward). | ||
| in_off_f, send_sz_f, out_off_f, recv_sz_f = compute_ragged_all_to_all_params( | ||
| state["all_shards_tokens_per_expert"], shard_id, num_ep | ||
| ) | ||
| recv_buf_for_bwd = jnp.zeros(post_a2a_buffer_shape, dtype=d_expert_outputs_global.dtype) | ||
| d_x_send_back = jax.lax.ragged_all_to_all( | ||
| d_expert_outputs_global, | ||
| recv_buf_for_bwd, | ||
| in_off_f, | ||
| send_sz_f, | ||
| out_off_f, | ||
| recv_sz_f, | ||
| axis_name=ep_axis, | ||
| ) | ||
| # Step 1 (EP) inverse: combine fwd applied is_forward=False; the | ||
| # bwd is is_forward=True with the SAME row_id_map. | ||
| recv_buffer_rows, hidden = d_x_send_back.shape | ||
| d_expert_outputs, _ = sort_chunks_by_map( | ||
| d_x_send_back, | ||
| state["local_perm_row_id_map"], | ||
| None, | ||
| recv_buffer_rows, |
Contributor
There was a problem hiding this comment.
ctx["expert_outputs"] is the wrong tensor for the combine backward under EP
_body_fwd stores the shard-local FFN output (shape [recv_buffer_rows, hidden]) in ctx["expert_outputs"] before calling _combine. However, the global-combine step (step 3) inside _combine operates on a different tensor: the globally-permuted FFN output (shape [num_real_tokens + padding_size, hidden]) produced after the EP inverse-local-permute and reverse-ragged-A2A steps (steps 1–2). _combine_bwd then receives ctx["expert_outputs"] (shard-local, expert-major) and uses it as if it were the globally-permuted tensor.
Under PURE_JAX + EP: unsort_intermediate = expert_outputs[unsort_indices] indexes the shard-local tensor with global argsort indices; the recovered intermediate has wrong content, so d_routing_weights (and the resulting gate-kernel gradient) is silently incorrect.
Under TRITON + EP: expert_outputs is passed directly to unpermute_bwd_with_merging_probs[_and_unpad] with num_out_tokens = expert_outputs.shape[0] = recv_buffer_rows instead of num_real_tokens + padding_size; both d_expert_outputs_global and d_merging_probs are wrong, corrupting expert-weight gradients as well.
The existing distributed tests only assert finiteness and non-zero gradients; because both backends suffer the same mismatch, the backend-parity test passes even when both produce wrong gradients.
Fix: either (a) run the EP steps (inverse local permute + reverse ragged-A2A) before calling _combine / _combine_bwd, save the resulting globally-permuted tensor in ctx["expert_outputs_global_sorted"], and pass it to step 3; or (b) re-execute the EP forward steps inside _combine_bwd to recover the globally-permuted output before delegating to the step-3 backward.
Collaborator
There was a problem hiding this comment.
@tdophung is this comment valid? I see your recent commits fixed a Triton+EP issue that looks similar. Do we also need to make a corresponding fix to the pure-JAX version?
transformer_engine/jax/moe.py: - Hoist 'import math' to module top (was two local imports). - Trim the verbose _moe_fwd_rule arg-order block comment. - Update PermutationBackend docstring: TRITON is the recommended default and is faster than PURE_JAX on current hardware. - Rename layer_w0 / layer_w1 to gate_proj_out / up_proj_out so the names reflect what they are (SwiGLU projection outputs, not weights). - moe() now rejects overlapping EP / FSDP axes up front instead of letting JAX produce a duplicate-axis PartitionSpec. transformer_engine/jax/permutation.py: - Drop reference to the temporary MaxText-fork compute_ragged_all_to_all_params helpers. tests/jax/test_moe_vjp.py, tests/jax/test_multiprocess_moe_vjp.py: - Add a module-level Blackwell (sm_100+) skip; grouped GEMM is Blackwell-only today. - Move the 'triton' pytest marker from the class onto the triton parametrize variant only, so the pure_jax variant still runs in environments without Triton. Signed-off-by: tdophung <tdophung@nvidia.com>
for more information, see https://pre-commit.ci
transformer_engine/jax/moe.py + permutation.py: - Wrap triton_extensions imports in try/except so the modules still load on builds without Triton. PermutationBackend.TRITON callsites now raise a clear ImportError up-front via the public moe() entry; permutation.py exposes _require_triton_permutation for the same purpose. - Replace the dict returned by _compute_static_shape_info with a frozen _StaticShapeInfo dataclass. Callers use attribute access. tests/jax/test_moe_vjp.py, tests/jax/test_multiprocess_moe_vjp.py: - Drop the _inject_moe fixture and # noqa: F821 suppressions now that the imports are safe to do at module top. Signed-off-by: tdophung <tdophung@nvidia.com>
jberchtold-nvidia
removed
the
community-contribution
Collaborator
|
fyi @tdophung I've removed the |
Replace the dispatch state dict carried _dispatch -> _combine / _combine_bwd / _dispatch_bwd with one of two NamedTuples: * _PureJaxDispatchState: group_sizes, sorted_indices, routing_weights + optional EP residuals. * _TritonDispatchState: group_sizes, row_id_map, pad_offsets (Optional when align_size == 0), merging_probs + optional EP residuals. _DispatchState is the Union. _build_dispatch_specs / _build_ctx_specs now build the matching NamedTuple of P()s (with None at fields that the value won't populate, e.g. EP-only fields on non-EP runs and pad_offsets when align_size == 0), so the shard_map spec mirrors the value tree leaf-for-leaf. Signed-off-by: tdophung <tdophung@nvidia.com>
Replace the residual dict carried _body_fwd -> _body_bwd with a typed
_BodyCtx NamedTuple. Optional fields (expert_bias, aux_*) live on the
tuple as Optional[...] = None when the matching feature is disabled,
and _build_ctx_specs returns a matching _BodyCtx of P() so the
shard_map value and spec trees line up leaf-for-leaf.
The static side-info ('has_wi_bias' etc.) that used to live under the
sentinel '__static__' key is now passed alongside ctx as a separate
plain dict (returned as the residuals tuple (ctx, static) from the
fwd rule, unpacked at the top of the bwd rule). This keeps Python
ints/tuples out of the pytree leaves.
Signed-off-by: tdophung <tdophung@nvidia.com>
Swap the NamedTuple form of _PureJaxDispatchState, _TritonDispatchState and _BodyCtx for @flax.struct.dataclass. The pytree behavior is the same (each annotated field is a leaf, None is a non-leaf sentinel, shard_map specs and values stay tree-aligned), but the call sites read as dataclasses and we can reuse flax.struct.dataclass's pytree registration instead of relying on the NamedTuple shape. Signed-off-by: tdophung <tdophung@nvidia.com>
Under EP, _combine reassigns expert_outputs locally to the post- ragged_all_to_all tensor before Step 3 (the global combine). Saving the input (pre-A2A) tensor in ctx.expert_outputs meant _combine_bwd's Step-3 inverse (unpermute_bwd_with_merging_probs) consumed a tensor with the wrong shape and contents, silently corrupting d_expert_outputs. _combine now returns (output, expert_outputs_post_ep). The caller stashes the second value as the bwd residual so the Step-3 inverse sees the same tensor the forward Step 3 saw. Signed-off-by: tdophung <tdophung@nvidia.com>
Replace the two grouped_gemm calls (gate_proj and up_proj) with a
single fused grouped_gemm on `wi := concat([wi_0, wi_1], axis=-1)`
producing [T, 2M], which is then sliced into gate_proj_out and
up_proj_out. `sorted_x` is now quantized once per FFN (rather than
twice -- once per kernel) and the FFN bwd dgrad/wgrad also fuse: the
upstream grads concat along M, one grouped_quantize + one dgrad GEMM
+ one wgrad GEMM, then split.
ctx residuals collapse from two casted_wi_{0,1}_rhs_trans to one
casted_wi_rhs_trans [E, H, 2M]. The public API (separate wi_0, wi_1
params + biases) is unchanged so existing checkpoints and Flax
callers don't have to migrate.
FP8 caveat: per-expert amax for the fused wi is computed over
[H, 2M] instead of [H, M] per kernel, which can slightly shift the
representable range for one of the halves. Block-scaled recipes
(MXFP8/NVFP4) are largely unaffected since scales are per-block.
Signed-off-by: tdophung <tdophung@nvidia.com>
Comment on lines
+2122
to
+2123
| if permutation_backend is PermutationBackend.TRITON: | ||
| _require_triton() |
Contributor
There was a problem hiding this comment.
_dispatch (and its bwd counterparts _combine, _dispatch_bwd, _combine_bwd) unconditionally call make_chunk_sort_map and sort_chunks_by_map — both Triton operations — in the EP path (lines 549–553), regardless of permutation_backend. But _require_triton() is only triggered here when permutation_backend is TRITON. A user who passes PermutationBackend.PURE_JAX with ep_axis set on a Triton-less install will see a confusing TypeError: 'NoneType' object is not callable inside JAX's JIT compilation rather than a clear error. The check should also fire when ep_axis is not None.
jberchtold-nvidia
previously approved these changes
May 27, 2026
phu0ngng
previously approved these changes
May 27, 2026
| # via tests/jax/run_multiprocess_moe_vjp.sh (mirrors the pattern in | ||
| # examples/jax/encoder/run_test_multiprocessing_encoder.sh). Requires | ||
| # >=4 visible GPUs. | ||
| TE_PATH=$TE_PATH bash $TE_PATH/tests/jax/run_multiprocess_moe_vjp.sh \ |
Collaborator
There was a problem hiding this comment.
Just to learn, why do we want/need to run this test with multiple processes?
Collaborator
Author
There was a problem hiding this comment.
I encounter a bug similar to one described here: https://docs.google.com/document/d/18XCO7WNbTvadIG1WXPpO409dmu7HQaRlHG02mm5TYDA/edit?tab=t.0#heading=h.k4ivwttoxnil, but for triton kernels, where it would hang. More details were discussed in the TE/JAX meeting
Running the test with multiple process (1 process per GPU) circumvents that
Collaborator
Author
|
/te-ci |
3 similar comments
Collaborator
Author
|
/te-ci |
Collaborator
|
/te-ci |
Collaborator
Author
|
/te-ci |
Two failures from the same PR run: * CI B200 (8 GPUs) -- the test hard-coded EP_SIZE=2, FSDP_SIZE=2 -> 4-device mesh, while jax.device_count()==8. mesh_utils.create_device_mesh rejected with 'Number of devices 8 must equal the product of mesh_shape (2, 2)'. FSDP_SIZE now derives from jax.device_count() // EP_SIZE, so the test scales to whatever the launcher gives us (4 GPUs -> FSDP=2, 8 GPUs -> FSDP=4). BATCH formula already depends on EP_SIZE*FSDP_SIZE so it scales automatically and stays MXFP8-aligned. * CI H100 (4 GPUs) -- the file emits pytest.skip(allow_module_level=True) when get_device_compute_capability(0) < 100 (Blackwell-only). pytest reports 'no tests collected' with exit code 5, which the multiprocess launcher was treating as a failure. Accept 5 alongside 0; anything else is still a failure. Signed-off-by: tdophung <tdophung@nvidia.com>
* Drop now-unused dataclasses.field / typing.Callable / jax imports.
* Move the .flax.module._convert_to_activation_function import to the
top of moe.py (no longer circular after the big VJP refactor).
* Remove the dead pre_a2a_buffer_shape assignment in _dispatch.
* Rename q_set_w1 to _q_set_w1 in _body_bwd (the fused FFN bwd uses
q_set_w0 only; the placeholder keeps the unpack readable).
* Convert the two body_kwargs = dict(...) calls to {...} literals.
* Mark _combine_bwd / _body_fwd / _body_bwd / _build_dispatch_specs /
_build_ctx_specs / _moe_fwd_rule with
'pylint: disable=unused-argument' on the def line. These are
custom_vjp / shard_map signature requirements: rules MUST accept
every nondiff arg even when one branch doesn't read them, and
shard_map bodies must accept every spec key for in_specs /
out_specs structural compatibility.
* flax/moe.py: drop unused 'import jax'; keep the re-exported
PartitionSpec with explicit pylint + flake8 suppressions and a
comment explaining why.
Signed-off-by: tdophung <tdophung@nvidia.com>
Collaborator
Author
|
/te-ci |
jberchtold-nvidia
approved these changes
May 29, 2026
Baibaifan
pushed a commit
to Baibaifan/TransformerEngine
that referenced
this pull request
Jun 1, 2026
…upedGEMM and communication (NVIDIA#2912) Refactor MoEBlock into a unified MoE custom_vjp, add tests Replace the per-primitive custom_vjp boundaries in MoEBlock with a single jax.custom_vjp covering routing, dispatch, expert FFN, and combine. Helper functions group permute -> ragged_all_to_all -> local-permute into a single dispatch / combine pair, with a hand- derived bwd that mirrors the forward and runs entirely inside the EP shard_map body. Add a multi-process (one-GPU-per-process) test suite for the new unified VJP under a 2x2 (ep, fsdp) mesh: * tests/jax/test_multiprocess_moe_vjp.py -- fwd/bwd + aux_loss + PURE_JAX vs TRITON parity at Mixtral-ish shapes (batch=16, seq=2048, hidden=1024, intermediate=4096, num_experts=8, topk=2). * tests/jax/run_multiprocess_moe_vjp.sh -- launcher; forks one pytest process per visible GPU (mirrors examples/jax/encoder/run_test_multiprocessing_encoder.sh). * tests/jax/conftest.py -- pytest --num-process / --process-id options for the launcher. * qa/L0_jax_distributed_unittest/test.sh -- CI hook for the multiprocess smoke. Signed-off-by: tdophung <tdophung@nvidia.com> * [JAX] Fix EP+TRITON combine bwd: save post-A2A expert_outputs Under EP, _combine reassigns expert_outputs locally to the post- ragged_all_to_all tensor before Step 3 (the global combine). Saving the input (pre-A2A) tensor in ctx.expert_outputs meant _combine_bwd's Step-3 inverse (unpermute_bwd_with_merging_probs) consumed a tensor with the wrong shape and contents, silently corrupting d_expert_outputs. _combine now returns (output, expert_outputs_post_ep). The caller stashes the second value as the bwd residual so the Step-3 inverse sees the same tensor the forward Step 3 saw. Signed-off-by: Teddy Do <tdophung@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Signed-off-by: yangfan.bai <yangfan.bai@shopee.com>
Baibaifan
pushed a commit
to Baibaifan/TransformerEngine
that referenced
this pull request
Jun 1, 2026
…upedGEMM and communication (NVIDIA#2912) Refactor MoEBlock into a unified MoE custom_vjp, add tests Replace the per-primitive custom_vjp boundaries in MoEBlock with a single jax.custom_vjp covering routing, dispatch, expert FFN, and combine. Helper functions group permute -> ragged_all_to_all -> local-permute into a single dispatch / combine pair, with a hand- derived bwd that mirrors the forward and runs entirely inside the EP shard_map body. Add a multi-process (one-GPU-per-process) test suite for the new unified VJP under a 2x2 (ep, fsdp) mesh: * tests/jax/test_multiprocess_moe_vjp.py -- fwd/bwd + aux_loss + PURE_JAX vs TRITON parity at Mixtral-ish shapes (batch=16, seq=2048, hidden=1024, intermediate=4096, num_experts=8, topk=2). * tests/jax/run_multiprocess_moe_vjp.sh -- launcher; forks one pytest process per visible GPU (mirrors examples/jax/encoder/run_test_multiprocessing_encoder.sh). * tests/jax/conftest.py -- pytest --num-process / --process-id options for the launcher. * qa/L0_jax_distributed_unittest/test.sh -- CI hook for the multiprocess smoke. Signed-off-by: tdophung <tdophung@nvidia.com> * [JAX] Fix EP+TRITON combine bwd: save post-A2A expert_outputs Under EP, _combine reassigns expert_outputs locally to the post- ragged_all_to_all tensor before Step 3 (the global combine). Saving the input (pre-A2A) tensor in ctx.expert_outputs meant _combine_bwd's Step-3 inverse (unpermute_bwd_with_merging_probs) consumed a tensor with the wrong shape and contents, silently corrupting d_expert_outputs. _combine now returns (output, expert_outputs_post_ep). The caller stashes the second value as the bwd residual so the Step-3 inverse sees the same tensor the forward Step 3 saw. Signed-off-by: Teddy Do <tdophung@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Signed-off-by: yangfan.bai <yangfan.bai@shopee.com>
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
4 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.
Description
Most of MoE building blocks integration work has been deeply coupled with Maxtext development. Now creating this MoE block to isolate the work from Maxtext and provide more room for experimentation.
MoEBlockis a self-contained Flax-Linen module that wires together TE's fused router, pluggable token-dispatch backends (pure-JAX argsort or Tritonsort_chunks_by_index),grouped_dense-based expert FFN, and ragged-all-to-all (A2Av) expert parallelism viashard_mapThis first iteration will start with ring-of-experts EP, sharding on batch dimention for FSDP, CUBLASLt groupedGEMM and 2 permutation backend: pure JAX or Triton kernels. The block also exposes pluggable knobs for: weight layout (
wi_kernel_axes/wo_kernel_axes), permutation backend, A2A vs no-EP (single GPU) path, data-parallelism axes for true FSDP (batch sharded across(ep, fsdp)simultaneously), top-K with optional grouped/sigmoid scoring (for DSv3 workload), and optional auxiliary load-balancing loss.Fixes #2895
Type of change
Changes
transformer_engine/jax/flax/moe.py--MoEBlockLinen module:gate -> fused topk -> global permute -> A2A EP shard_map (ragged_a2a fwd, local permute, 3x grouped GEMM SwiGLU FFN, local unpermute, ragged_a2a rev) -> global combine.
transformer_engine/jax/permutation.pywith A2A param helpers (compute_ragged_all_to_all_params,compute_reverse_ragged_all_to_all_params,local_permute_after_a2a,local_unpermute_before_a2a) and the pure-JAXunfused_token_dispatch/unfused_token_combinepathswith custom VJPs.
tests/jax/test_moe_block.py-- single-device shape, backward,cross-backend equivalence, aux-loss, group-topk, JIT determinism.
tests/jax/test_distributed_moe_block.py-- EP=2 x FSDP=2 mesh test using the canonical Flax-Linen sharded-init pattern (eval_shape -> get_partition_spec -> logical_to_mesh_sharding -> jit(init, out_shardings=...)) anddata_parallelism_axes=("fsdp",)to exercise true FSDP (batch sharded across both axes).Checklist: