Integration of fused moe kernel (e.g., megablocks) for efficient moe training

[zinccat]

⚠️ Please check that this feature request hasn't been suggested before.

• I searched previous Ideas in Discussions didn't find any similar feature requests.
• I searched previous Issues didn't find any similar feature requests.

🔖 Feature description

The default huggingface moe implementation is inefficient due to for loop on experts, which causes low utilization when training models like Qwen3 30B-A3B, we might want a drop in kernel patch for more efficient moe calculation
A previous issue on this #930

✔️ Solution

Integrating solutions like https://github.com/databricks/megablocks/

❓ Alternatives

No response

📝 Additional Context

No response

Acknowledgements

• My issue title is concise, descriptive, and in title casing.
• I have searched the existing issues to make sure this feature has not been requested yet.
• I have provided enough information for the maintainers to understand and evaluate this request.

[NanoCode012]

Yep, this is something I was just checking. I saw that upstream transformers EP PR was merged huggingface/transformers#39501 . It uses kernels-community/megablocks (not sure if same as databrick's one)

I just need to read more on how it can be used and apply it for qwen3 too.

[zinccat]

Seems that original kernel is binded to gpt oss, I have made it work for qwen3, but it seems deepspeed is causing trouble if I try to merge the experts weight after loading, otherwise it requires merge the weight beforehand

[NanoCode012]

@zinccat , could you share how you got it working for qwen3 for reference purposes?

This is currently a WIP for us.

[zinccat]

sure, I'll provide my version later

[zinccat]

you may reference the commits in https://github.com/zinccat/qwen3_moe_megablocks

[NanoCode012]

@zinccat , correct me if I'm wrong but is the shape for the router mixed up?

        self.weight = nn.Parameter(torch.empty(config.num_experts, config.hidden_size, dtype=torch.bfloat16))

Should it be:

self.weight = nn.Parameter(torch.empty(
            config.hidden_size, config.num_experts, dtype=torch.bfloat16
        ))

? Ref: https://github.com/huggingface/transformers/blob/9b4bd96e2b3e0c65bac21706dfeb2fc5ff7e3c22/src/transformers/models/llama4/modeling_llama4.py#L133

[zinccat]

it's either self.weight = nn.Parameter(torch.empty(config.num_experts, config.hidden_size, dtype=torch.bfloat16)) or the class as a subclass ofnn.Linear(config.hidden_size, config.num_experts), as specified in your reference

[lhl]

Just an FYI, I've been poking around w/ the kernel-commmunity/megablocks code myself, here's a HIP-compatible version I've been bringing up (all numerics are valid, however I get RCCL hangs on my MI300X...) https://huggingface.co/shisa-ai/megablocks-hip

[NanoCode012]

@zinccat , late getting back to this. I've just been checking transformers v5 implementation on this with your megablocks. Have you done numerical checks? I'm finding that there's a large difference in forward and backward pass against baseline in my version (set bias to 0, permuted params to follow gptoss, changed to silu following yours).

  megablocks forward MISMATCH: max_diff=0.511719
  megablocks backward MISMATCH: max_diff=0.832031

I also pinged you on Discord (hopefully it's you).