[shardformer]: support Qwen3-MoE with expert parallelism

[LiRunGuo]

📌 Checklist before creating the PR

• I have created an issue for this PR for traceability
• The title follows the standard format: [doc/gemini/tensor/...]: A concise description
• I have added relevant tags if possible for us to better distinguish different PRs
• I have installed pre-commit: pip install pre-commit && pre-commit install (ran the pinned hooks manually instead: black 24.10.0 --line-length=120, isort 5.13.2 --profile black, autoflake 2.3.1, all clean)

🚨 Issue number

Closes the Qwen part of #6180 ("Expert Parallel for qwen/deepseek"; DeepSeek EP already exists).

Depends on #6449: this branch contains that fix as its first commit, because the new test saves a Qwen3-MoE checkpoint with a padded vocab under TP. I will rebase once it is merged; please review only the second commit here.

📝 What does this PR do?

Adds shardformer support for Qwen3-MoE (Qwen3MoeModel, Qwen3MoeForCausalLM, transformers >= 4.51, e.g. Qwen3-30B-A3B / Qwen3-235B-A22B) with MoeHybridParallelPlugin.

feature	status
expert parallelism	✅ EPQwen3MoeSparseMoeBlock: all-to-all token dispatch, experts sharded over the EP group, optional TP inside each expert, Qwen3's norm_topk_prob
dense layers (mlp_only_layers, decoder_sparse_step)	✅ kept as they are (TP-sharded if TP is on)
tensor parallelism	✅ attention, router, dense MLP, vocab parallel embedding, lm_head
pipeline parallelism	✅ 1F1B and interleaved; router logits are carried across stages for the load balancing loss, dense layers' None router logits are skipped
sequence parallelism	✅ all_to_all, reusing the Qwen3 attention forward (Qwen3MoeAttention is identical to Qwen3Attention)
ZeRO 1/2	✅ through the plugin
other SP modes, SP + PP, zero bubble schedule	❌ raise NotImplementedError

The EP block is adapted from EPMixtralSparseMoeBlock; the policy mirrors the Mixtral / Qwen3 policies.

Tests

• tests/test_shardformer/test_model/test_shard_qwen3_moe.py (4 GPUs, plus an 8-GPU largedist variant like Mixtral's):
  • training against a single-GPU reference for 11 (ZeRO, EP, PP, TP, SP) combinations, on a model with GQA and both dense and sparse layers; the sharded checkpoint is reloaded with Qwen3MoeModel.from_pretrained and compared.
  • Qwen3MoeForCausalLM loss including the router aux loss for EP=4, PP=4, TP=4 and EP=2×PP=2 (the reference is computed on the same micro batches, since the aux loss is not linear in the batch).
• tests/test_moe/test_qwen3_moe_layer.py (4 GPUs): fp32 comparison of the EP block with the transformers block — output, router logits, input grad and every parameter grad — for 1 and 2 experts per rank and both norm_topk_prob values.

A note on test sensitivity

The model-level test (same design as test_shard_mixtral.py) cannot detect wrong expert math: at the default initializer_range the expert outputs are ~1e-4 against a residual stream of ~1, so e.g. flipping norm_topk_prob in the EP block still passes it (max output error 0.0156 with or without the mutation, i.e. bf16 noise). The layer test uses a larger init and fails on that mutation on every rank, so I added it and did not skip it. test_mixtral_layer.py is currently skipped as "tested in corresponding shardformer" — the same blind spot applies to Mixtral.

Results (4x H200, torch 2.5.1+cu124, transformers 4.51.3)

test	result
test_shard_qwen3_moe.py::test_qwen3_moe (11 model configs + 4 causal LM configs)	passes
test_qwen3_moe_layer.py (ep 2/4 × norm_topk_prob on/off)	passes
layer test with norm_topk_prob flipped in the EP block	fails on all ranks (as intended)

Sharding check: with ep=4 each rank holds 1 of the 4 experts in each of the 3 sparse layers.

Also found (upstream transformers, not changed here)

In transformers 4.51.3, Qwen3MoeModel.forward collects the router logits of dense layers (None) and load_balancing_loss_func then fails when output_router_logits=True and mlp_only_layers is non-empty. The pipeline / SP forward in this PR skips them; the causal LM test uses only sparse layers so that the transformers reference can compute the aux loss.

💥 Checklist before requesting a review

• I have linked my PR to an issue (instruction)
• My issue clearly describes the problem/feature/proposal, with diagrams/charts/table/code if possible
• I have performed a self-review of my code
• I have added thorough tests.
• I have added docstrings for all the functions/methods I implemented

⭐️ Do you enjoy contributing to Colossal-AI?

• 🌝 Yes, I do.
• 🌚 No, I don't.