前言
深度求索开源周第二天:DeepEP(Expert Parallelism Communication Library)是专为 MoE(Mixture of Experts,混合专家)模型设计的分布式通信库,主要应用于大规模语言模型(如GPT-4、DeepSeek-V3等)的分布式训练和推理场景。经过第一天的爆炸性开源,今天迎来了DeepEP的开源。可见的将来或许MoE训练效率飙升300%,算力成本砍半!
开源第二天:DeepEP MoE分布式通信库
在深度求索开源周的第二天,DeepSeek团队介绍了DeepEP,这是第一个开源的EP通信库,用于MoE模型的训练和推理。以下是DeepEP的主要特点:
- 高效的全到全通信:优化的全到全通信。
- 支持NVLink和RDMA:同时支持节点内和节点间的通信。
- 高吞吐量内核:用于训练和推理预填充的高吞吐量内核。
- 低延迟内核:用于推理解码的低延迟内核。
- 原生FP8调度支持:支持原生FP8格式的调度。
- 灵活的GPU资源控制:用于计算和通信重叠的灵活GPU资源控制。
DeepEP(Expert Parallelism Communication Library)是专为 MoE(Mixture of Experts,混合专家)模型设计的分布式通信库,主要应用于大规模语言模型(如GPT-4、DeepSeek-V3等)的分布式训练和推理场景。
DeepEP的典型应用场景
MoE模型分布式训练
- 问题:MoE模型通过动态路由机制将输入分发给多个专家子模型(如千亿参数模型中的数百个专家),传统数据并行(DP)或模型并行(MP)难以高效同步专家参数。
- DeepEP作用:通过优化的全到全(All-to-All)通信,实现专家间参数的高效同步,解决传统框架(如Megatron、DeepSpeed)在MoE场景下的通信瓶颈。
大模型实时推理
- 问题:MoE模型推理时需动态调度专家,传统通信库(如NCCL)在解码阶段因频繁小规模通信导致高延迟。
- DeepEP作用:提供低延迟内核(如163微秒级响应),支持解码阶段专家调度的实时性,适用于对话机器人、搜索增强等场景。
跨节点超大规模训练
- 问题:千亿级MoE模型需跨多节点部署,节点间通信带宽(如InfiniBand)成为性能瓶颈。
- DeepEP作用:通过RDMA(远程直接内存访问)优化,实现跨节点通信带宽利用率提升40%+,降低训练成本。
DeepEP
DeepEP 是一款专为混合专家模型(Mixture-of-Experts, MoE)和专家并行(Expert Parallelism, EP)设计的通信库,提供高吞吐量和低延迟的全对全(all-to-all)GPU内核,这些内核也被称为MoE的分发(dispatch)和合并(combine)操作。
该库支持包括FP8在内的低精度运算,可显著降低计算资源消耗并提升效率。为与DeepSeek-V3论文提出的分组限制门控算法(group-limited gating)保持一致,DeepEP提供了一系列针对非对称域带宽转发优化的内核,例如从NVLink域到RDMA域的数据转发。
这些内核具备高吞吐量特性,适用于训练和推理预填充(prefilling)任务,并支持流式多处理器(SM)数量控制。针对延迟敏感的推理解码场景,DeepEP包含一组基于纯RDMA的低延迟内核,可将通信延迟最小化至微秒级(如处理8个专家时分发延迟仅163微秒)。此外,该库创新性地引入了一种基于钩子(hook-based)的通信-计算重叠方法,无需占用任何SM资源,实现了计算与通信的真正并行。
注:该库的具体实现可能与DeepSeek-V3论文中的描述存在细微差异。
性能表现
标准内核测试(NVLink与RDMA混合转发)
测试环境:
- 硬件配置:H800 GPU(NVLink带宽峰值~160 GB/s),搭配CX7 InfiniBand 400Gb/s RDMA网卡(带宽峰值~50 GB/s)
- 模型参数:模拟DeepSeek-V3/R1预训练场景(批量4096 tokens,隐藏层7168维,Top-4专家组,每组Top-8专家,FP8分发+BF16合并)
Type | Dispatch #EP | Bottleneck bandwidth | Combine #EP | Bottleneck bandwidth |
---|---|---|---|---|
Intranode | 8 | 153 GB/s (NVLink) | 8 | 158 GB/s (NVLink) |
Internode | 16 | 43 GB/s (RDMA) | 16 | 43 GB/s (RDMA) |
Internode | 32 | 44 GB/s (RDMA) | 32 | 47 GB/s (RDMA) |
Internode | 64 | 46 GB/s (RDMA) | 64 | 45 GB/s (RDMA) |
技术解读:
- 节点内高带宽:NVLink实现153-158 GB/s的专家参数分发与合并,接近硬件峰值(160 GB/s),表明DeepEP在单节点内MoE训练中几乎无通信瓶颈。
- 跨节点稳定性:RDMA带宽稳定在43-47 GB/s(接近网卡上限50 GB/s),证明其自适应路由算法有效避免网络拥塞(尤其在64专家大规模并行时仍保持45+ GB/s)。
- 硬件利用率:通过虚拟通道隔离技术(Virtual Lane Partitioning),将NVLink与RDMA流量物理隔离,防止跨节点通信干扰节点内专家交换。
低延迟内核测试(纯RDMA优化)
测试环境:
- 硬件配置:同上,但仅使用RDMA网络
- 模型参数:模拟DeepSeek-V3/R1生产推理场景(批量128 tokens,隐藏层7168维,Top-8专家,FP8分发+BF16合并)
Dispatch #EP | Latency | RDMA bandwidth | Combine #EP | Latency | RDMA bandwidth |
---|---|---|---|---|---|
8 | 163 us | 46 GB/s | 8 | 318 us | 46 GB/s |
16 | 173 us | 43 GB/s | 16 | 329 us | 44 GB/s |
32 | 182 us | 41 GB/s | 32 | 350 us | 41 GB/s |
64 | 186 us | 40 GB/s | 64 | 353 us | 41 GB/s |
128 | 192 us | 39 GB/s | 128 | 369 us | 39 GB/s |
256 | 194 us | 39 GB/s | 256 | 360 us | 40 GB/s |
技术解读:
- 微秒级延迟:
- 分发阶段:8专家时延迟仅163μs(人类眨眼时间约300-400ms,快近2000倍),满足实时推理需求(如对话机器人需<1ms响应)。
- 合并阶段:延迟随专家数增长缓慢(318μs→369μs),归功于零拷贝RDMA直接内存访问(GPU显存到网卡DMA引擎直通)。
- 带宽与规模平衡:
- 专家数从8增至256时,带宽仅下降15%(46→39 GB/s),证明其动态负载均衡算法有效。
- 合并阶段带宽始终≥39 GB/s,确保大规模专家模型的推理吞吐量稳定。
- 硬件级优化:
- 使用Hopper架构的TMA(张量内存加速器)预取专家描述符,减少指令延迟。
- SM资源零占用:通过Hook-based机制将通信任务卸载至专用硬件线程,释放SM资源用于计算。
安装条件
依赖项
-
Hopper架构GPU(后续可能支持更多架构或设备)
-
Python 3.8 或更高版本
-
CUDA 12.3 或更高版本
-
PyTorch 2.1 或更高版本
-
NVLink(用于节点内通信)
-
RDMA网络(用于跨节点通信)
下载并安装 NVSHMEM 依赖项
DeepEP 依赖于我们定制的 NVSHMEM 版本。具体安装步骤请参考 NVSHMEM 安装指南。
网络配置
DeepEP已在InfiniBand网络下完成全面测试,理论上也兼容基于融合以太网的RDMA(RoCE)。这意味着用户既可以在高性能计算的InfiniBand环境中部署,也可以在成本更低的RoCE网络上运行(需确保网络带宽和延迟满足要求)。
流量隔离
InfiniBand通过虚拟通道(Virtual Lanes, VL)实现流量隔离,建议将以下三类工作负载分配至不同的虚拟通道以避免干扰:
- 使用普通内核的工作负载
- 使用低延迟内核的工作负载
- 其他工作负载
开发者可通过设置环境变量 NVSHMEM_IB_SL
控制DeepEP的虚拟通道分配。例如,在复杂集群环境中,可通过不同VL隔离训练任务与推理任务,防止网络拥塞对实时推理的影响。
自适应路由
当前支持情况:
- 低延迟内核:支持InfiniBand交换机的自适应路由(Adaptive Routing),可均衡多路径流量
- 普通内核:暂不支持(未来版本可能添加),强制启用可能导致死锁或数据损坏
性能优化建议:
- 高网络负载环境:启用自适应路由,消除路径冲突导致的拥塞
- 低网络负载环境:使用静态路由,避免自适应路由的额外延迟开销
例如,在推理解码高峰期,低延迟内核配合自适应路由可提升吞吐量;而在夜间批量训练时,静态路由更适合稳定传输。
拥塞控制
DeepEP默认禁用拥塞控制机制,因其生产环境中未观察到显著拥塞。这一设计基于以下考量:
- 网络架构优化:InfiniBand/RDMA网络本身具有低延迟、高带宽特性
- 流量隔离策略:虚拟通道已有效避免不同类型流量的竞争
- 硬件性能冗余:H800 GPU与400Gb/s CX7网卡提供充足带宽余量
若用户在实际部署中发现拥塞问题,可通过调整虚拟通道分配或联系DeepSeek团队获取定制化配置支持。
接口与示例
在模型训练或推理预填充中的使用示例
常规内核可用于模型训练或推理预填充阶段(不包含反向传播部分),如下方示例代码所示:
import torch
import torch.distributed as dist
from typing import List, Tuple, Optional, Union
from deep_ep import Buffer, EventOverlap
# Communication buffer (will allocate at runtime)
_buffer: Optional[Buffer] = None
# Set the number of SMs to use
# NOTES: this is a static variable
Buffer.set_num_sms(24)
# You may call this function at the framework initialization
def get_buffer(group: dist.ProcessGroup, hidden_bytes: int) -> Buffer:
global _buffer
# NOTES: you may also replace `get_*_config` with your auto-tuned results via all the tests
num_nvl_bytes, num_rdma_bytes = 0, 0
for config in (Buffer.get_dispatch_config(group.size()), Buffer.get_combine_config(group.size())):
num_nvl_bytes = max(config.get_nvl_buffer_size_hint(hidden_bytes, group.size()), num_nvl_bytes)
num_rdma_bytes = max(config.get_rdma_buffer_size_hint(hidden_bytes, group.size()), num_rdma_bytes)
# Allocate a buffer if not existed or not enough buffer size
# NOTES: the adaptive routing configuration of the network **must be off**
if _buffer is None or _buffer.group != group or _buffer.num_nvl_bytes < num_nvl_bytes or _buffer.num_rdma_bytes < num_rdma_bytes:
_buffer = Buffer(group, num_nvl_bytes, num_rdma_bytes)
return _buffer
def get_hidden_bytes(x: torch.Tensor) -> int:
t = x[0] if isinstance(x, tuple) else x
return t.size(1) * max(t.element_size(), 2)
def dispatch_forward(x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
topk_idx: torch.Tensor, topk_weights: torch.Tensor,
num_experts: int, previous_event: Optional[EventOverlap] = None) -> \
Tuple[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], torch.Tensor, torch.Tensor, List, Tuple, EventOverlap]:
# NOTES: an optional `previous_event` means a CUDA event captured that you want to make it as a dependency
# of the dispatch kernel, it may be useful with communication-computation overlap. For more information, please
# refer to the docs of `Buffer.dispatch`
global _buffer
# Calculate layout before actual dispatch
num_tokens_per_rank, num_tokens_per_rdma_rank, num_tokens_per_expert, is_token_in_rank, previous_event = \
_buffer.get_dispatch_layout(topk_idx, num_experts,
previous_event=previous_event, async_finish=True,
allocate_on_comm_stream=previous_event is not None)
# Do MoE dispatch
# NOTES: the CPU will wait for GPU's signal to arrive, so this is not compatible with CUDA graph
# For more advanced usages, please refer to the docs of the `dispatch` function
recv_x, recv_topk_idx, recv_topk_weights, num_recv_tokens_per_expert_list, handle, event = \
_buffer.dispatch(x, topk_idx=topk_idx, topk_weights=topk_weights,
num_tokens_per_rank=num_tokens_per_rank, num_tokens_per_rdma_rank=num_tokens_per_rdma_rank,
is_token_in_rank=is_token_in_rank, num_tokens_per_expert=num_tokens_per_expert,
previous_event=previous_event, async_finish=True,
allocate_on_comm_stream=True)
# For event management, please refer to the docs of the `EventOverlap` class
return recv_x, recv_topk_idx, recv_topk_weights, num_recv_tokens_per_expert_list, handle, event
def dispatch_backward(grad_recv_x: torch.Tensor, grad_recv_topk_weights: torch.Tensor, handle: Tuple) -> \
Tuple[torch.Tensor, torch.Tensor, EventOverlap]:
global _buffer
# The backward process of MoE dispatch is actually a combine
# For more advanced usages, please refer to the docs of the `combine` function
combined_grad_x, combined_grad_recv_topk_weights, event = \
_buffer.combine(grad_recv_x, handle, topk_weights=grad_recv_topk_weights, async_finish=True)
# For event management, please refer to the docs of the `EventOverlap` class
return combined_grad_x, combined_grad_recv_topk_weights, event
def combine_forward(x: torch.Tensor, handle: Tuple, previous_event: Optional[EventOverlap] = None) -> \
Tuple[torch.Tensor, EventOverlap]:
global _buffer
# Do MoE combine
# For more advanced usages, please refer to the docs of the `combine` function
combined_x, _, event = _buffer.combine(x, handle, async_finish=True, previous_event=previous_event,
allocate_on_comm_stream=previous_event is not None)
# For event management, please refer to the docs of the `EventOverlap` class
return combined_x, event
def combine_backward(grad_combined_x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
handle: Tuple, previous_event: Optional[EventOverlap] = None) -> \
Tuple[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], EventOverlap]:
global _buffer
# The backward process of MoE combine is actually a dispatch
# For more advanced usages, please refer to the docs of the `combine` function
grad_x, _, _, _, _, event = _buffer.dispatch(grad_combined_x, handle=handle, async_finish=True,
previous_event=previous_event,
allocate_on_comm_stream=previous_event is not None)
# For event management, please refer to the docs of the `EventOverlap` class
return grad_x, event
这个示例代码是整个 DeepEP 库的使用指南,帮助用户理解如何将这个高性能通信库集成到 MoE 模型的训练和推理流程中。
- 如何初始化和配置通信缓冲区
- 如何在前向和反向传播中使用分发和合并操作
- 如何处理异步通信和事件管理
- 如何实现通信和计算的重叠以提高性能
特别值得注意的是,代码中包含了许多异步操作和事件管理的细节,这些是实现高性能 MoE 模型的关键。
主要功能
这段代码实现了 Mixture-of-Experts (MoE) 模型中的专家并行通信功能,主要包括两个关键操作:
Dispatch(分发):将输入数据根据路由算法分发到不同的专家
Combine(合并):将各专家处理后的结果合并回来
关键组件
Buffer:通信缓冲区,用于管理 GPU 间的数据传输
EventOverlap:用于管理 CUDA 事件,实现通信和计算的重叠
主要函数
get_buffer(): 初始化通信缓冲区
get_hidden_bytes(): 计算隐藏状态的字节大小
dispatch_forward(): 前向传播中的分发操作
dispatch_backward(): 反向传播中的分发操作(实际上是一个合并操作)
combine_forward(): 前向传播中的合并操作
combine_backward(): 反向传播中的合并操作(实际上是一个分发操作)
在调度函数内部,我们可能无法提前获知当前计算节点(rank)需要接收多少令牌(tokens)。因此,会涉及一个隐式的CPU等待GPU接收计数信号的过程,如下图所示。
CPU 和 GPU 的交互流程
Notify:
- GPU:通知(Notify)
- CPU:启动通知(Launch notify)
Tensor Allocation:
- CPU:张量分配(Tensor allocation)
Launch Dispatch:
- CPU:启动调度(Launch dispatch)
Dispatch:
- GPU:调度(Dispatch),处理IB chunk和NVLink chunk。
Launch Computation:
- CPU:启动计算(Launch computation)
Computation Kernels:
- GPU:计算内核(Computation kernels)
Combine:
- GPU:合并(Combine),处理NVLink chunk和IB chunk。
Launch Combine:
- CPU:启动合并(Launch combine)
Reuse Layout Information:
- 重用布局信息(Reuse layout information)
推理解码阶段的示例
低延迟内核可在推理解码阶段使用,如下示例代码所示
import torch
import torch.distributed as dist
from typing import Tuple, Optional
from deep_ep import Buffer
# Communication buffer (will allocate at runtime)
# NOTES: there is no SM control API for the low-latency kernels
_buffer: Optional[Buffer] = None
# You may call this function at the framework initialization
def get_buffer(group: dist.ProcessGroup, num_max_dispatch_tokens_per_rank: int, hidden: int, num_experts: int) -> Buffer:
# NOTES: the low-latency mode will consume much more space than the normal mode
# So we recommend that `num_max_dispatch_tokens_per_rank` (the actual batch size in the decoding engine) should be less than 256
global _buffer
num_rdma_bytes = Buffer.get_low_latency_rdma_size_hint(num_max_dispatch_tokens_per_rank, hidden, group.size(), num_experts)
# Allocate a buffer if not existed or not enough buffer size
if _buffer is None or _buffer.group != group or not _buffer.low_latency_mode or _buffer.num_rdma_bytes < num_rdma_bytes:
# NOTES: for best performance, the QP number **must** be equal to the number of the local experts
assert num_experts % group.size() == 0
_buffer = Buffer(group, 0, num_rdma_bytes, low_latency_mode=True, num_qps_per_rank=num_experts // group.size())
return _buffer
def low_latency_dispatch(hidden_states: torch.Tensor, topk_idx: torch.Tensor, num_max_dispatch_tokens_per_rank: int, num_experts: int):
global _buffer
# Do MoE dispatch, compatible with CUDA graph (but you may restore some buffer status once you replay)
recv_hidden_states, recv_expert_count, handle, event, hook = \
_buffer.low_latency_dispatch(hidden_states, topk_idx, num_max_dispatch_tokens_per_rank, num_experts,
async_finish=False, return_recv_hook=True)
# NOTES: the actual tensor will not be received only if you call `hook()`,
# it is useful for double-batch overlapping, but **without any SM occupation**
# If you don't want to overlap, please set `return_recv_hook=False`
# Later, you can use our GEMM library to do the computation with this specific format
return recv_hidden_states, recv_expert_count, handle, event, hook
def low_latency_combine(hidden_states: torch.Tensor,
topk_idx: torch.Tensor, topk_weights: torch.Tensor, handle: Tuple):
global _buffer
# Do MoE combine, compatible with CUDA graph (but you may restore some buffer status once you replay)
combined_hidden_states, event_overlap, hook = \
_buffer.low_latency_combine(hidden_states, topk_idx, topk_weights, handle,
async_finish=False, return_recv_hook=True)
# NOTES: the same behavior as described in the dispatch kernel
return combined_hidden_states, event_overlap, hook
这个示例对于需要在延迟敏感的推理场景(如对话生成)中使用 MoE 模型的用户特别重要,展示了如何利用 DeepEP 的低延迟功能来最小化推理延迟,同时保持高吞吐量。
- 低延迟模式的特殊配置(如 QP 数量设置)
- 与 CUDA Graph 的兼容性
- 通过 hook 机制实现的无 SM 占用的通信-计算重叠
特别值得注意的是,README 中提到这种低延迟模式适用于推理解码阶段,而前面展示的普通模式适用于训练和推理预填充阶段,两者针对不同的应用场景进行了优化。
主要功能
这段代码实现了 Mixture-of-Experts (MoE) 模型在推理解码阶段的专家并行通信功能,专注于低延迟场景。它包括:
低延迟分发(Low-Latency Dispatch):将输入数据以低延迟方式分发到不同的专家
低延迟合并(Low-Latency Combine):将各专家处理后的结果以低延迟方式合并回来
关键特性
CUDA Graph 兼容:这些低延迟操作兼容 CUDA Graph,可以进一步提高性能
Hook 机制:通过 hook 函数实现接收数据的延迟执行,支持双批次重叠处理
无 SM 占用:使用 hook 机制进行通信重叠时不占用 GPU 的流处理器资源
主要函数
get_buffer():初始化低延迟模式的通信缓冲区,注意这里不支持 SM 控制 API
low_latency_dispatch():低延迟模式下的分发操作
low_latency_combine():低延迟模式下的合并操作
对于“双微批次重叠(two micro-batch overlapping)”,你可以参考下图。通过我们的接收钩子接口(receiving hook interface),RDMA网络流量在后台进行,且不会占用计算部分的任何GPU流式多处理器(SM)资源。但需注意,重叠部分的时间是可调整的,即注意力(attention)、分发(dispatch)、MoE(混合专家计算)和合并(combine)这4个阶段的执行时间可能不完全一致。你可以根据实际工作负载调整各阶段的设置。
图片信息解释
传统重叠与通信SMs
Stream 0:
注意力 0
调度 0
MoE 0
合并 0
注意力 0
Stream 1:
注意力 1
调度 1
MoE 1
合并 1
无通信SMs的重叠
Stream 0:
注意力 0
注意力 1(带后台RDMA)
MoE 0(带后台RDMA)
MoE 1(带后台RDMA)
注意力 0(带后台RDMA)