PaperlayerView exact revision on arXiv

Source-linked AI summary

Tutel: Adaptive Mixture-of-Experts at Scale

Changho Hwang, Wei Cui, Yifan Xiong, Ziyue Yang, Ze Liu, Han Hu, Zilong Wang, Rafael Salas, Jithin Jose, Prabhat Ram, Joe Chau, Peng Cheng, Fan Yang, Mao Yang, Yongqiang Xiong

arXiv:2206.03382v2cs.DCcs.CLcs.CV

TL;DR

MoE's token routing creates workload variation that static parallelism and pipelining do not accommodate efficiently. TUTEL provides adaptive parallelism and pipelining through a shared distribution layout that enables zero-cost runtime switching. It achieves up to 5.75× single-layer speedup at 2,048 GPUs and accelerates SwinV2-MoE in both training and inference while retaining superior accuracy over its dense counterpart.

  • Problem

    MoE expert workload changes with token routing, while existing static execution and runtime parallelism switching do not efficiently accommodate that variation.

  • Method

    TUTEL uses a shared distribution layout to support adaptive parallelism and pipelining without tensor migration or mathematical inequivalence.

  • Results

    TUTEL achieves up to 5.75× speedup for a single MoE layer and up to 1.55× training and 2.11× inference speedup for SwinV2-MoE over Fairseq.

  • Takeaways & Limitations

    TUTEL supports training and inference of real-world MoE models, and SwinV2-MoE shows effectiveness compared with its counterpart dense model.

  • Takeaways & Limitations

    The evaluation considers system-side solutions that apply regardless of load-balancing loss and uses the minimum required capacity factor without dropped tokens or unnecessary computation.

Abstract

from arXiv · show

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.

1 INTRODUCTION

MoE scales model capacity through sparse expert routing, but changing token workloads make static execution inefficient. TUTEL addresses this with zero-cost adaptive parallelism and additional acceleration techniques, yielding substantial MoE-layer and end-to-end speedups.

  • Motivation: MoE routes each input token to selected parallel experts, but expert workload changes with the input data and gating function.Experiments report workload changes of up to 4.38× within a single training run.
  • Motivation: Static execution can miss the best parallelism because the optimal strategy varies with dynamic workload, while runtime switching adds redistribution and memory overhead.Manipulating the algorithm with load-balancing loss can also harm model accuracy.
  • TUTEL Design: TUTEL uses one distribution layout for all possible parallelism strategies, enabling iteration-level adaptive switching without tensor reformatting or migration overhead.Communication-cost analysis is used to preserve the optimal parallelism strategy.
  • TUTEL Design: TUTEL combines adaptive parallelism with adaptive pipelining, hierarchical All-to-All communication, and sparse GPU encode/decode kernels.The framework is implemented for diverse MoE algorithms and integrated into Fairseq and DeepSpeed.
  • Results: TUTEL accelerates SwinV2-MoE by up to 1.55× in training and 2.11× in inference compared with Fairseq.The system is evaluated on a real-world sparse MoE vision model.
  • Results: 4.96× and 5.75× speedup is achieved for a single MoE layer over 16 and 2,048 A100 GPUs, respectively.The reported gains aggregate TUTEL's acceleration techniques.

2 BACKGROUND & MOTIVATION

MoE dynamically routes tokens to experts, causing workload and capacity to vary across iterations. Existing static parallelism and pipelining can therefore waste computation or miss better execution strategies.

  • Dynamic Workload of MoE: MoE routes each token to multiple experts, often unevenly, so each expert’s workload changes at every iteration.Expert capacity depends on batch tokens T, global experts E, top-k routing, and capacity factor f.
  • Dynamic Workload of MoE: A larger capacity factor f indicates more imbalanced routing and requires an expert to process more tokens.f = 1 represents the minimum value and the most even token distribution.
  • Dynamic Workload of MoE: Static upper-bound capacity fupper causes unnecessary computation and may drop training tokens when set too low.TUTEL instead supports the minimum required f that avoids both unneeded computation and dropped tokens.
  • Static Parallelism: Runtime parallelism switching is costly because existing methods use incompatible data and weight layouts and require parameter migration.These systems also differ in data splits, weight splits, gradient momentum management, and training interfaces.
  • Static Pipelining: Static pipelining cannot consistently select the best All-to-All algorithm and pipelining degree across MoE settings and scales.Communication and computation interference makes separate optimization of these choices unreliable.

