PaperlayerView exact revision on arXiv

Source-linked AI summary

MegaBlocks: Efficient Sparse Training with Mixture-of-Experts

Trevor Gale, Deepak Narayanan, Cliff Young, Matei Zaharia

arXiv:2211.15841v1cs.LGcs.AIcs.DC

TL;DR

Existing MoE systems constrain dynamic routing to fit hardware and software, forcing a tradeoff between dropping tokens and wasting computation on padding. MegaBlocks reformulates MoE computation with block-sparse operations and GPU kernels, never drops tokens, and achieves up to 40% faster MoE training than Tutel and 2.4× faster DNN training than Megatron-LM.

  • Problem

    Existing MoE frameworks constrain dynamic routing to fixed-size expert assignments, forcing users to trade token dropping against computation and memory wasted on padding.

  • Method

    MegaBlocks reformulates MoE computation as block-sparse operations and provides GPU kernels that handle dynamic, imbalanced routing.

  • Results

    MegaBlocks enables up to 40% end-to-end training speedup over Tutel MoEs and 2.4× over Megatron-LM DNNs.

  • Takeaways & Limitations

    MegaBlocks provides dropless MoE training that maps efficiently to modern hardware accelerators.

  • Takeaways & Limitations

    Weight-gradient DSTD operations have reduced throughput from poor spatial locality, though their end-to-end impact is minimal.

Abstract

from arXiv · show

We present MegaBlocks, a system for efficient Mixture-of-Experts (MoE) training on GPUs. Our system is motivated by the limitations of current frameworks, which restrict the dynamic routing in MoE layers to satisfy the constraints of existing software and hardware. These formulations force a tradeoff between model quality and hardware efficiency, as users must choose between dropping tokens from the computation or wasting computation and memory on padding. To address these limitations, we reformulate MoE computation in terms of block-sparse operations and develop new block-sparse GPU kernels that efficiently handle the dynamism present in MoEs. Our approach never drops tokens and maps efficiently to modern hardware, enabling end-to-end training speedups of up to 40% over MoEs trained with the state-of-the-art Tutel library and 2.4x over DNNs trained with the highly-optimized Megatron-LM framework.

1 INTRODUCTION

MegaBlocks addresses the hardware and software difficulty of efficiently training dynamically routed MoEs. It reformulates MoE computation with block-sparse operations and GPU kernels, avoiding token dropping while improving training speed.

  • Structured sparsity reduces DNN computation, but fine-grained sparse training remains inefficient on GPUs and TPUs.
  • MoEs dynamically route tokens to expert-specific DNNs, enabling reduced training times and models exceeding 1 trillion parameters.
  • Dynamic routing and load imbalance do not map cleanly to existing accelerator hardware, compilers, and deep-learning software primitives.
  • Rigid routing implementations trade model quality against hardware efficiency by dropping tokens or wasting computation and memory on padding.
  • MegaBlocks never drops tokens and achieves up to 40% and 2.4× end-to-end training speedups over Tutel MoEs and Megatron-LM DNNs, respectively.
  • MegaBlocks expresses MoE computation as block-sparse operations and develops GPU kernels for dynamic, imbalanced expert assignments.

2 BACKGROUND: MOE LAYERS

MoE layers use learned routers to assign tokens dynamically among small expert networks, but fixed-shape batched computation creates capacity, padding, and token-dropping tradeoffs. Capacity increases can improve loss while adding computation.

  • MoE layers: MoE layers comprise small expert networks, with each token dynamically routed to a small subset of experts.
  • Routing: Tokens are commonly routed to between 1 and 4 experts, whose outputs are combined using router-produced assignment probabilities.
  • Routing: Learned routing projects token representations to expert scores, applies softmax normalization, and selects the top k experts.
  • Expert computation: Batched matrix multiplication requires equal expert shapes and equal token counts, although learned routing provides no load-balance guarantee.
  • Capacity: Increasing expert capacity significantly lowers loss but increases computation; dynamic maximum capacity avoids token dropping.
  • Capacity: Fixed expert capacity drops excess tokens and pads underfilled experts, creating a tradeoff between additional computation and model quality.

