Source-linked AI summary
Accelerating Distributed MoE Training and Inference with Lina
Jiamin Li, Yimin Jiang, Yibo Zhu, Cong Wang, Hong Xu
TL;DR
Distributed MoE training and inference are hindered by interleaved all-to-all communication, motivating more efficient communication and scheduling. Lina analyzes these bottlenecks, prioritizes all-to-all over allreduce during training, and dynamically balances inference resources using expert-selection patterns. On A100 GPUs, Lina improves training step time and 95%ile inference time over state-of-the-art systems.
Problem
Interleaved all-to-all communication makes distributed MoE training and inference inefficient despite MoE's lower computation scaling than dense models.
Method
Lina prioritizes all-to-all over concurrent allreduce with tensor partitioning during training and dynamically schedules inference resources using expert-selection patterns.
Results
Lina reduces training step time by up to 1.73x and 95%ile inference time by an average of 1.63x over state-of-the-art systems.
Takeaways & Limitations
Lina improves distributed MoE training efficiency and inference time by targeting all-to-all communication with training prioritization and inference workload balancing.
Takeaways & Limitations
Naively prioritizing all-to-all is inefficient because current multi-GPU communication libraries cannot preempt an allreduce after its transmission strategy is launched.
Abstract
from arXiv · showhide
Scaling model parameters improves model quality at the price of high computation overhead. Sparsely activated models, usually in the form of Mixture of Experts (MoE) architecture, have sub-linear scaling of computation cost with model size, thus providing opportunities to train and serve a larger model at lower cost than their dense counterparts. However, distributed MoE training and inference is inefficient, mainly due to the interleaved all-to-all communication during model computation. This paper makes two main contributions. First, we systematically analyze all-to-all overhead in distributed MoE and present the main causes for it to be the bottleneck in training and inference, respectively. Second, we design and build Lina to address the all-to-all bottleneck head-on. Lina opportunistically prioritizes all-to-all over the concurrent allreduce whenever feasible using tensor partitioning, so all-to-all and training step time is improved. Lina further exploits the inherent pattern of expert selection to dynamically schedule resources during inference, so that the transfer size and bandwidth of all-to-all across devices are balanced amid the highly skewed expert popularity in practice. Experiments on an A100 GPU testbed show that Lina reduces the training step time by up to 1.73x and reduces the 95%ile inference time by an average of 1.63x over the state-of-the-art systems.
1 Introduction
MoE reduces computation scaling by activating only selected experts, but distributed training and inference remain inefficient because all-to-all communication becomes a bottleneck. Lina addresses training contention with allreduce and inference imbalance through prioritized communication and popularity-aware resource scheduling.
- Motivation: MoE activates only a few experts per input, yielding sub-linear FLOP scaling with model size.The architecture selects a small number of experts rather than applying all parameters.
- Distributed MoE: Distributed MoE uses data and expert parallelism, exchanging tokens with all-to-all and aggregating non-expert gradients with allreduce.All-to-all serves expert-parallel communication, while allreduce handles non-expert gradients during backward propagation.
- Bottleneck: All-to-all is a bottleneck because training overlaps it with bandwidth-contending allreduce, while inference faces skewed expert popularity and imbalanced transfers.The two phases have different primary causes of all-to-all inefficiency.
- Training: Lina prioritizes all-to-all over allreduce during training using tensor partitioning, micro-op scheduling, and pipelining.Partitioned tensors allow allreduce to make progress when all-to-all is absent and allow expert computation to overlap with communication.
- Inference: Lina dynamically schedules inference resources using estimated expert popularity, allocating more resources to popular experts and refining allocations when estimates deviate.Its two-phase scheduler limits fine-tuning to cases where the actual popularity differs substantially from the estimate.
- Results: Lina achieves 1.57x average training-step speedup and reduces median and 95%ile inference time by 1.45x and 1.63x, respectively, versus DeepSpeed.The evaluation uses up to 16 A100 GPUs with 100Gbps InfiniBand.
2 Background and Motivation
MoE replaces dense feed-forward computation with sparsely selected experts, but its required all-to-all exchanges consume substantial time. Training suffers from contention with allreduce, while inference suffers from skewed expert popularity and uneven device utilization.
- MoE Primer: MoE layers replace Transformer feed-forward networks with multiple expert FFNs and a gating network that dispatches each token to a small number of experts.The output combines the selected experts’ outputs using their gating weights.
- MoE Primer: Load-balancing loss encourages a uniform training-time token distribution across experts but does not enforce perfect balance.Inference uses the trained gating network to dispatch tokens based on their embeddings.
- All-to-All Bottleneck: All-to-all is a synchronous MoE component used twice per layer to dispatch tokens and restore their positions, taking an average of 34.1% of step time.Its large data transfers make it a shared bottleneck in both training and inference.
- Training Bottleneck: During training, concurrent allreduce and all-to-all contend for network bandwidth, prolonging the blocking communication needed for expert gradients.Hybrid data and expert parallelism creates this contention in the backward pass.
- Inference Bottleneck: During inference, expert popularity is highly skewed: the most popular expert receives 4.02x and 5.56x the tokens of the least popular expert in 4- and 16-expert tasks.The resulting imbalance leaves devices hosting less popular experts idle while waiting for others.
3 Design Overview
Lina targets distributed MoE all-to-all communication by coordinating it differently in training and inference. It prioritizes communication over competing allreduce during training and balances inference workload using expert-selection patterns.
- Training: Lina prioritizes all-to-all over concurrent allreduce during training to improve bandwidth and reduce communication blocking.The design uses tensor partitioning to support this priority strategy.
- Inference: Lina dynamically balances inference workload using token-level expert-selection patterns to address unbalanced all-to-all bandwidth.Resource scheduling targets the straggler effect created by uneven expert popularity.
4 Prioritizing All-to-All Training
Lina addresses training inefficiency by partitioning communication into micro-operations that prioritize all-to-all while preserving useful allreduce and computation overlap. This design avoids the poor scheduling behavior of naive priority schemes and pipelines communication with computation.
- Design Challenge: Naively prioritizing all-to-all can lengthen the first all-to-all and overall training step because NCCL operations cannot be preempted.Communication strategies are fixed when primitives are launched, preventing later all-to-all operations from preempting an active allreduce.
- Design Challenge: Deferring allreduce until multiple all-to-all operations finish can delay optimization and undermine a wait-free backward pass in unfavorable cases.If allreduce remains blocked across the current step’s all-to-all operations, devices must wait before optimization begins.
- Tensor Partitioning: Tensor partitioning breaks communication operations into small, uniform micro-ops that can be prioritized independently.Lina partitions gradient tensors into equal-sized chunks rather than fusing them into variable-sized allreduce buckets.
- Scheduling: Lina schedules allreduce micro-ops in gaps around all-to-all operations, preserving work conservation while reducing bandwidth contention.The scheduler can launch some allreduce work before an all-to-all arrives and defer subsequent micro-ops until all-to-all has priority.
- Pipelining: Micro-ops also enable Lina to pipeline expert computation with all-to-all communication.The scheduling design coordinates partitioned communication with computation instead of treating the operations as indivisible phases.
5 Scheduling Resources in Inference
Lina schedules inference resources around skewed and layer-varying expert popularity. It estimates future selections from token-level cross-layer patterns, allocates devices accordingly, and fine-tunes the mapping after actual gating results arrive.
- Design Challenge: Skewed expert popularity creates unequal token processing times and imbalanced all-to-all bandwidth in MoE inference.The mismatch between token-level expert processing and sequence-level attention contributes to the imbalance.
- Design Challenge: Expert popularity varies across sequences and MoE layers, so resource allocation must be dynamic and layer-specific.Measurements across inference batches show different popular experts across layers of the same task.
- Selection Pattern: Tokens selecting the same expert in one layer tend to select the same expert again in the next layer, providing a basis for prediction.Lina exploits this token-level selection pattern to estimate the next layer’s expert distribution before gating completes.
- Popularity Estimation: A longer profiling path improves next-layer popularity estimates but increases data-collection and computation costs.The path length controls the accuracy–cost tradeoff in profiling.
- Popularity Estimation: Lina estimates popularity from top-k experts on each sample path before MoE computation, focusing resources on experts expected to demand the most.The remaining experts are treated as lower-popularity candidates for scheduling.
- Two-Phase Scheduling: In phase one, popular experts are replicated across more devices while unpopular experts are packed onto fewer devices.The allocation uses estimated batch-level popularity and a first-fit-decreasing heuristic for packing.
- Two-Phase Scheduling: In phase two, Lina compares estimated and actual top-2k experts and recomputes resource allocation only when they differ significantly.When the lists match, inference continues without fine-tuning.
6 Implementation
Lina integrates with DeepSpeed MoE and PyTorch through local communication scheduling and a dedicated inference resource scheduler. Its implementation coordinates expert mappings and load balancing through existing communication operations.
- Implementation: Lina is implemented in C++ and Python on DeepSpeed MoE and PyTorch, using PyTorch 1.10, CUDA 11, and NCCL 2.10.The implementation comprises approximately 7,500 lines of code.
- Training Scheduler: Each device runs a local communication scheduler with a priority queue for micro-operations.The scheduler requires no coordination across devices because its scheduling scope is local.
- Inference Scheduler: The inference resource scheduler runs on device 0 and manages expert placement using layer-wise popularity distributions.Devices retain expert weights in host DRAM and can limit GPU loading when memory is scarce.
- Inference Coordination: Lina piggybacks popularity estimates and expert-device mapping information on regular all-to-all operations to reduce coordination overhead.The scheduler communicates mappings and next-layer coordination data through the first and second all-to-all operations.
- Inference Coordination: Phase-two fine-tuning sends actual popularity separately and broadcasts a revised mapping only when needed.If no fine-tuning is required, the scheduler broadcasts a resume signal; otherwise computation waits for the updated command.
- All-to-All Coordination: Inference uses unequal-split all-to-all, including placeholder pointers when no tokens target a device.This avoids initializing multiple process groups for different transfer sizes.
7 Evaluation
Lina is evaluated on A100 GPUs across training and inference MoE workloads, with results showing faster communication, training, and inference alongside scheduling trade-offs.
- Overall Performance: Lina reduces training step time over DeepSpeed by 1.37x–1.73x across models with 2–16 experts.The largest gains occur with 2- and 8-expert models, where expert packing uses two experts per device.
- Overall Performance: Lina improves GPU utilization by 17.6%, while expert packing raises BERT2GPT2 peak memory by 19.5%.Transformer-XL and GPT-2 use all memory and apply DRAM-offloading for packed expert parameters.
- Communication Scheduler: 2.21x, 2.39x, and 2.31x average all-to-all speedups are achieved for 4-, 8-, and 16-expert cases, respectively.The backward-pass scheduler prioritizes all-to-all and avoids concurrent allreduce; expert packing also reduces transfer size.
- Communication Scheduler: Tensor partitioning reduces step time over Baseline by 1.36x, 1.36x, 1.41x, and 1.42x for 2-, 4-, 8-, and 16-expert cases.Priority scheduling alone averages 24% gain, while pipelining adds limited benefit without expert packing.
- Resource Scheduler: Lina reduces 95%ile inference time by 1.82x for 16-expert Transformer-XL and 1.68x for 16-expert BERT-Large.The scheduler balances device load using estimated expert popularity, bringing inference time closer to Ideal.
- Resource Scheduler: Without fine-tuning, tail inference time increases by 26.7% and 33.1% for 16-expert Transformer-XL and BERT-Large, respectively.The importance of fine-tuning depends on estimation accuracy and the number of experts.
8 Discussion
The discussion identifies training-parallelism scope and expert-popularity estimation as important boundaries for Lina’s applicability and performance.
- Lina focuses on sparsely activated MoE models using data and expert parallelism, which are orthogonal to pipeline and tensor parallelism work for dense models.The paper distinguishes Lina’s scope from existing coordination methods for dense-model parallelism.
- Lina’s expert-popularity estimation relies on data collected during training, leaving room for more accurate and confident inference-time predictions.The discussion suggests predicting each token’s selected expert ahead of time as one possible improvement.
9 Related Work
Prior systems improve distributed MoE through flexible parallelism, adaptive switching, performance modeling, expert-parameter movement, topology-aware selection, and heterogeneous-resource utilization. Lina instead prioritizes all-to-all over allreduce in MoE training.
- DeepSpeed enables distributed MoE training through flexible combinations of parallelism strategies and introduces Pyramid-Residual MoE.
- Tutel extends DeepSpeed with adaptive parallelism switching specialized for MoE training and hierarchical all-to-all support.
- FasterMoE models end-to-end MoE training performance with a roofline model, then uses dynamic shadowing and topology-aware expert selection.
- Lina differs from prior work by prioritizing all-to-all over allreduce in MoE training, while remaining compatible with these acceleration techniques.
- BytePS reduces communication traffic by using heterogeneous GPU/CPU resources, an approach that can leave more bandwidth for all-to-all operations.
10 Conclusion
Lina accelerates distributed MoE by improving all-to-all communication in training and balancing workloads during inference. Its A100 testbed evaluation shows improved training efficiency and inference time.
- Lina prioritizes all-to-all over allreduce using tensor partitioning and pipelining to improve communication bandwidth during training.
- Lina dynamically balances inference workload using token-level expert selection patterns.
- Lina was implemented over DeepSpeed and evaluated on A100 GPUs with 100Gbps InfiniBand.
- The evaluation shows that Lina significantly improves training efficiency and inference time.