  Sparsely-gated mixture-of-experts (MoE) has been widely adopted to scale deep
learning models to trillion-plus parameters with fixed computational cost. The
algorithmic performance of MoE relies on its token routing mechanism that
forwards each input token to the right sub-models or experts. While token
routing dynamically determines the amount of expert workload at runtime,
existing systems suffer inefficient computation due to their static execution,
namely static parallelism and pipelining, which does not adapt to the dynamic
workload. We present Flex, a highly scalable stack design and implementation
for MoE with dynamically adaptive parallelism and pipelining. Flex designs an
identical layout for distributing MoE model parameters and input data, which
can be leveraged by all possible parallelism or pipelining methods without any
mathematical inequivalence or tensor migration overhead. This enables adaptive
parallelism/pipelining optimization at zero cost during runtime. Based on this
key design, Flex also implements various MoE acceleration techniques.
Aggregating all techniques, Flex finally delivers huge speedup at any scale --
4.96x and 5.75x speedup of a single MoE layer over 16 and 2,048 A100 GPUs,
respectively, over the previous state-of-the-art. Our evaluation shows that
Flex efficiently and effectively runs a real-world MoE-based model named
SwinV2-MoE, built upon Swin Transformer V2, a state-of-the-art computer vision
architecture. On efficiency, Flex accelerates SwinV2-MoE, achieving up to 1.55x
and 2.11x speedup in training and inference over Fairseq, respectively. On
effectiveness, the SwinV2-MoE model achieves superior accuracy in both
pre-training and down-stream computer vision tasks such as COCO object
detection than the counterpart dense model, indicating the readiness of Flex
for end-to-end real-world model training and inference.
