Mixture of Experts: Advanced Architecture, Training, and Scaling
Chapter 1: Foundations of Sparse Expert Models
Conditional Computation Principles
The Sparse MoE Approach
Contrasting Dense vs. Sparse Activation
Mathematical Formulation of Basic MoE Layers
Chapter 2: Advanced MoE Architectures
Designing Effective Gating Networks
Hierarchical MoE Structures
Router Architectures: Linear, Non-Linear, Attention-Based
Expert Capacity and Sizing Considerations
Stabilization Techniques for Routers
Hands-on Practical: Implementing Custom Gating Mechanisms
Chapter 3: Training Dynamics and Optimization
The Load Balancing Problem in MoE
Auxiliary Loss Functions for Load Balancing
Router Optimization Strategies
Handling Dropped Tokens
Expert Specialization Collapse and Prevention
Impact of Optimizer Choice and Hyperparameters
Hands-on Practical: Implementing and Tuning Load Balancing Losses
Chapter 4: Scaling MoE Models: Distributed Training
Challenges in Distributed MoE Training
Expert Parallelism: Distributing Experts Across Devices
Integrating Expert Parallelism with Data Parallelism
All-to-All Communication Patterns
Pipeline Parallelism for MoE Models
Communication Optimization Techniques (e.g., Overlapping)
Frameworks and Libraries for Distributed MoE (e.g., DeepSpeed, Tutel)
Practice: Configuring Distributed MoE Training
Chapter 5: Inference Optimization and Deployment
Inference Challenges with Sparse Models
Batching Strategies for MoE Inference
Model Compression Techniques for MoE
Hardware Acceleration Considerations
Router Caching and Optimization
Deployment Patterns for Large Sparse Models
Hands-on Practical: Profiling MoE Inference
Hardware Acceleration Considerations
Was this section helpful?
- Sparsely-Gated Mixture-of-Experts Layers, Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, Jeff Dean, 2017 arXiv preprint arXiv:1701.06538 DOI: 10.48550/arXiv.1701.06538 - Introduces the concept of sparsely-gated Mixture-of-Experts, outlining the architecture and the challenges associated with their sparse nature and dynamic routing.
- Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity, William Fedus, Barret Zoph, Noam Shazeer, 2022 Journal of Machine Learning Research, Vol. 23 (JMLR) DOI: 10.55986/joxg921 - Describes a simplified MoE architecture and discusses practical engineering challenges for scaling, including communication and memory optimization for distributed training and inference.
- A Domain-Specific Architecture for Deep Neural Networks, Norman P. Jouppi, Cliff Young, David Patil, Dustin Patterson, Gaetano Agostini, Shumeet Baluja, Keren Bergman, Ry Chiang, Sheng Li, Mike Ni, Vijay Nivargi, Paul Norman, Mike Reddi, Kevin Smith, David Sprague, Greg Thorson, Rajat Wadia, Kevin Walker, David Wang, Hongbo Wei, Christof Zabriskie, 2017 ACM SIGARCH Computer Architecture News, Vol. 45 (ACM) DOI: 10.1145/3144819.3144824 - Describes the architecture of Google's Tensor Processing Unit (TPU), explaining its systolic array design and high-bandwidth memory, which are beneficial for accelerating machine learning workloads.
- LLM.int8(): 8-bit Matrix Multiplication for Transformers at Scale, Tim Dettmers, Mike Lewis, Younes Belkada, Luke Zettlemoyer, 2022 NeurIPS 2022 DOI: 10.48550/arXiv.2208.07339 - Presents a method for performing 8-bit matrix multiplication for large transformer models, significantly reducing memory footprint and computation for inference, with implications for MoE models.