3 ADAPTIVE MOE WITH TUTEL

TUTEL enables runtime adaptation by using compatible layouts and execution flows for switchable parallelism, while jointly adapting pipelining and communication. It also overlaps communication with expert computation through fine-grained token partitioning and caches optimal configurations.

  • Adaptive Parallelism: TUTEL uses one distribution layout covering possible parallelism strategies, avoiding input or weight reformatting during switching.The design preserves mathematically equivalent MoE execution and keeps switching overhead within O(1) for migration.
  • Adaptive Parallelism: Communication-complexity analysis reduces seven possible EP, DP, and MP combinations to a smaller efficiency-equivalent set of execution flows.Because GPU computation is identical, communication complexity determines relative parallelism efficiency.
  • Adaptive Parallelism: Switchable EP+DP+MP uses control parameter r to group GPUs, perform DP within groups, and MP across groups while sharing DP’s input and weight format.Increasing r reduces group size; at r = W/E, within-group all-gather is eliminated.
  • Adaptive Pipelining: Adaptive pipelining jointly selects pipelining degree and Linear or 2DH All-to-All because communication latency affects the optimal strategy.TUTEL partitions only the two All-to-All operations and the intervening expert, avoiding whole-layer partitioning problems.
  • Adaptive Pipelining: Fine-grained capacity partitions overlap All-to-All communication with expert computation on separate GPU streams.Customized operations perform partitioning, reshaping, and data shuffling inline without extra copying overhead.
  • Configuration Search: TUTEL stores optimal adaptive parallelism, pipelining degree, and All-to-All algorithm by quantized expert-capacity ranges.The search uses ternary search for r, four pipelining degrees, and two All-to-All algorithms.

4 IMPLEMENTATION

TUTEL supports dynamic MoE behavior through adaptable routing settings, flexible communication layouts, and optimized encode/decode kernels. These mechanisms target throughput, memory use, and runtime adaptability across configurations.

  • Framework support: TUTEL provides broader support for devices, data types, and MoE-related features than DeepSpeed MoE, Fairseq MoE, and FastMoE.The implementation includes adaptive parallelism switching and related MoE functionality.
  • Dynamic routing: TUTEL supports dynamic top-ANY routing, allowing the top-k sparsity setting to change across iterations.This enables users to fine-tune MoE sparsity dynamically.
  • Dynamic capacity: TUTEL dynamically adjusts the capacity factor to control capacity under changing token imbalance.A positive setting applies a fixed factor, while zero or negative settings enable bounded adaptation.
  • Communication: Flexible All-to-All produces an (Eg, C, D) layout, preserving the same expert matrix-multiplication shape at any GPU scale.This differs from conventional All-to-All, whose Cg depends on the number of GPUs W.
  • Kernel optimization: Fast encode and decode replace multiple einsum and matrix-multiplication operations with SIMT-efficient sparse GPU operations.The optimization reduces non-expert latency and saves 20% ∼90% of GPU memory in most cases.

5 EVALUATION

Evaluation across model configurations and scales shows that TUTEL’s adaptive execution and kernel, communication, and pipelining optimizations improve MoE-layer and end-to-end performance. SwinV2-MoE also outperforms its dense counterpart on the reported vision tasks.

  • Adaptive parallelism: The optimal parallelism method changes with expert configuration and capacity factor, motivating adaptive switching.As capacity decreases, preference can shift from DP to EP+DP and then EP+DP+MP.
  • Adaptive pipelining: Adaptive pipelining improves average performance by 9% ∼101% over the baseline and 1% ∼107% over static strategies across 243 models.Worst-case improvement reaches 23% ∼599%.
  • Adaptive pipelining: 39% improvement at f = 4 and 57% improvement at f = 8 are achieved over degree-1 pipelining under dynamic workloads.Adaptive pipelining selects the best strategy for the tested capacity-factor patterns.
  • Single-layer scaling: 4.96×, 3.11×, and 5.75× speedups are achieved over the Fairseq / DeepSpeed baseline on 16, 128, and 2,048 GPUs, respectively.The breakdown combines kernel, 2DH All-to-All, Flexible All-to-All, and adaptive pipelining optimizations.
  • End-to-end speed: TUTEL accelerates SwinV2-MoE training by 1.14× ∼1.55× and inference by 1.95× ∼2.11× versus Fairseq.The comparison covers GPU counts from 8 to 128.
  • End-to-end accuracy: SwinV2-MoE-B exceeds the dense counterpart on ImageNet pre-training, fine-tuning, few-shot classification, and COCO detection accuracy.Reported gains include +1.3% ImageNet-22K top-1 accuracy and +0.4/+0.4 COCO box/mask AP.

