[JazyJiang]

Related issue

No tracking issue yet. This is the first step of the EP / MoE layer refactor. Background: #6464, #6465 and #6466 (fixed in #6467, #6469, #6468) came from the EP logic being copied into every model's forward. A follow-up that unifies the gradient sync groups of expert parameters builds on the ExpertGradScale tag added here.

Background

Today every EP model (Mixtral, DeepSeek, DeepSeek-V3, and Qwen3-MoE in #6450) carries its own copy of routing, permutation, all-to-all, the expert loop and the combine. That has two costs:

• Performance: a Python loop over (source rank × expert) chunks with several host syncs per layer. In a single-layer benchmark on H20 (EP=4, forward + backward), this is 3–7× slower than Megatron-Core with grouped GEMM on Qwen3-style shapes (e.g. Qwen3.5-like, 256 experts top-8: 87.1 vs 13.3 ms).
• Maintainability: fixing or replacing one step (e.g. grouped GEMM, DeepEP) means editing every model, and gradient-scaling bugs were copied the same way.

This PR adds the shared layer and the contract between its parts, so that new backends and model ports can be built on one definition. Nothing uses it yet.

Changes

• colossalai/moe/dispatcher.py

  • DispatchOutput, the data contract handed to the experts:
    • rows grouped by local expert; within an expert, rows keep the global token order of the EP group (source rank, then position), so expert grads can be compared bitwise with a single-GPU reference;
    • per-expert row counts on the device, plus the host copy when the dispatcher has one;
    • optional fp32 routing weights for dispatchers that apply them in the experts.
  • The Dispatcher protocol (dispatch / combine).
  • NCCLAllToAllDispatcher: token-major sort, one host sync per layer (the split sizes NCCL needs), all_to_all_uneven, regroup by expert on the GPU. combine applies the fp32 routing weights on the source rank, sums in fp32 and casts once.

• colossalai/moe/experts.py

  • The Experts protocol: SwiGLU experts on the dispatched rows; outputs are not weighted.
  • GroupedMMExperts: two torch._grouped_mm calls for all local experts. Needs torch ≥ 2.8; before 2.11 each expert's rows are padded to a multiple of 8.
  • ReferenceExperts: one F.linear pair per expert, for debugging and tests.

• colossalai/moe/layer.py — EPMoELayer:

  • forward(hidden_states, topk_idx, topk_weights). The router stays in the model.
  • Local experts are stored in the transformers v5 layout: gate_up_proj [E_local, 2I, H], down_proj [E_local, H, I].
  • load_full_weights copies this rank's experts by writing .data only, so parameter objects held by optimizers / ZeRO stay valid.
  • With the default ep_group=None and distributed initialized, the layer resolves the group to dist.group.WORLD before building the default dispatcher and tagging the parameters, so get_ep_size(param) / get_ep_rank(param) work on its parameters ([BUG] Default EPMoELayer group is rejected by get_ep_size #6485, fixed in [moe] resolve the default EP group before tagging parameters #6486, which is merged into this branch).

• colossalai/tensor/moe_tensor/api.py — ExpertGradScale with set_expert_grad_scale / get_expert_grad_scale. This records where the 1/ep factor of an expert grad is applied:

  • IN_FORWARD: in the layer's autograd, like the existing EPGradScaler;
  • IN_REDUCTION: by the gradient reduction.

  Untagged expert params count as IN_FORWARD. EPMoELayer defaults to IN_FORWARD, so it is correct with the current plugins. IN_REDUCTION is for plugins that read the tag; that support is being designed separately.

• colossalai/moe/__init__.py: the new names are exported lazily, like the existing _operation ones; importing colossalai.moe still loads no submodule.

• New tests/test_moe/test_ep_moe_layer.py (4 GPUs):

  • Dispatch contract, ep = 1/2/4: rows reach the right local expert, grouped by expert and in global token order; the GPU and host counts agree; an identity-expert dispatch + combine returns the input. Also covers a rank with no tokens and an expert that is never chosen.
  • Layer vs fp32 single-GPU reference: output, input grad and per-expert weight grads, over ep = 1/2/4 × {grouped_mm, reference} × {IN_FORWARD, IN_REDUCTION} × {uniform, half the experts unused}. Unused experts must get zero (not missing) grads.
  • Parameter contract, with an explicit EP subgroup and with the default group: the layer, its dispatcher and both parameters carry the same group, get_ep_size(param) matches, the ExpertGradScale tag is set, and load_full_weights keeps the parameter objects.

• New tests/test_moe/test_ep_moe_layer_cpu.py (CPU, Gloo): default-group construction with distributed initialized and not initialized. Added to both CPU test lists (run_cpu_unit_tests.sh, run_colossalai_batch.sh).

Not in this PR (follow-ups on top of this contract):

• the Triton fused combine kernel (bitwise equal to the torch combine, ~15–20% faster for the layer);
• DeepEP and Transformer Engine backends;
• shared experts, overlap hooks;
• porting models / ShardFormer policies to EPMoELayer (after support Transformers 5.17 in ShardFormer #6474's v4/v5 split settles);
• plugin support for IN_REDUCTION and unified expert sync groups.

User and compatibility impact

• Python API: new public classes in colossalai.moe (EPMoELayer, DispatchOutput, Dispatcher, NCCLAllToAllDispatcher, Experts, GroupedMMExperts, ReferenceExperts) and new helpers in colossalai.tensor.moe_tensor.api. Nothing is removed or renamed.
• Default behavior, model accuracy, performance, existing training scripts: none. No model, policy or plugin uses the new code.
• Checkpoint format, Python/PyTorch/CUDA versions, C++/CUDA kernels, Docker images: none. GroupedMMExperts raises a clear error below torch 2.8; ReferenceExperts works on any version.

Test environment

• OS: Ubuntu 22.04.5 LTS
• Python: 3.10.12
• GPU: 8 × NVIDIA H20, 1 node (4 used)
• NVIDIA driver: 595.71.05 (nvidia-smi CUDA Version: 13.2)
• PyTorch: 2.11.0+cu128 (torch.version.cuda: 12.8, NCCL 2.28.9)
• CUDA toolkit (nvcc --version): 12.8 (V12.8.61)
• Transformers: 4.51.3
• flash-attn: 2.7.4.post1
• ColossalAI: the first commit on top of main 0ac6a922 (rebased after the Transformers 5.17 upgrade, [ci]: upgrade Transformers to 5.17 and align CI dependencies #6476). The branch has since been rebased onto main d414d565; the CI results below are for that head.

The local environment still has transformers 4.51.3, which main supports through the v4/v5 split of #6474; the new code does not use transformers. CI uses torch 2.13 / CUDA 13.0 with transformers 5.17; CI results will be added below.

Test commands and results

FAST_TEST unset, 4 GPUs reserved via gpu-run.

python -m pytest -p no:cacheprovider -rA tests/test_moe/test_ep_moe_layer.py
python -m pytest -p no:cacheprovider -rA tests/test_moe

Run	Result
test_ep_moe_layer.py	1 passed, 0 failed (before and after the rebase onto 0ac6a922)
tests/test_moe	5 passed, 7 skipped, 0 failed

• Tolerance. The largest relative L2 error over all checks of the new test is 6.0e-3; the tolerance is 2e-2.
• Sensitivity check. With EPGradScalerOut removed from the layer, the test fails (input grad error 0.93 at ep=2).
• Skips in tests/test_moe are pre-existing:
  • 4 × "tested in corresponding sharderformer";
  • 2 × "requires MOE_TENSOR_PATH reference tensors";
  • test_moe_ep_zero_grad_scale[8] needs 8 GPUs.
• Lazy exports. import colossalai.moe loads no submodule, the old names still resolve, and the new names load on first access.
• torch._grouped_mm edge cases (checked on torch 2.11, H20): bf16 / fp16 / fp32, an expert with no rows, zero rows in total, and row counts that are not a multiple of 8 all work, and empty experts get zero weight grads.

Logs and JUnit XML are kept on the team cluster in test-results/ep_moe_layer_20261022/.

CI (torch 2.13.0+cu130, transformers 5.17, head 9a1c8b8b on main d414d565):

Layer	Result
Layer 1 (formatting, lint, CPU tests, including test_ep_moe_layer_cpu.py)	pass
Layer 2 (batches 2, 5)	pass
Layer 3, batch 8, fast_mode=false (run 40622247693)	157 passed, 3 failed, 28 skipped; the 3 failures are unrelated to this PR, see below

Earlier head 05585301 (before the default-group fix): Layer 1, Draft smoke (batch 4a) and Layer 2 passed; Layer 3 batch 8 (run 39793007918) had the same 3 unrelated failures.

• test_ep_moe_layer[4] and both test_ep_moe_layer_cpu cases passed. The other MoE tests (test_moe_checkpoint, test_moe_ep_sp_grad_sync, test_moe_ep_zero_grad_norm, test_moe_ep_zero_grad_scale[4]) and the MoE optimizer-checkpoint tests added on main by [checkpoint]: remove embedded pickle from async optimizer checkpoints #6479 also passed.
• The 3 failures come from main after the Transformers 5.17 upgrade (support Transformers 5.17 in ShardFormer #6474, [ci]: upgrade Transformers to 5.17 and align CI dependencies #6476) and touch none of the files changed here:
  • test_booster/test_plugin/test_3d_plugin.py::test_3d_plugin: 'GPT2Model' object has no attribute 'get_head_mask'. shardformer/modeling/gpt2.py:186 still calls an API removed in Transformers v5.
  • test_checkpoint_io/test_gemini_checkpoint_io.py::test_gemini_ckpIO and test_gemini_torch_compability.py::test_gemini_ckpIO[2]: GeminiDDP.named_parameters() got an unexpected keyword argument 'remove_duplicate'. general_checkpoint_io.py:308 (added in support Transformers 5.17 in ShardFormer #6474) passes remove_duplicate, which GeminiDDP.named_parameters(prefix, recurse) does not accept.
• An earlier Layer 3 attempt (run 39761270429) hung for 3.5 h on a Gemini test in tests/test_booster, after one rank hit an NCCL out-of-memory error on the shared host, and was cancelled. This PR's tests had not started yet.

Known issues

• The first Draft smoke run failed before any test ran: the CI image had already been switched to the Transformers 5.17 recipe while this branch still had the old Dockerfile ("CI image does not match Dockerfile"). Fixed by rebasing onto main.
• The default ep_group=None was stored as None on the expert parameters, so get_ep_size(param) asserted ([BUG] Default EPMoELayer group is rejected by get_ep_size #6485). Fixed in this PR (see Changes) and covered by both the GPU and the CPU test.
• The first local run failed on a rank with zero tokens: combine reshaped with an inferred -1 dimension. This is fixed in this PR and covered by the test, not a pre-existing issue.
• With expert_grad_scale=IN_REDUCTION, current plugins do not divide expert grads by ep_size, so they would be ep_size times too large. The default is IN_FORWARD, and the docstring says to use IN_REDUCTION only with a plugin that reads the tag.
• Bitwise equality of the NCCL forward with transformers v5's grouped_mm MoE path was checked in standalone experiments with transformers 5.17. The test compares against its own fp32 reference so that it does not depend on the transformers version; a comparison with the HF v5 MoE block can follow now that CI runs 5.17.
• No new skip, xfail or --ignore.

Risk and rollback

• Risk: low. New files plus additive changes to two existing ones and one line in each CPU test list. The colossalai.moe lazy-export table is rewritten as a name → submodule map; the old names map to _operation as before.
• Rollback: revert the two commits of this PR.

Documentation and release impact

API docstrings are included. No README, installation, Docker, version.txt, PyPI or DockerHub changes. A changelog line: "Add colossalai.moe.EPMoELayer, a model-independent expert-parallel MoE layer with pluggable dispatchers and experts backends (experimental)."