PaperlayerView exact revision on arXiv

Source-linked AI summary

Pre-gated MoE: An Algorithm-System Co-Design for Fast and Scalable Mixture-of-Expert Inference

Ranggi Hwang, Jianyu Wei, Shijie Cao, Changho Hwang, Xiaohu Tang, Ting Cao, Mao Yang

arXiv:2308.12066v3cs.LGcs.AIcs.AR

TL;DR

MoE architectures scale model capacity without proportionally scaling computation, but their large memory footprint and sequential, dynamic expert activation hinder efficient inference. Pre-gated MoE co-designs a pre-gate function with a CPU-offloading system to reduce GPU memory use and improve throughput while maintaining comparable model accuracy across NLP tasks.

  • Problem

    MoE inference is constrained by large expert-parameter memory demands and sequential dynamic expert activation, which makes single-GPU deployment difficult and causes CPU-offloading overhead.

  • Method

    Pre-gated MoE uses a pre-gate function that selects experts for the next MoE block, eliminating the within-block dependency between expert selection and execution.

  • Results

    Pre-gated MoE improves inference throughput, significantly reduces GPU memory consumption, and achieves comparable model accuracy across various natural language processing tasks.

  • Takeaways & Limitations

    The co-design supports scalable, high-performance MoE inference and facilitates deployment across a wide range of real-world applications.

Abstract

from arXiv · show

Large language models (LLMs) based on transformers have made significant strides in recent years, the success of which is driven by scaling up their model size. Despite their high algorithmic performance, the computational and memory requirements of LLMs present unprecedented challenges. To tackle the high compute requirements of LLMs, the Mixture-of-Experts (MoE) architecture was introduced which is able to scale its model size without proportionally scaling up its computational requirements. Unfortunately, MoE's high memory demands and dynamic activation of sparse experts restrict its applicability to real-world problems. Previous solutions that offload MoE's memory-hungry expert parameters to CPU memory fall short because the latency to migrate activated experts from CPU to GPU incurs high performance overhead. Our proposed Pre-gated MoE system effectively tackles the compute and memory challenges of conventional MoE architectures using our algorithm-system co-design. Pre-gated MoE employs our novel pre-gating function which alleviates the dynamic nature of sparse expert activation, allowing our proposed system to address the large memory footprint of MoEs while also achieving high performance. We demonstrate that Pre-gated MoE is able to improve performance, reduce GPU memory consumption, while also maintaining the same level of model quality. These features allow our Pre-gated MoE system to cost-effectively deploy large-scale LLMs using just a single GPU with high performance.

I. INTRODUCTION

LLMs gain accuracy by scaling model capacity, but their compute and memory demands make deployment difficult. Pre-gated MoE co-designs preemptive expert selection with CPU offloading to reduce memory use while hiding migration latency.

  • I. INTRODUCTION: MoE scales model capacity with sublinear compute by sparsely activating experts, but this increases memory requirements and creates dynamic execution challenges.MoE models can have substantially more parameters than FLOPs-equivalent dense models, while runtime expert activation can leave GPUs underutilized.
  • I. INTRODUCTION: CPU offloading reduces GPU requirements, but migrating activated experts from CPU to GPU adds substantial inference latency and harms quality of service.The migration overhead arises because expert selection and execution remain sequentially dependent.
  • I. INTRODUCTION: Pre-gated MoE proposes an algorithm-system co-design that targets low GPU memory consumption, high performance, and reduced total cost of ownership.The stated design goal is to balance memory efficiency with performance for scalable MoE inference.
  • I. INTRODUCTION: The pre-gate function selects experts for the next MoE block, eliminating the same-block dependency between expert selection and execution.This shifts the data dependency across adjacent MoE blocks, enabling performance optimization.
  • I. INTRODUCTION: Pre-gated MoE prefetches only the next block’s activated experts, overlapping CPU-to-GPU migration with current expert execution.The system stores expert parameters in CPU memory and transfers selected experts to GPU memory before their use.

A. Dense LLMs using Transformers

