DeepSeek开源周第二日-DeepEP

🚀deepseek开源周第二天，DeepEP：专为MoE和专家并行打造的高性能通信库

🔥DeepEP 主要特点

💡 高效 GPU 通信内核：提供高吞吐、低延迟的 all-to-all GPU 内核（MoE dispatch & combine）。

💾 低精度支持：支持 FP8 计算，提升计算效率。

🚀 优化 DeepSeek-V3 组限流门控算法：特别针对 NVLink 和 RDMA 之间的数据传输进行优化。

⚡ 低延迟推理解码：采用 纯 RDMA 方案，最小化延迟，并且支持计算-通信重叠，避免占用 GPU 计算资源。

📊DeepEP 性能测试
🚀 标准内核（NVLink & RDMA）
测试环境：H800 GPU + CX7 InfiniBand 400Gb/s RDMA

https://github.com/deepseek-ai/DeepEP

DeepEP

DeepEP 是一个专为 Mixture-of-Experts (MoE) 和专家并行 (EP) 通信优化的库。它提供了高吞吐量和低延迟的全对全 GPU 内核（也称为 MoE 分发和合并）。该库还支持低精度操作，包括 FP8。

为了与 DeepSeek-V3 论文中提出的组限制门控算法保持一致，DeepEP 提供了一套针对不对称域带宽转发优化的内核，例如从 NVLink 域到 RDMA 域的数据转发。这些内核提供高吞吐量，适用于训练和推理预填充任务。此外，它们还支持 SM（流式多处理器）数量控制。

对于延迟敏感的推理解码，DeepEP 包含了一套使用纯 RDMA 的低延迟内核，以最大限度地减少延迟。该库还引入了一种基于钩子的通信计算重叠方法，不会占用任何 SM 资源。

注意：该库中的实现可能与 DeepSeek-V3 论文有一些细微的差异。

性能

带 NVLink 和 RDMA 转发的普通内核

我们在 H800（~160 GB/s NVLink 最大带宽）上测试了普通内核，每个 H800 连接到一个 CX7 InfiniBand 400 Gb/s RDMA 网卡（~50 GB/s 最大带宽）。我们遵循 DeepSeek-V3/R1 预训练设置（每批 4096 个 token，7168 个隐藏单元，top-4 组，top-8 专家，FP8 分发和 BF16 合并）。

在这里插入图片描述

pure RDMA 的低延迟内核

我们在 H800（~160 GB/s NVLink 最大带宽）上测试了低延迟内核，每个 H800 连接到一个 CX7 InfiniBand 400 Gb/s RDMA 网卡（~50 GB/s 最大带宽）。我们遵循典型的 DeepSeek-V3/R1 生产设置（每批 128 个 token，7168 个隐藏单元，top-8 专家，FP8 分发和 BF16 合并）。

在这里插入图片描述

快速入门

要求

Hopper GPU（可能以后支持更多架构或设备）
Python 3.8 及以上
CUDA 12.3 及以上
PyTorch 2.1 及以上
节点内通信的 NVLink
节点间通信的 RDMA 网络

下载并安装 NVSHMEM 依赖项

DeepEP 还依赖于我们修改过的 NVSHMEM。请参考我们的 NVSHMEM 安装指南获取安装说明。

开发

# 构建并为 SO 文件创建符号链接
NVSHMEM_DIR=/path/to/installed/nvshmem python setup.py build
# 你可以根据自己的平台修改具体的 SO 名称
ln -s build/lib.linux-x86_64-cpython-38/deep_ep_cpp.cpython-38-x86_64-linux-gnu.so

# 运行测试用例
# 注意：你可以根据自己的集群设置修改 `tests/utils.py` 中的 `init_dist` 函数，并启动多个节点
python tests/test_intranode.py
python tests/test_internode.py
python tests/test_low_latency.py

安装

NVSHMEM_DIR=/path/to/installed/nvshmem python setup.py install

然后，在你的 Python 项目中导入 deep_ep，并享受！

网络配置

DeepEP 已经在 InfiniBand 网络上进行了全面测试。然而，它理论上也兼容 RDMA over Converged Ethernet (RoCE)。

流量隔离

InfiniBand 通过虚拟车道 (VL) 支持流量隔离。

为了防止不同类型流量之间的干扰，我们建议按照以下方式在不同的虚拟车道中隔离工作负载：

使用普通内核的工作负载
使用低延迟内核的工作负载
其他工作负载

对于 DeepEP，你可以通过设置 NVSHMEM_IB_SL 环境变量来控制虚拟车道的分配。

自适应路由

自适应路由是 InfiniBand 交换机提供的一种高级路由功能，可以将流量均匀分布在多条路径上。目前，低延迟内核支持自适应路由，而普通内核不支持（可能很快会添加支持）。为普通节点间内核启用自适应路由可能会导致死锁或数据损坏问题。

对于低延迟内核，启用自适应路由可以完全消除由路由冲突引起的网络拥塞，但也会引入额外的延迟。我们建议以下配置以获得最佳性能：

在网络负载较重的环境中启用自适应路由
在网络负载较轻的环境中使用静态路由

拥塞控制

我们没有在生产环境中观察到显著的拥塞，因此禁用了拥塞控制。

接口和示例

在模型训练或推理预填充中的示例用法

普通内核可以用于模型训练或推理预填充阶段（不包括反向部分），如下示例代码所示。

import torch
import torch.distributed as dist
from typing import List, Tuple, Optional, Union

from deep_ep import Buffer, EventOverlap