3 MOTIVATION: TOKEN DROPPING IN MOES

Token routing remains imbalanced despite load-balancing methods, making fixed expert capacities costly: dropping tokens harms quality, while avoiding drops increases computation, memory, and tuning burden.

  • Experimental Setup: Experiments trained Transformer MoEs on The Pile for 10B tokens across capacity factors 1, 1.5, 2, and dynamic capacity.The models used 64-expert layers, top-1 routing, sequence length 1024, and a single A100 GPU.
  • Token Dropping and Model Quality: 1.73× larger validation-loss reduction came from avoiding token dropping than from capacity factor 1, exceeding Transformer-Medium quality.The capacity-factor-1 MoE reduced validation loss by 0.15, whereas the no-dropping MoE reduced it by 0.26.
  • Capacity Trade-offs: Over 2× more MoE-layer math operations were required to avoid dropping tokens in this example.Some MoEs reportedly require capacity factors as high as 11, and the necessary factor can spike unpredictably during training.
  • Capacity Trade-offs: Capacity-factor tuning trades additional computation and memory against model quality and reduces the chance of dropping tokens.The extra hyperparameter can substantially increase the number of models trained for a target task.

4 NO-TOKEN-LEFT-BEHIND WITH BLOCK SPARSITY2

MegaBlocks formulates dynamically routed MoE computation as block-sparse operations, allowing imbalanced token assignments without dropping tokens while targeting efficient accelerator execution.

  • Motivation: Block-sparse operations provide a flexible representation of dynamic, load-imbalanced MoE computation and map to GPU and TPU systolic-array hardware.The formulation is designed to avoid dropping tokens while retaining hardware-compatible matrix operations.
  • Block-Sparse Formulation: The formulation exposes MoEs as dynamic, structured activation sparsity.
  • Block-Sparse Formulation: MoE computation can be represented as an SDD with block-diagonal sparse output, using variable row counts to represent imbalanced expert assignments.Each variable-sized block is computed from many smaller fixed-size blocks using block-sparse operations.
  • Scope Boundary: MegaBlocks did not explore variable-sized experts because more research is needed to determine how that capability could improve efficiency.
  • Hardware Mapping: MoE blocks are substantially larger than dense-matrix tile dimensions, providing flexibility to choose block sizes that match dense-kernel throughput.Transformer MoE blocks commonly use hidden dimensions from 1024 to 8192 and thousands to tens of thousands of tokens per expert.
  • Kernel Configuration: 128x128 tiles perform consistently on-par or better than other supported tile configurations on an A100 SXM4 80GB GPU.The benchmark used CUDA 11.5 and CUTLASS 2.5.

5 MEGABLOCKS: A FRAMEWORK FOR EFFICIENT MOE TRAINING

MegaBlocks implements MoE training with custom block-sparse GPU primitives that accommodate dynamic, imbalanced expert assignments and transposed computations. Its design combines suitable block sizes with hybrid sparse metadata to support efficient forward and backward operations.

  • 5.1.1 Existing Block-Sparse Primitives: MegaBlocks uses custom SDD, DSD, and DDS GPU primitives because existing libraries lack required operations or transposition support for dynamic MoE computation.The kernels support transposed and non-transposed inputs.
  • 5.1.2 Selecting Block Size for MoEs: 128x128 blocks were selected because they consistently performed on-par or better than other tested tile configurations.The highest-performing tile dimensions in the studied workloads were also 128x128 for 128x128 blocks.
  • 5.1.3 Sparse Matrix Format: BCSR enables efficient row-wise iteration over nonzero blocks, while block metadata keeps identifying block positions inexpensive.The format supports operations such as DSD and DDST.
  • 5.1.3 Sparse Matrix Format: Materialized row indices let threadblocks directly locate sparse output blocks for parallel SDD operations with negligible additional storage.The hybrid encoding preserves row-wise ordering and supports both BCSR and blocked coordinate formats.
  • 5.1.4 Block-Sparse Transposition With Transpose Indices: A dMoE constructs sparse topology from expert assignments, pads expert batches to block-size multiples, and computes experts by iterating SDD and DSD operations.These changes adapt the standard MoE computation to the block-sparse implementation.
  • 5.1.4 Block-Sparse Transposition With Transpose Indices: Transpose indices enable BCSR matrices to be iterated in transposed order without copying the nonzero values.The indices store each nonzero block’s memory offset in transposed order and provide access through indirection.

