[WIP] Qwen3 MoE support

[intervitens]

This PR adds support for Qwen3 MoE (30B-A3B and 235B-A22B) models. Loss looked reasonable from a simple test with 30B-A3B on the Alpaca dataset.

TODO:

• Tensor/Expert parallel
• Test 235B model
• Verify loss curves against HF implementation
• LoRA support
• Documentation
• Tests

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchtune/2820

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[intervitens]

This is kinda hacky and slow, but stacking experts into a single tensor lets us use the existing MoE implementation and _grouped_mm for a very significant speed boost.

[SalmanMohammadi]

I found this to be prohibitively slow for larger models when testing with DeepSeek. I ended up following the torchtitan approach which uses a nn.ModuleDict for storing expert weights and grouping activated experts on-the-fly for grouped gemm.

[intervitens]

As far as I can see, torchtitan also needs to stack expert weights to use grouped gemm, they just do it at a slightly later point. They also don't support grouped gemm for training for now.

[SalmanMohammadi]

Oh you're totally right - I was thinking of EP.
If this doesn't take ages to convert the state dict it may be fine, otherwise only stacking the activated experts at runtime does make more sense to me.

[SalmanMohammadi]

YMMV but I ended up having to do something like

            denominator = top_scores.sum(dim=-1, keepdim=True) + 1e-20
            top_scores /= denominator

[dz1iang]

@intervitens Hi，Does it support training with fp8-compatible checkpoints?

[dz1iang]

Hello, I see the training process prompts: "Saving Qwen3 MoE adapter weights to PEFT format is not supported, saving to torchtune format instead."
May I ask how to obtain the Hugging Face (hf) checkpoint? Is there any code example for reference?
Additionally, I only set lora_attn_modules: ['q_proj', 'v_proj', 'output_proj'], apply_lora_to_mlp: False, and apply_lora_to_output: False—can the tune_to_peft_adapter_weights logic be used in this case? Thank you.