MoE (mixture of experts) support for expert parallel

[MeteorsHub]

Description & Motivation

Hi, is there any tutorial on how to setup the expert parallel in lightning trainer? Maybe the shortest way is to use DeepSpeed MoE (https://www.deepspeed.ai/tutorials/mixture-of-experts/) since lightning support deep speed strategy. But it needs to set some optimizer parameter groups or something. Can anyone give me some hints to achieve this?

Pitch

No response

Alternatives

No response

Additional context

No response

cc @lantiga @Borda

[Captain-Named]

Hi, have you solve the problem that how to inference an moe model with lightning-deepspeed? When I try to conduct that(llama4-17b16e) with lightning+ds standard stage3, I found it stuck at the initialization of GPUs. I wonder if it is because of the unsupport of moe models of lightning.

[MeteorsHub]

I also encountered NCCL timeout with large MoE models in DeepSpeed stage 3 (small MoE models work). Now I always only run it on stage 2.