Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Revisiting SMoE Language Models by Evaluating Inefficiencies with Task Specific Expert Pruning

Soumajyoti Sarkar, Leonard Lausen, Volkan Cevher, Sheng Zha, Thomas Brox, George Karypis

TL;DR

An adaptive task-aware pruning technique UNCURL is introduced to reduce the number of experts per MoE layer in an offline manner post-training, revealing a threshold pruning factor for the reduction that depends on the number of experts used in pretraining, above which, the reduction starts to degrade model performance.

Abstract

Sparse Mixture of Expert (SMoE) models have emerged as a scalable alternative to dense models in language modeling. These models use conditionally activated feedforward subnetworks in transformer blocks, allowing for a separation between total model parameters and per-example computation. However, large token-routed SMoE models face a significant challenge: during inference, the entire model must be used for a sequence or a batch, resulting in high latencies in a distributed setting that offsets the advantages of per-token sparse activation. Our research explores task-specific model pruning to inform decisions about designing SMoE architectures, mainly modulating the choice of expert counts in pretraining. We investigate whether such pruned models offer advantages over smaller SMoE models trained from scratch, when evaluating and comparing them individually on tasks. To that end, we introduce an adaptive task-aware pruning technique UNCURL to reduce the number of experts per MoE layer in an offline manner post-training. Our findings reveal a threshold pruning factor for the reduction that depends on the number of experts used in pretraining, above which, the reduction starts to degrade model performance. These insights contribute to our understanding of model design choices when pretraining with SMoE architectures, particularly useful when considering task-specific inference optimization for later stages.

Revisiting SMoE Language Models by Evaluating Inefficiencies with Task Specific Expert Pruning

TL;DR

An adaptive task-aware pruning technique UNCURL is introduced to reduce the number of experts per MoE layer in an offline manner post-training, revealing a threshold pruning factor for the reduction that depends on the number of experts used in pretraining, above which, the reduction starts to degrade model performance.

Abstract

Sparse Mixture of Expert (SMoE) models have emerged as a scalable alternative to dense models in language modeling. These models use conditionally activated feedforward subnetworks in transformer blocks, allowing for a separation between total model parameters and per-example computation. However, large token-routed SMoE models face a significant challenge: during inference, the entire model must be used for a sequence or a batch, resulting in high latencies in a distributed setting that offsets the advantages of per-token sparse activation. Our research explores task-specific model pruning to inform decisions about designing SMoE architectures, mainly modulating the choice of expert counts in pretraining. We investigate whether such pruned models offer advantages over smaller SMoE models trained from scratch, when evaluating and comparing them individually on tasks. To that end, we introduce an adaptive task-aware pruning technique UNCURL to reduce the number of experts per MoE layer in an offline manner post-training. Our findings reveal a threshold pruning factor for the reduction that depends on the number of experts used in pretraining, above which, the reduction starts to degrade model performance. These insights contribute to our understanding of model design choices when pretraining with SMoE architectures, particularly useful when considering task-specific inference optimization for later stages.
Paper Structure (18 sections, 9 figures, 8 tables, 1 algorithm)

This paper contains 18 sections, 9 figures, 8 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Technical Preliminaries on SMoE
  3. Tradeoffs in Performance, Training and Inference Costs
  4. Can expert pruning retain the performance benefits?
  5. Experiment Setup
  6. Observations
  7. Pruning with Clustering and Merging
  8. UnCuRL: Unified Clustering Based on Router Logits (Laywerise)
  9. Time Complexity
  10. Results
  11. Performance of SMoE models on downstream tasks
  12. Impact of Cluster-Merging experts
  13. Cluster Visualization
  14. Comparing Baselines
  15. Frequency Based Merging
  16. ...and 3 more sections

Figures (9)

  • Figure 1: MoE models over a distributed setup with Expert Parallelism.
  • Figure 2: Loss vs Inference costs
  • Figure 3: (Performance Equivalent) Dense Model Size when increasing # experts
  • Figure 5: Pruning and Evaluation
  • Figure 6: Wall clock latencies comparing the time spent in All2All and the expert computation in one SMoE layer of a transformer block as we increase the number of experts. Expert computation time denotes the time spent in the expert FFN operations of one SMoE layer. Here, we consider a GPT2 SMoE model with backbone size of 124M parameters. Latencies are for computing the output logits of a single sequence of 512 tokens, batch size of 1 per GPU (we do not do any token generation here). Setup used with Deepspeed Zero 2 data parallelism with rank 8 and expert parallelism with rank 2.
  • ...and 4 more figures