MoonEP logo

MoonEP

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

Developer ToolsPythonMIT 888 91Health 77
GitHub

What is it?

What it is

MoonEP is an Expert Parallelism communication library that maintains perfectly balanced token loads across ranks by dynamically creating redundant experts and prefetching their weights, enabling zero-copy, static-shape communication.

Why it exists

To eliminate communication bottlenecks and memory fragmentation caused by imbalanced expert routing in Mixture-of-Experts models, ensuring flat iteration time and avoiding OOM under skew.

Who should use it

Developers and researchers building large transformer models that require expert parallelism with balanced token distribution across GPUs, especially those with multi‑GPU clusters connected via NVLink and who need zero‑copy communication and static memory shapes for efficient training and inference.

Who should avoid it

Users running on single GPUs, machines without NVLink or other high‑bandwidth interconnects, or those who are not comfortable with low‑level CUDA programming and distributed training concepts. Also, teams that do not need expert parallelism or are on unsupported hardware (currently only NVIDIA GPUs, Zhenwu PPU is not yet released).

How it works

A quick walkthrough in plain English

How MoonEP works

Step 1 of 3

You interact with it

Open MoonEP, send a request, or connect it to your stack.

Features

Dynamic redundant experts for perfect token load balancing
Online near-optimal GPU planning kernel
Zero-copy communication with fused permute/unpermute
Static memory shapes to eliminate per-layer host synchronization
Weight prefetching via symmetric memory mapping
Asynchronous communication via CUDA events
Support for NVIDIA GPUs and Zhenwu PPU

Advantages

  • Immunity to routing imbalance (constant iteration time)
  • Elimination of GPU memory fragmentation
  • Prevention of Out-of-Memory (OOM) errors caused by activation shape variance
  • Reduced communication latency compared to DeepEP v2
  • Higher efficiency through direct buffer views for computation

Disadvantages

  • Requires contiguous symmetric-memory weight tensors
  • Requires manual management of buffer destruction
  • Zero-copy mode requires careful autograd handling to avoid aliasing issues

Installation

native

pip install -e .

FAQ

How does MoonEP achieve perfect load balancing in Expert Parallelism?

MoonEP uses dynamic redundant experts. It plans a near-optimal GPU kernel online to identify skewed routing and prefetches redundant experts to ensure every rank receives exactly S × K tokens, regardless of routing imbalance.

What are the benefits of MoonEP's zero-copy mechanism?

By using zero-copy, tokens are written directly to their expert-grouped positions on remote ranks. This eliminates the need for permute operations and avoids the expensive copy from communication buffers to user buffers, significantly reducing latency.

How should I configure the weight buffer for training versus inference?

For training, you must set B = E/R to ensure all experts touched by the group GEMM are local. For inference, B < E/R is allowed, and a value of 3–4 is recommended for optimal prefetching.

Does MoonEP help prevent Out-of-Memory (OOM) errors during training?

Yes. Because MoonEP uses static memory shapes and ensures every rank computes exactly the same number of tokens per layer, it prevents GPU memory fragmentation and keeps iteration time flat even under high routing imbalance.

What is the requirement for the weight tensor layout in MoonEP?

MoonEP requires one contiguous symmetric-memory weight tensor per expert projection (e.g., gate, up, down). The tensor must be laid out as [E+B, H, H'] so that the group GEMM can address experts via row indices.

Loading documentation…
View on GitHub

Featured in Videos

YouTube tutorials and walkthroughs for MoonEP

Alternatives

Similar projects ranked by category, topics, and text overlap.

Compare
MoonEP | MushyBook