Transformer-based dense LLMs face rapidly increasing compute and memory costs, while MoE architectures scale model size by activating only selected experts. This sparsity preserves compute efficiency but creates a large, dynamic memory footprint that complicates deployment.

  • A. Dense LLMs using Transformers: Dense transformer models scale capacity by stacking transformer blocks, but increasing model size raises compute and memory costs.A dense block includes self-attention, an FFN layer, normalization, and residual connections.
  • B. Sparse LLMs using Mixture-of-Experts (MoE): MoE replaces the dense transformer block’s FFN layer with multiple experts and a gate that selects experts for each input token.Each expert has the same dimension as the corresponding dense FFN layer.
  • B. Sparse LLMs using Mixture-of-Experts (MoE): Expert selection and execution form a sequential, input-dependent inference process, while sparse activation means only a small subset of experts is usually used.SwitchTransformer activates top-1 experts and NLLB-MoE activates top-2 experts.
  • B. Sparse LLMs using Mixture-of-Experts (MoE): MoE compute remains constant as the number of available experts increases, enabling model-size scaling without proportional FLOP growth.Figure 2 compares SwitchTransformer with FLOPs-equivalent dense T5 models across Base and Large configurations.
  • A. Key Challenges of MoE inference: 75× higher memory consumption than dense T5 can occur for SwitchTransformer, making single-GPU deployment difficult and reducing utilization in multi-GPU systems.MoE expert parameters account for most of the model’s memory consumption, and sparse activation can leave allocated parameters unused.

B. Prior Solution: CPU Offloading of Expert Parameters

CPU offloading addresses MoE’s large GPU-memory footprint by storing expert parameters outside the GPU, but transfer latency remains a central performance bottleneck. Prior designs either fetch activated experts on demand or prefetch all experts for the next block.

  • B. Prior Solution: CPU Offloading of Expert Parameters: CPU- and SSD-offloading methods reduce the GPUs needed for MoE deployment, but parameter-transfer latency can degrade end-to-end performance.These methods do not fundamentally resolve the dynamic and sparse nature of expert activation.
  • B. Prior Solution: CPU Offloading of Expert Parameters: MoE-OnDemand migrates only selected experts after runtime selection, improving GPU memory utilization but serializing migration with expert execution.This serialization causes noticeable performance overhead.
  • B. Prior Solution: CPU Offloading of Expert Parameters: MoE-Prefetch proactively migrates the next block’s expert parameters to hide CPU-to-GPU transfer latency before execution.The approach offloads all expert parameters to CPU memory and prefetches parameters before their actual use.
  • B. Prior Solution: CPU Offloading of Expert Parameters: The sequential dependency between expert selection and execution motivates system designs that overlap data movement with computation.The paper identifies balancing memory efficiency and high performance as the central system objective.

IV. PRE-GATED MOE: CO-DESIGNING ALGORITHM AND SYSTEM FOR FAST & SCALABLE MOE INFERENCE

Pre-gated MoE decouples expert selection from execution by predicting the next block’s active experts, enabling selective migration and overlapped computation.

  • A. High-level Overview: The approach is designed to address MoE’s large memory footprint and dynamic sparse activation while supporting large-scale LLM deployment on a single GPU.The paper presents Pre-gated MoE as an algorithm-system co-design for scalable, high-performance inference.
  • C. (System) Preemptive Expert Migration: Pre-gated MoE stores expert parameters in CPU memory and migrates only the activated experts to GPU memory before their block executes.This hierarchical storage approach reduces GPU memory requirements while retaining access to the full expert set.
  • C. (System) Preemptive Expert Migration: Pre-gating removes the sequential dependency between expert selection and execution, allowing expert migration or selection to overlap with computation in subsequent blocks.The resulting overlap targets the communication overhead introduced by CPU-offloaded experts.
  • B. (Algorithm) Pre-gated MoE Architecture: Pre-gated MoE’s pre-gate function selects experts for the next MoE block, rather than the current block, using a learned activation mask.This changes the gate’s role from immediate selection to preemptive selection for the subsequent block.
  • B. (Algorithm) Pre-gated MoE Architecture: The first MoE block uses separate gates for current and next-block experts, while the last block has no pre-gate because no subsequent block executes in that iteration.The design does not use pre-gating across decoder iterations.