6 EXPERIMENTS

MegaBlocks is evaluated against Tutel and Megatron-LM through microbenchmarks and end-to-end language-model training, showing higher efficiency while avoiding token dropping. Its advantage grows with model size, although limited memory constrains medium-model efficiency.

  • 6.1 MoE Training Without Dropping Tokens: 1.38×, 2.0×, and 4.35× end-to-end speedups were achieved for MoE-XS, MoE-Small, and MoE-Medium versus Tutel’s padding-based approach.Tutel’s activation-memory overhead also reduced its maximum micro batch size by 2×, 4×, and 8× respectively.
  • 6.1 MoE Training Without Dropping Tokens: MegaBlocks’ advantage over Tutel increased with model size across MoE-XS, MoE-Small, and MoE-Medium.The experiments used The Pile and compared dMoEs with padding-based and token-dropping MoEs.
  • 6.1 MoE Training Without Dropping Tokens: 1.8×–2.4× less training time was required for MegaBlocks dMoEs to reach a given validation loss than for Megatron-LM Transformers.The range reflects differing weight-memory usage and micro batch sizes across MoE models.
  • 6.1 MoE Training Without Dropping Tokens: Megatron-LM sustained 21%–48% of the 2.5 petaFLOP peak throughput on the evaluated eight-GPU system.Efficiency increased with model size, while MegaBlocks dMoEs reduced time to a given validation loss relative to these Transformers.
  • 6.2 MoE Training With Token Dropping: MoE-Medium incurred some implementation-efficiency loss because limited GPU memory forced a relatively small micro batch size.Smaller kernel tile dimensions or increased memory capacity could improve performance in this setting.
  • 6.3 Block-Sparse Matrix Multiplication Performance: 98.6% of cuBLAS throughput was achieved on the evaluated block-sparse matrix-multiplication problems, with a 4% standard deviation.Relative throughput ranged from 91% to 104%, and the kernels slightly outperformed cuBLAS on half the problems.

7 RELATED WORK

Related work addresses MoE routing, distributed training, and sparse kernels through balancing, communication, and matrix-format techniques. MegaBlocks is positioned as complementary to improved routing and distributed-training methods.

  • MoE Routing: BASE layers formulate routing as a linear assignment problem that maximizes token–expert affinities under perfectly balanced assignments.This guarantees that tokens are not dropped by rerouting them as needed.
  • MoE Routing: Approximate Sinkhorn routing reduces the overhead of BASE layers but no longer guarantees avoiding token dropping, so experiments use capacity factor 2.Hash-based static mappings and expert-selected top-k routing provide alternatives but have reported performance or token-dropping issues.
  • High-Performance MoEs: Tutel and FasterMoE improve scalable MoE training through optimized communication, communication hiding, routing changes, and network-congestion avoidance.The authors state that MegaBlocks could additionally benefit from these techniques.
  • Sparse Kernels: Sparse matrix formats supporting efficient transposed access provide prior foundations for adapting sparse computation to large block sparsity on modern GPUs.The paper identifies this adaptation as a direction for future research.

8 CONCLUSION

MegaBlocks reformulates MoE computation with block-sparse operations and GPU kernels that handle dynamic routing without dropping tokens. The system reports end-to-end speedups over Tutel-trained MoEs and Megatron-LM-trained dense DNNs.

  • 8 CONCLUSION: MegaBlocks reformulates MoEs as block-sparse operations and introduces GPU kernels for dynamic MoE computation.The system is designed to avoid token dropping while mapping MoE computation efficiently to modern hardware accelerators.
  • 8 CONCLUSION: Up to 40% and 2.4× end-to-end training speedups were reported over Tutel MoEs and Megatron-LM DNNs, respectively.These are the paper’s reported headline comparisons.
Loading 2211.15841v1…