ZeRO integration
ZeRO technique enables data scientists to scale PyTorch models efficiently by providing memory savings.
DeepSpeed ZeRO and FairScale FSDP are some of the common independent implementations of the ZeRO
technique. ZeRO techniques can be applied on an MoE model that is using ORT MoE. ZeRO technique reduces
memory consumption which could enable data scientists to increase batch size. We were able to increase
batch size of a CLIP model by applying ZeRO using Fairscale FSDP implementation. Figure 11 shows the result.
We applied FSDP on 2 replicas of CLIP Vision model running on 8 GPUs. In this configuration expert
parameters are sharded in two partitions and non-expert parameters are sharded in eight partitions. When at
batch 36, with FSDP, it uses 36% less peak memory. Without FSDP, the max batch size per gpu can reach 36,
while with FSDP, it increases to 48.

Figure 11 : ZeRO helps improve batch size for the CLIP Model

Ongoing Effort
Scaling Speech, Language and Vision Models with Mixture of Experts Technique - Microsoft Community Hub