🤖 machine learning

DisagMoE: Computation-Communication overlapped MoE Training via Disaggregated AF-Pipe Parallelism

DisagMoE is a novel Mixture-of-Experts training system that enhances efficiency on large-scale clusters by disaggregating attention and feed-forward layers into separate GPU groups with a multi-stage pipeline, effectively overlapping communication and computation to overcome all-to-all communication bottlenecks and achieve up to 1.8x training speedup.

Original authors: Zhichen Zeng, Chi-Chih Chang, Jiayi Wang, Zezhou Wang, Ningxin Zheng, Zheng Zhong, Cesar A. Stuardo, Dongyang Wang, Mohamed S. Abdelfattah, Haibin Lin, Banghua Zhu, Ang Li, Ziheng Jiang

Published 2026-05-13

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Zhichen Zeng, Chi-Chih Chang, Jiayi Wang, Zezhou Wang, Ningxin Zheng, Zheng Zhong, Cesar A. Stuardo, Dongyang Wang, Mohamed S. Abdelfattah, Haibin Lin, Banghua Zhu, Ang Li, Ziheng Jiang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are running a massive, high-speed factory that builds giant "thinking machines" (AI models). This factory has two main types of workers:

The "Context" Team (Attention): These workers are like detectives. They look at the whole story so far to understand what's happening. They are very good at thinking, but they don't need to move around much.
The "Specialist" Team (Experts/FFN): These are the actual builders. There are thousands of them, but for any single task, only a few are needed. They are great at doing the heavy lifting, but they are scattered all over the factory floor.

The Problem: The Traffic Jam

In the old way of running this factory (called Expert Parallelism), the Context Team and the Specialist Team were stuck in the same room.

Here's what happened:

The Context Team finished their thinking.
They had to shout across the room to the Specialists: "Hey, you, you, and you! Come here and work on this!"
The Specialists did their work.
They had to shout back: "Here is the result!"

The problem? The "shouting" (sending data between computers) was slow, especially when the factory got huge and the workers were on different floors (different server nodes). The specialists often sat idle, waiting for the Context Team to finish shouting, and vice versa. The factory was stuck in a traffic jam, spending more time talking than working.

Previous attempts to fix this were like trying to tell the workers to "talk while they work." But because the Context Team takes a long time to think (especially with long stories) and the Specialists are quick to build, the timing never quite matched up. There was always leftover shouting time that couldn't be hidden.

The Solution: DisagMoE (The Disaggregated Factory)

The authors of this paper, DisagMoE, decided to stop trying to make the two teams work in the same room. Instead, they built a two-stage assembly line.

1. Separate the Teams (Disaggregation)
They moved the Context Team to one set of machines and the Specialist Team to a completely different set of machines.

Team A (Context): All the "detective" layers are on one group of computers.
Team B (Specialists): All the "builder" layers are on another group.

2. The New Assembly Line (AF-Pipe)
Instead of a chaotic shout-out, they created a smooth, one-way conveyor belt system:

The Context Team finishes a batch of thinking and slides the work down the belt to the Specialists.
While the Specialists are working on Batch #1, the Context Team is already starting Batch #2.
While the Specialists finish Batch #1 and send the results back, the Context Team is already working on Batch #3.

This is called pipelining. It's like a relay race where the baton is passed so smoothly that no runner ever stops running.

3. The Smart Manager (Adaptive Allocation)
Here is the clever part. The paper realized that the Context Team and the Specialist Team have different needs.

The Context Team is like a heavy thinker; they need powerful brains (computing power) but don't need a fast internet connection as much.
The Specialist Team is like a runner; they need a super-fast highway (network bandwidth) to get their data to the right person quickly.

In the old factory, everyone got the same amount of internet and computing power. In DisagMoE, the manager is smart. If the story is very long, they give more "internet lanes" to the Specialists and more "brain power" to the Context Team. They constantly adjust the resources to make sure neither team is ever waiting for the other.

The Results

By separating the teams and letting the manager adjust the resources, the factory stopped wasting time waiting for messages.

Speed: They found that this new system made the training 1.8 times faster than the old standard methods.
Efficiency: They reduced the time the computers spent just "talking" (waiting for data) by up to 88%.

The Big Takeaway

The paper argues that you can't just throw more internet bandwidth at a problem to make it faster. You have to look at the specific job each part of the AI is doing. By splitting the "thinking" from the "building" and giving each team exactly the right amount of resources they need, you can build these giant AI models much, much faster.

In short: They stopped trying to force a square peg into a round hole by separating the workers, building a better conveyor belt, and hiring a manager who knows exactly how to balance the workload.

Technical Summary: DisagMoE

1. Problem Statement

Large-scale Mixture-of-Experts (MoE) architectures enable the training of trillion-parameter Large Language Models (LLMs) by sparsely activating a subset of experts. However, training these models faces severe bottlenecks due to all-to-all communication required for expert parallelism (EP). As models scale across multiple GPU nodes, the limited inter-node network bandwidth (e.g., InfiniBand) exacerbates communication latency, which can consume up to 50% of a training step's time.

Existing optimization strategies primarily focus on overlapping communication with computation at the operator level (e.g., chunking FFN computations) or using micro-batch pipelining. However, these approaches suffer from two fundamental limitations:

Residual Stalls: Inherent imbalances in computation-communication ratios between attention and FFN layers prevent perfect overlap. Attention computation scales quadratically with sequence length, while FFN and communication scale linearly. This creates "communication tails" that cannot be hidden, especially as sequence lengths increase.
Resource Inefficiency: Prior disaggregated approaches (separating attention and FFN onto different hardware) often require each group to hold full model parameters, leading to Out-of-Memory (OOM) errors at scale, or fail to adapt resource allocation to the specific compute-communication characteristics of each component.

2. Methodology: DisagMoE

DisagMoE is a training system designed to mitigate all-to-all bottlenecks by jointly optimizing model placement and scheduling through a disaggregated architecture.

2.1 Disaggregated Component Placement

Unlike traditional Expert Parallelism where attention and FFN layers coexist on the same GPU group, DisagMoE partitions the model by component type:

Attention Groups (A-Workers): These groups contain self-attention layers from multiple transformer layers. They replicate dense components and distribute training data via Data Parallelism (DP).
FFN Groups (F-Workers): These groups contain FFN layers (experts) from multiple transformer layers. Experts are sharded across GPUs within the group via Expert Parallelism (EP).
Layer Assignment: Layers are assigned to groups in an interleaved manner (e.g., Group 0 handles layers 0, $p$ , $2p$ ; Group 1 handles layers 1, $p+1$ , etc.), maximizing device utilization under memory constraints.

2.2 AF-Pipe: Multi-Stage Pipeline with Overlapping

DisagMoE introduces AF-Pipe, a pipeline schedule that treats communication as a first-class stage rather than a side effect.

Many-to-Many (M2N) Communication: Instead of standard all-to-all exchanges within a single node group, AF-Pipe establishes uni-directional, many-to-many token exchanges between the Attention and FFN groups.
Stage Alignment: The pipeline aligns boundaries across groups, fusing the traditional "dispatch" and "combine" operations into a single unified communication stage. This reduces communication overhead by roughly $1/k$ (where $k$ is the top- $k$ experts) and eliminates redundant data transfers.
Asynchronous Execution: The system utilizes three coordinated streams (forward, backward, and communication) within each worker. This allows the FFN group to send gradients (N2M) while simultaneously computing the forward pass on available data, ensuring non-blocking execution and near-continuous GPU utilization.

2.3 Adaptive Worker Allocation

To address the imbalance between attention (compute-bound) and FFN (communication-bound), DisagMoE employs an adaptive allocator guided by a Compute-Communication Roofline Model.

Modeling: The system models the turning point ( $\hat{I}$ ) where a component shifts from being network-bound to compute-bound. Attention layers reach the compute roof faster than FFN layers as sequence length increases.
Optimization: The system formulates GPU and Network Interface Card (NIC) partitioning as a Mixed Integer Linear Programming (MILP) problem. It dynamically allocates more network bandwidth to FFN groups (to lift their communication roof) and more compute resources to Attention groups, balancing the pipeline stages to minimize bubbles.
Search Strategy: An algorithm uses the roofline model to generate a seed configuration, which is then refined via local search using measured per-step times to account for unmodeled overheads.

3. Key Contributions

Disaggregated Component Placement: A novel architecture that separates attention and FFN layers into disjoint worker groups, allowing for independent scaling and resource allocation.
AF-Pipe Schedule: A multi-stage pipeline that replaces conventional all-to-all with inter-group M2N exchanges, systematically overlapping communication with both attention and FFN computations.
Adaptive Resource Allocation: A roofline-guided mechanism that jointly sizes GPU and NIC resources for attention and FFN groups, effectively balancing the compute-communication ratio across varying sequence lengths and model configurations.

4. Experimental Results

The authors evaluated DisagMoE on clusters of 16 nodes (128× H800 GPUs) using three MoE models: DeepSeek-MoE, GPT-OSS-120B, and Qwen3, with sequence lengths ranging from 4K to 32K.

Throughput Improvement: DisagMoE consistently outperformed strong baselines, including Megatron-LM (1F1B interleaved), Tutel, Comet, and DualPipe.
- Achieved up to 1.81× speedup over Megatron-1F1B.
- Achieved up to 1.34× speedup over state-of-the-art MoE overlap systems.
Communication Reduction: Compared to baselines, DisagMoE reduced non-overlapped communication time by up to 88% (vs. Tutel), 75% (vs. Comet), and 45% (vs. DualPipe).
Scalability: The system demonstrated effectiveness across different top- $k$ values and Expert Parallelism (EP) sizes.
Resource Sensitivity: Experiments confirmed that optimal GPU ratios between attention and FFN groups vary with sequence length (e.g., longer sequences favor more attention resources), validating the need for adaptive allocation.

5. Significance and Claims

The paper claims that DisagMoE addresses a fundamental limitation in current MoE training: the inability to perfectly overlap communication due to the structural imbalance between attention and FFN layers. By disaggregating these components and adaptively allocating resources based on a roofline model, the system expands the effective overlap window beyond what symmetric bandwidth provisioning can achieve.

The authors emphasize that their approach is particularly significant for large-scale pre-training workloads with fixed sequence lengths. They note that while their current design assumes stable shapes (fixed sequence length and micro-batch), the core insight—that future hardware and software should target favorable compute-communication ratios across components rather than simply scaling raw bandwidth—offers a new direction for scaling MoE training efficiently. The system is implemented on Megatron-LM and is shown to be a viable solution for training models with hundreds of billions of parameters without the memory constraints that plagued previous disaggregated training attempts.