C. (System) Preemptive Expert Migration

Preemptive expert migration overlaps CPU-to-GPU transfers with expert execution and transfers only selected experts, targeting lower migration overhead and GPU memory use.

  • C. (System) Preemptive Expert Migration: Pre-gated MoE overlaps communication-bound expert migration with compute-bound expert execution for all but the first MoE block.The first block remains sequential because it lacks a preceding pre-gate function.
  • C. (System) Preemptive Expert Migration: Pre-gated MoE addresses MoE-OnDemand’s exposed migration latency and MoE-Prefetch’s excessive transfer volume through selective, overlapped migration.The comparison is presented against GPU-only, MoE-OnDemand, and MoE-Prefetch execution timelines.
  • C. (System) Preemptive Expert Migration: Only activated experts are migrated from CPU memory, while the complete expert parameters remain offloaded, reducing the GPU memory needed for MoE deployment.Peak GPU memory is dominated by non-MoE parameters and active experts for current and subsequent blocks.
  • C. (System) Preemptive Expert Migration: The memory design facilitates deployment of considerably larger LLMs on systems with limited GPU memory resources.The paper identifies desktop and edge devices as examples of such systems.

V. METHODOLOGY

The evaluation uses single-GPU GPU-only and CPU-GPU systems, SwitchTransformer models, and downstream summarization and question-answering tasks.

  • System configuration: Evaluation compares GPU-only and CPU-GPU designs using an AMD EPYC 7V12 CPU, 1.8TB DDR4 memory, and one NVIDIA A100 GPU with 80GB HBM.The CPU and GPU communicate over PCIe Gen4 at 32 GB/sec.
  • System configuration: GPU-only stores all model parameters on the GPU, whereas CPU-GPU persistently stores dense non-MoE parameters on GPU and offloads sparse MoE parameters to CPU memory.Pre-gated MoE and both CPU-offloading baselines use the CPU-GPU configuration.
  • Model and dataset: Experiments use Google’s SwitchTransformer with pretrained HuggingFace weights fine-tuned for downstream tasks.The datasets cover summarization and closed-book question answering.
  • Model and dataset: The study evaluates summarization with ROUGE-2 and question answering with ExactMatch and F1 scores.These metrics correspond to the reported downstream task types.
  • Model training (fine-tuning): All model architectures receive the same fine-tuning configuration, including a minibatch of 256 sequences of length 256 and 2,048 training steps.The constant learning rate is 0.0001.
  • Software implementation: Inference performance is implemented with NVIDIA FasterTransformer and reported using a model fine-tuned for closed-book question answering on SQuAD.The implementation uses the same GPU-only and CPU-GPU system designs described for evaluation.

VI. EVALUATION

Pre-gated MoE improves inference performance while substantially reducing peak GPU memory, enabling scalable deployment of large MoE models.

  • A. Performance: Pre-gated MoE reduces MoE block latency by 1.7× versus MoE-OnDemand and 42× versus MoE-Prefetch, with only 19% overhead versus GPU-only.For Switch-Large, it achieves 1.9× and 125× latency reductions against the two CPU-GPU designs, while GPU-only runs out of memory.
  • A. Performance: 111 tokens/sec average throughput on Switch-Base, 1.5× over MoE-OnDemand and 27× over MoE-Prefetch, reaching 81% of GPU-only.For Switch-Large with 128 experts, throughput reaches 42 tokens/sec, 1.6× and 52× higher than the two CPU-offloading baselines.
  • B. Scalability: 23% of GPU-only peak GPU memory is consumed, only 0.2% more than memory-optimal MoE-OnDemand.CPU offloading becomes increasingly beneficial as the number of experts grows, while MoE-Prefetch still averages 51% of GPU-only memory usage.
  • B. Scalability: Pre-gated MoE combines near-GPU-only performance with MoE-OnDemand-like memory efficiency for deploying large LLMs.The evaluation targets single-batch inference, where performance depends primarily on hiding CPU-to-GPU expert migration within MoE execution.