6 CONCLUSION

TUTEL addresses dynamic MoE execution with adaptive parallelism and pipelining, and evaluation shows gains at large scale. Its use in SwinV2-MoE demonstrates end-to-end training and inference on computer-vision workloads.

  • Conclusion: TUTEL addresses dynamic MoE execution through adaptive parallelism for expert execution and adaptive pipelining for dispatch and combine operations.The conclusion presents these as the system’s two major design aspects.
  • Conclusion: 5.75× speedup is achieved for a single MoE layer on an Azure A100 cluster with 2,048 GPUs.The reported result is part of the conclusion’s large-scale evaluation.
  • Conclusion: The SwinV2-MoE example shows TUTEL supporting training and inference for a real-world state-of-the-art deep-learning model.The comparison is made against the counterpart dense model for computer-vision effectiveness.

A.2 Approach and Challenges

The 2DH All-to-All approach aggregates local chunks before network transfer to improve bandwidth utilization. Its motivation is the poor scaling of naïve small-message aggregation in multi-GPU systems.

  • 2DH All-to-All: 2DH All-to-All aggregates chunks sent from multiple local GPUs to the same remote GPU before networking.Merging small chunks into larger messages improves link bandwidth utilization.
  • Aggregation challenge: Naïve local aggregation requires intra-node exchange of S/m chunks among m local GPUs, creating additional communication work.The process is equivalent to performing the intra-node All-to-All n/m times.
  • Aggregation challenge: The intra-node aggregation latency increases with scale because GPUs perform n/m non-contiguous memory accesses.For S = 128 MiB and m = 8, latency rises from ∼600µs at n = 8 to ∼5ms at n = 2048.
  • Communication context: The discussion focuses on NCCL because it outperforms MPI in most deep-learning scenarios, while MPI targets smaller traditional HPC workloads.This bounds the communication comparison to the stated deep-learning setting.
  • Linear All-to-All: Linear All-to-All splits S bytes across n GPUs and communicates S/n-byte chunks between GPU pairs.The algorithm uses point-to-point NCCL sends and receives.

A.3 Algorithm

2DH All-to-All reorganizes communication through stride copies and hierarchical intra- and inter-node exchanges, reducing non-contiguous memory-access costs as systems scale. Its benefits grow for smaller per-GPU data sizes and larger GPU counts.

  • 2DH All-to-All: 2DH All-to-All first aligns local-destination chunks, performs intra-node exchange, aligns remote-destination chunks, and then performs inter-node exchange.Stride memory copies make chunks contiguous before each communication phase.
  • 2DH All-to-All: 2DH All-to-All keeps latency constant and low during its first three phases by using stride memory copies with high memory-bandwidth utilization.The design also avoids cross-rail communication, benefiting rail-optimized InfiniBand networks.
  • Scalability: 2DH All-to-All becomes increasingly advantageous as S/n decreases, corresponding to smaller data sizes or larger GPU counts.The improvement is attributed to avoiding inefficient non-contiguous accesses and cross-rail communication.
  • Implementation: The NCCL implementation separates local communication and inter-node communication, making the latter depend on n/m after local chunks are merged.The first step covers intra-node exchange and two stride copies; the second performs inter-node exchange.
  • Extensions: 2DH All-to-All can be extended to larger local GPU groups using NVSwitch and adapted to three-dimensional network hierarchies such as dragonfly.NVSwitch enables configurations with m = 256 GPUs connected through high-speed NVLink.

A.5 Evaluation

