Shortcut-connected Expert Parallelism for Accelerating Mixture-of-Experts

TL;DR

The paper tackles the bottleneck of All-to-All communication in sparsely-gated MoE models by introducing ScMoE, a shortcut-connected architecture that decouples communication from computation and uses an overlapping parallelization strategy. By processing preceding-layer representations with a top-1 MoE via shortcuts while a shared expert handles current-layer representations, ScMoE achieves substantial training and inference speedups across vision and language models, while maintaining or exceeding baseline model quality. The work also analyzes gating behavior and representation similarity to justify the design and presents memory-limited inference techniques that further extend practical deployment. Overall, ScMoE offers a device-agnostic method to accelerate MoE models on both single and multi-node setups, with broad applicability and potential for further optimization.

Abstract

Expert parallelism has emerged as a key strategy for distributing the computational workload of sparsely-gated mixture-of-experts (MoE) models across multiple devices, enabling the processing of increasingly large-scale models. However, the All-to-All communication inherent to expert parallelism poses a significant bottleneck, limiting the efficiency of MoE models. Although existing optimization methods partially mitigate this issue, they remain constrained by the sequential dependency between communication and computation operations. To address this challenge, we propose ScMoE, a novel shortcut-connected MoE architecture integrated with an overlapping parallelization strategy. ScMoE decouples communication from its conventional sequential ordering, enabling up to 100% overlap with computation. Compared to the prevalent top-2 MoE baseline, ScMoE achieves speedups of 1.49 times in training and 1.82 times in inference. Moreover, our experiments and analyses indicate that ScMoE not only achieves comparable but in some instances surpasses the model quality of existing approaches.

Figures (15)

• Figure 1: The overhead of MLP and top-2/top-1 MoE in a transformer block of SwinV2-MoE-S tutel model, allocating one expert per GPU with expert parallelism. The All-to-All communication takes up 60% of total time on a single node with 8$\times$A30 GPUs, but drops to 15% on 8$\times$A800 due to the latter's 6$\times$ higher bandwidth provided by GPU-to-GPU NVLink NVLink. Despite benefiting from NVLink, communication still approaches 50% due to the lower-bandwidth inter-node Ethernet GPUInterconnect when scaling across multiple nodes.
• Figure 2: Illustrations of the standard top-2 MoE architecture (a) and the corresponding shared-expert MoE architecture (b). "SE" in (b) denotes the shared expert.
• Figure 3: Illustration of scaling MoE transformer layer across multiple devices with expert parallelism.
• Figure 4: Illustrations of various ScMoE architectures with shortcut connections to different positions of the preceding layer: (a) "Pos-1" output, (b) "Pos-2" intermediate, and (c) "Pos-3" input. The red line indicates the transmission of the preceding-layer representations to the MoE via a shortcut connection. Details regarding pre-layer normalization and dropout procedures have been excluded for simplicity.
• Figure 5: Illustration showcasing the application of the ScMoE (Pos-1) architecture to the MoE model, wherein the MoE module is integrated into each Transformer block.