[REQUEST] Support for Expert Optimizer State Partitioning with ZeRO Optimization in DeepSpeed MoE

Is your feature request related to a problem? Please describe.
I am currently training a Transformer model using DeepSpeed's MoE (Mixture of Experts) layer. After enabling Expert Parallelism (EP) and ZeRO-2 optimization, the total activation size of the experts is 2B, and the total parameter count is 6B. Since the expert parameters account for the majority of the parameters (5B), while the shared parameters are minimal, I observed that the memory usage reduction after enabling ZeRO-2 is not significant. From my understanding, as mentioned in the DeepSpeed MoE documentation, in the EP + ZeRO-2 setup, only the gradients and optimizer states of non-expert parameters are partitioned, while the gradients and optimizer states of the MoE layer are not optimized. This results in a substantial memory overhead. As a comparison, I tested the training of a non-MoE 2B model. With almost the same activation size, the non-MoE model can use a batch size 6-7 times larger than the MoE model, and it also trains faster. Are there any plans to optimize the memory consumption of DeepSpeed MoE in the future? Alternatively, do you have any suggestions to help optimize the training process for my model?

Describe the solution you'd like
A clear and concise description of what you want to happen.
I would like DeepSpeed to extend ZeRO optimization to the MoE layer parameters, enabling more efficient memory usage and allowing the training of larger models with MoE architectures.

Describe alternatives you've considered
Enable zero optimization in Deepspeed moe layer parameters so that we can train larger model.

Additional context
Add any other context or screenshots about the feature request here.

The issue names no implementation files, tests, or entry points. Start with the DeepSpeed MoE documentation and the described EP + ZeRO-2 behavior, then identify where expert parameters, gradients, and optimizer states are handled. Done means expert-layer state partitioning works without breaking distributed training and its memory benefit is validated.