[REQUEST] ZeRO stage 3 support for mixture-of-experts (MoE) layer

Hello everyone,

I've always wanted to run large models using minimal GPUs, as I only have a few at my disposal. That is why I was impressed that ZeRO-3 can support the running of large models by offloading parameters from GPU memory to CPU memory.

However, I recently discovered that ZeRO-3 does not support the use of MoE models, which came as a shock to me. Personally, I believe that the MoE model is a well-known and effective way to increase model capacity. Therefore, I think it would make sense for ZeRO to support the MoE model.

I'm wondering if it's true that ZeRO-3 does not support MoE model inference?

  File "/root/deepspeed/ds/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 1291, in _configure_optimizer
    self.optimizer = self._configure_zero_optimizer(basic_optimizer)
  File "/root/deepspeed/ds/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 1575, in _configure_zero_optimizer
    assert not self.has_moe_layers, "MoE not supported with Stage 3"
AssertionError: MoE not supported with Stage