Mixture of Experts (MoEs) in Transformers

Hugging Face has overhauled the transformers library to provide first-class support for Mixture of Experts (MoE) architectures. This redesign addresses the fundamental mismatch between how MoE models are stored in checkpoints and how they are executed on hardware, enabling significantly faster weight loading, more efficient inference via pluggable backends, and scalable distribution through expert parallelism.

MoE Fundamentals: Capacity vs. Active Parameters

Mixture of Experts models replace dense feed-forward layers with a set of learnable sub-networks called experts. A router selects a small subset of these experts to process each token, decoupling the model's total capacity from its inference cost.

  • Total Parameters: Determines the model's overall capacity and quality.
  • Active Parameters: Determines the inference speed and compute cost per token.

For example, the gpt-oss-20b model contains 21B total parameters but only activates ~3.6B parameters per token (using 4 out of 32 experts). This allows the model to maintain the quality of a 21B parameter system while operating at the speed of a 3.6B parameter model, achieving approximately 115 tokens per second on an M3 Ultra Mac.

Weight Loading Refactor and WeightConverter

Traditional weight loading in transformers assumed a one-to-one mapping between checkpoint tensors and runtime parameters. MoE checkpoints typically serialize experts independently (e.g., 256 separate tensors), but optimized runtime kernels require experts to be packed into a single contiguous tensor for grouped GEMM operations.

To solve this, Hugging Face introduced a conversion pipeline via the WeightConverter abstraction. This shifts the loading process from a simple key-by-key copy to a dynamic transformation:

  • MergeModulelist: Stacks multiple expert tensors into a single contiguous tensor.
  • SplitModulelist: Splits a packed tensor back into individual experts.
  • Lazy Materialization: The loader scans keys once and materializes tensors via a thread pool only when dependencies are ready, reducing memory peaks and avoiding repeated scans.

Weight Loading Benchmarks

Benchmarks using Qwen/Qwen1.5-110B-Chat on a single A100 (80GB) demonstrate significant speedups in the v5 pipeline compared to v4:

| Version | Strategy | Loading Mode | Time | | :--- | :--- | :--- | | v4.57.6 | device_map="auto" | Threadpool | 66.24s | | v5 | device_map="auto" | Async (default) | 20.71s | | v5 | TP | Async | 10.1s |

Pluggable Expert Backend

To decouple expert computation from model implementation, Hugging Face introduced an Experts Backend system using the @use_experts_implementation decorator. This allows models to switch between three execution strategies at runtime:

  1. eager: Loops over selected experts; used primarily for debugging and correctness.
  2. batched_mm: Uses torch.bmm to duplicate expert weights per token; optimized for small batches and GPU-heavy workloads.
  3. grouped_mm: Uses torch._grouped_mm to sort tokens by expert ID and perform a single grouped GEMM; optimized for large batches and memory-constrained environments.

Expert Parallelism (EP)

Expert Parallelism allows models with hundreds of billions of parameters to scale across multiple GPUs by distributing experts across devices. Unlike standard tensor parallelism, each device loads only its assigned subset of experts (num_experts / num_devices).

This is implemented via DistributedConfig(enable_expert_parallel=True) and relies on two core components:

  • GroupedGemmParallel: Handles the sharding of expert weights along the expert dimension (dim=0).
  • RouterParallel: Remaps global expert indices to local indices and uses an all-reduce operation to combine partial outputs across devices.

Optimized MoE Training with Unsloth

Through collaboration with Unsloth, Hugging Face has enabled faster MoE training by leveraging the Expert Backend abstraction and PyTorch's torch._grouped_mm API, combined with custom Triton grouped-GEMM and LoRA kernels. These optimizations provide:

  • Up to 12x faster MoE training.
  • Over 35% reduction in VRAM usage.
  • Approximately 6x longer context windows.
  • An overall speedup of 12-30x compared to the v4 transformers implementation.

Sources