The evaluation benchmarks All-to-All implementations from 64 to 4096 GPUs and shows that 2DH, optimized compilation, and protocol selection each affect scalability by workload size and system scale. These results motivate dynamic adaptation among algorithms and protocols.

  • All-to-All Evaluation: 64- to 4096-GPU experiments compare the proposed 2DH All-to-All with NCCL’s default linear All-to-All algorithm.The benchmark uses nccl-tests and measures latency across All-to-All sizes from 1 KiB to 16 GiB.
  • All-to-All Evaluation: 2DH All-to-All scales better with lower gradient than the linear algorithm and can scale to 4096 GPUs.For 1 MiB, 2DH has lower latency from small scales; for 32 MiB and 256 MiB, extra copies initially increase latency.
  • Adaptive Selection: Dynamic selection between linear and 2DH All-to-All is required because the faster algorithm depends on message size and GPU count.2DH improves as GPU counts scale out, while extra copies can hurt performance for larger messages at smaller scales.
  • Protocol Selection: LL128 has lower latency for 1 MiB and 32 MiB messages, whereas the default protocol performs better for 256 MiB messages.Protocol choice therefore must adapt to message size.

B SIMT-EFFICIENT FAST ENCODE AND DECODE

TUTEL replaces dense MoE encode/decode operations with differentiable sparse GPU kernels designed for SIMT efficiency. The implementation reduces latency and memory use while supporting end-to-end vision evaluation, including COCO detection.

  • Motivation: Existing MoE encode and decode implementations use computationally heavy einsum and matrix-multiplication operations for dispatch and combine.The dense encode path is identified as unnecessarily dense because only selected token-expert positions are relevant.
  • Limitation: Sparse GPU kernels remain challenging to implement efficiently because they cannot leverage matrix-multiplication accelerators as effectively as dense computation.The paper notes that existing sparsity support does not efficiently address this coarse-grained sparsity pattern.
  • Fast Operators: TUTEL’s fast encode and decode use three differentiable GPU kernels, K0, K1, and K2, for sparse dispatch and combine computation.The operators expose moe.fast_encode for dispatch and moe.fast_decode for combine.
  • Kernel Design: Assigning different token indices to thread arrays makes computation along dimension M SIMT-efficient and enables warp shuffling, Blelloch scans, and low-precision vectorization.These optimizations allow sparse computation to use techniques commonly applied to dense computation.
  • Memory Efficiency: 20%–90% memory saving is achieved in most tested cases for a single MoE layer.The memory comparison uses static settings D = H = 4096, top-k = 2, and E_g = 2.
  • COCO Evaluation: Direct COCO fine-tuning causes -1.7/-1.4 box/mask AP drops, while fixing all MoE layers yields +0.4/+0.4 box/mask AP improvements over dense counterparts.The fixed-MoE fine-tuning strategy alleviates the degradation observed with direct fine-tuning.

C.2 Ablation Study

The ablation study examines expert count, routing with batch prioritized routing, and router capacity settings across SwinV2-MoE configurations and vision tasks. It identifies settings that improve accuracy or the speed-accuracy trade-off.

  • Number of Experts: 32 and 64 experts perform best across SwinV2-S and SwinV2-B models and a variety of vision tasks.This pattern is consistent with previous works cited by the paper.
  • Routing Algorithms: Batch prioritized routing is crucial for computer-vision MoE models, especially at lower capacity factors.The comparison evaluates routing methods with and without BPR.
  • Router Settings: A top-1 router provides a better speed-accuracy trade-off, motivating the default settings k = 1 and f = 1.0.The ablation varies both router k and capacity factor f.

C.3 A New Cosine Router Supported in TUTEL

TUTEL supports a cosine router as an algorithmic alternative for MoE routing, aiming to improve numerical stability as model size increases. Preliminary results find it matches a common linear router in image classification, while its normalization may improve routing stability under feature scaling.

  • The cosine router is introduced as an algorithmic option in TUTEL, inspired by prior work and intended to improve numerical stability with increased model size.
  • P denotes routing scores used to select experts after projecting input features into dimension D and comparing them with expert parameters.W projects x from dimension C to D; M contains expert representations, and τ is a learnable temperature bounded below by 0.01.
  • With 32 experts, the cosine router is as accurate as a common linear router for image classification.The comparison uses E = 32, k = 1, and f = 1.25.
  • Input normalization may yield more stable routing when feature amplitude or dimension is scaled.The passage presents this as a motivation for trying the cosine router in other problems, not as a demonstrated superiority in image classification.
Loading 2206.03382v2…