# 通信缓冲区（将在运行时分配）
_buffer: Optional[Buffer] = None

# 设置要使用的 SM 数量
# 注意：这是一个静态变量
Buffer.set_num_sms(24)

# 你可以在框架初始化时调用此函数
def get_buffer(group: dist.ProcessGroup, hidden_bytes: int) -> Buffer:
    global _buffer

    # 注意：你也可以将 `get_*_config` 替换为通过所有测试自动调整的结果
    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)

    # 如果不存在或缓冲区大小不足，则分配一个缓冲区
    # 注意：网络的自适应路由配置**必须关闭**
    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]:
    # 注意：可选的 `previous_event` 表示你希望作为分发内核依赖项的 CUDA 事件
    # 它可能对通信计算重叠有用。更多信息请参考 `Buffer.dispatch` 的文档
    global _buffer

    # 在实际分发之前计算布局
    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)
    # 执行 MoE 分发
    # 注意：CPU 将等待 GPU 的信号到达，因此这与 CUDA 图不兼容
    # 更多高级用法，请参考 `dispatch` 函数的文档
    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)
    # 关于事件管理，请参考 `EventOverlap` 类的文档
    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

    # MoE 分发的反向过程实际上是合并
    # 更多高级用法，请参考 `combine` 函数的文档
    combined_grad_x, combined_grad_recv_topk_weights, event = \
        _buffer.combine(grad_recv_x, handle, topk_weights=grad_recv_topk_weights, async_finish=True)

    # 关于事件管理，请参考 `EventOverlap` 类的文档
    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

    # 执行 MoE 合并
    # 更多高级用法，请参考 `combine` 函数的文档
    combined_x, _, event = _buffer.combine(x, handle, async_finish=True, previous_event=previous_event,
                                           allocate_on_comm_stream=previous_event is not None)

    # 关于事件管理，请参考 `EventOverlap` 类的文档
    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

    # MoE 合并的反向过程实际上是分发
    # 更多高级用法，请参考 `combine` 函数的文档
    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)

    # 关于事件管理，请参考 `EventOverlap` 类的文档
    return grad_x, event

此外，在 dispatch 函数中，我们可能不知道当前排名要接收多少 token。因此，将涉及隐式 CPU 等待 GPU 接收计数信号，如下图所示。

在这里插入图片描述

在推理解码中的示例用法

低延迟内核可以用于推理解码阶段，如下示例代码所示。

import torch
import torch.distributed as dist
from typing import Tuple, Optional

from deep_ep import Buffer

# 通信缓冲区（将在运行时分配）
# 注意：低延迟内核没有 SM 控制 API
_buffer: Optional[Buffer] = None

# 你可以在框架初始化时调用此函数
def get_buffer(group: dist.ProcessGroup, num_max_dispatch_tokens_per_rank: int, hidden: int, num_experts: int) -> Buffer:
    # 注意：低延迟模式将消耗比普通模式更多的空间
    # 因此我们建议 `num_max_dispatch_tokens_per_rank`（解码引擎中的实际批量大小）应小于 256
    global _buffer
    num_rdma_bytes = Buffer.get_low_latency_rdma_size_hint(num_max_dispatch_tokens_per_rank, hidden, group.size(), num_experts)

    # 如果不存在或缓冲区大小不足，则分配一个缓冲区
    if _buffer is None or _buffer.group != group or not _buffer.low_latency_mode or _buffer.num_rdma_bytes < num_rdma_bytes:
        # 注意：为了获得最佳性能，QP 数量**必须**等于本地专家的数量
        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

    # 执行 MoE 分发，与 CUDA 图兼容（但在重放时可能会恢复一些缓冲区状态）
    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)

    # 注意：只有在调用 `hook()` 时，实际张量才会被接收，
    # 这对于双批量重叠很有用，但**不占用任何 SM**
    # 如果你不想重叠，请设置 `return_recv_hook=False`
    # 稍后，你可以使用我们的 GEMM 库以这种特定格式进行计算
    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

    # 执行 MoE 合并，与 CUDA 图兼容（但在重放时可能会恢复一些缓冲区状态）
    combined_hidden_states, event_overlap, hook = \
        _buffer.low_latency_combine(hidden_states, topk_idx, topk_weights, handle,
                                    async_finish=False, return_recv_hook=True)

    # 注意：与分发内核描述的行为相同
    return combined_hidden_states, event_overlap, hook

对于两个微批量重叠，你可以参考以下图表。通过我们的接收钩子接口，RDMA 网络流量在后台发生，不会占用计算部分的任何 GPU SM。但请注意，重叠的部分可以调整，即注意力/分发/MoE/合并的四个部分可能没有完全相同的执行时间。你可以根据工作负载调整阶段设置。

在这里插入图片描述

注意事项

为了极致性能，我们发现并使用了一个未在文档中说明的 PTX 指令：ld.global.nc.L1::no_allocate.L2::256B。此指令将导致未定义行为：使用非一致性只读 PTX 修饰符 .nc 访问易失性 GPU 内存。但在 Hopper 架构上，使用 .L1::no_allocate 可以保证正确性，并且性能会更好。如果在其他平台上发现内核无法正常工作，你可以通过在 setup.py 中添加 DISABLE_AGGRESSIVE_PTX_INSTRS=1 来禁用此功能，或者提交问题。
为了在你的集群上获得更好的性能，我们建议运行所有测试并使用最佳自动调整的配置。默认配置已在 DeepSeek 的内部集群上进行了优化。