C. Model Accuracy

Pre-gated MoE largely preserves model accuracy while its performance depends on sparse expert activation and its benefits extend to caching and SSD offloading.

  • C. Model Accuracy: Pre-gated MoE maintains competitive accuracy across model sizes and downstream tasks, with only small degradations on some larger configurations.For the smallest Switch-Base configuration, it consistently shows slightly higher accuracy; observed variances do not indicate substantial capability changes.
  • D. Discussion: N=1 pre-gating achieves the highest accuracy among conventional MoE and N=2/3 pre-gating configurations.Accuracy gradually decreases as the pre-gate selects experts farther ahead.
  • D. Discussion: Increasing activated experts reduces the performance advantage of CPU offloading, while Pre-gated MoE remains superior for sparse activation.As activation approaches 100%, MoE behavior becomes dense and MoE-Prefetch overfetches fewer unnecessary experts relative to the activated set.
  • D. Discussion: Expert caching benefits both Pre-gated MoE and MoE-OnDemand, but helps MoE-OnDemand more because Pre-gated MoE already hides most migration latency.The comparison includes multiple cache replacement policies and varying GPU cache fractions.
  • D. Discussion: SSD offloading reduces Pre-gated MoE’s relative performance benefit because slower SSD-to-GPU migration becomes the dominant bottleneck.Even under this constraint, Pre-gated MoE consistently outperforms the evaluated baselines.

VII. RELATED WORKS

Prior work spans MoE training and inference systems, parameter offloading, and compact model architectures; Pre-gated MoE targets scalable CPU-offloaded inference through pre-gating.

  • VII. RELATED WORKS: Prior MoE inference systems optimize multi-GPU communication and kernels or offload tensors and parameters to CPU memory and NVMe SSD.HuggingFace Accelerate and SE-MoE implement the MoE-OnDemand and MoE-Prefetch baselines evaluated in this paper.
  • VII. RELATED WORKS: Compact MoE architectures reduce model size through architectural changes and knowledge distillation, but may degrade accuracy or target GPU-only deployment.These approaches differ from Pre-gated MoE’s CPU-offloading design.
  • VII. RELATED WORKS: Pre-gated MoE addresses MoE’s memory footprint and dynamic sparse activation, improving throughput and reducing GPU memory while maintaining comparable accuracy.The paper positions this algorithm-system co-design for scalable, high-performance inference across NLP tasks.

APPENDIX

The artifact provides a C++ and Python implementation of Pre-gated MoE built on SwitchTransformer and FasterTransformer, with substantial hardware and storage requirements.

  • APPENDIX: The artifact implements Pre-gated MoE using Google’s SwitchTransformer and NVIDIA’s FasterTransformer in C++ and Python.It is intended to reproduce latency, throughput, and peak GPU memory results against three baselines.
  • APPENDIX: Reproduction requires at least one GPU with 40GB memory, a CPU with 128GB memory, and more than 100GB of disk storage.Recent GPUs such as the NVIDIA A100 with 80GB HBM are recommended for optimal performance.
  • APPENDIX: The provided Docker image and installation script automate software setup for compiling and running the artifact.Further execution details are available in the repository documentation.

4) Data sets:

The artifact workflow prepares SwitchTransformer models, builds the system for the target GPU, installs dependencies, and runs the provided evaluation scripts.

  • SwitchTransformer model weights must be downloaded from HuggingFace before using the artifact.
  • The setup uses a Docker container with GPU access and a mounted model-preparation directory.
  • The build configuration should match the GPU’s compute capability, such as SM=80 for an A100.
  • After model preparation and dependency installation, evaluation is launched with the provided script.

F. Evaluation and expected results

The evaluation script records MoE block latency, inference throughput, and peak memory, with results presented in Figures 10–12.

  • The evaluation outputs CSV files for MoE block latency, inference throughput, and peak memory usage.
  • These measurements are presented in Figure 10, Figure 11, and Figure 12, respectively.
  • The experiment configuration can be customized in scripts/eval_all.py.
Loading 2308.12066v3…