MoonshotAI/MoonEP

MoonEP: A Perfectly Balanced Expert Parallelism Library via Dynamic Redundant Experts

What it solves

MoonEP addresses the efficiency bottlenecks in Expert Parallelism (EP) for Mixture-of-Experts (MoE) models, specifically the problem of load imbalance. In standard EP, some GPU ranks may receive significantly more tokens than others (the "hottest rank" problem), leading to increased latency, memory fragmentation, and potential Out-of-Memory (OOM) errors during training.

How it works

MoonEP ensures that every GPU rank processes an identical number of tokens regardless of how the router distributes them. It achieves this through:

  • Dynamic Redundant Experts: The library plans a small number of redundant experts online and prefetches them to balance the load perfectly across ranks.
  • Zero-Copy Communication: It uses fused permute/unpermute operations to send tokens directly to their final expert-grouped positions on remote ranks, returning buffer views to the computation to eliminate unnecessary data copies.
  • Online Planning: A high-performance GPU kernel determines the optimal distribution of experts and tokens with negligible overhead.
  • Static Memory Shapes: By maintaining a fixed buffer size for tokens, it eliminates the need for per-layer host synchronization and prevents memory fragmentation.

Who it’s for

This library is designed for developers and researchers building large-scale MoE models who need to optimize communication overhead and ensure stable, balanced training and inference across multiple NVIDIA GPUs (and eventually Zhenwu PPUs).

Highlights

  • Perfect Load Balance: Every rank receives exactly the same number of tokens, making iteration time immune to routing imbalance.
  • Zero-Copy Dispatch: Eliminates the communication-buffer to user-buffer copy, speeding up raw communication.
  • OOM Prevention: Static memory shapes prevent the memory fragmentation that typically occurs with varying activation shapes in imbalanced MoE training.
  • Integrated Weight Prefetching: Includes a dedicated mechanism to prefetch redundant expert weights into local slots for efficient computation.

Related

  • Project
    deepseek-ai/DeepEP

    A high-performance communication library for MoE models that provides optimized GPU kernels for expert parallelism to maximize throughput and minimize SM resource usage.

  • Project
    microsoft/Tutel

    An optimized Mixture-of-Experts (MoE) implementation that accelerates training and inference for massive LLMs on NVIDIA and AMD GPUs using advanced parallelism and quantization.

  • Dispatch
    Mixture of Experts (MoE) Explained

    Mixture of Experts (MoE) allows transformer models to scale parameters while maintaining efficient pretraining and faster inference by activating only a subset of neural network experts per token.

  • Dispatch
    vLLM Elastic Expert Parallelism

    vLLM introduces Elastic Expert Parallelism (Elastic EP), enabling Mixture-of-Experts (MoE) deployments to scale the number of workers up or down at runtime without requiring a server restart.

  • Project
    EfficientMoE/MoE-Infinity

    MoE-Infinity is a library for cost-effective MoE inference that enables running large Mixture-of-Experts models on memory-constrained GPUs by offloading expert weights to host memory and SSD.