DeepSeek 开源周：DeepEP 项目详解，GPU 压榨计划启动！

news2025/3/1 4:35:39

引言

就在今天，2025年2月25日，DeepSeek 再次为人工智能社区带来了一场技术盛宴——DeepEP 项目的开源。这个旨在优化 GPU 性能的工具一经发布便迅速获得了广泛的关注和赞誉，短短两小时内就斩获了超过1000个 Star。本文将详细介绍 DeepEP 的功能、应用场景以及如何使用它来提升 AI 训练和推理的效率。

DeepEP 概述

功能与作用

DeepEP 是一个专门针对 Mixture-of-Experts (MoE) 训练方法设计的通信工具，其核心目标是通过优化 GPU 之间的数据传输来加速 AI 的训练和推理过程。具体来说，DeepEP 能够：

提升 GPU 间的数据传输速度，无论是 NVLink（同一机器内的 GPU）还是 RDMA（跨机器的 GPU）。
减少推理解码阶段的延迟，这对于实时应用如 ChatGPT 至关重要。
实现计算和数据传输的并行执行，避免因等待数据传输而导致的计算停滞。

技术细节

DeepEP 采用了特殊的 hook 机制，使得 GPU 可以在进行数据传输的同时继续执行计算任务，从而提高了整体效率。此外，DeepEP 还支持低延迟模式，特别适合需要快速响应的大规模模型推理场景。

实际性能表现

根据测试，在 NVIDIA H800 GPU 加上 400Gb/s RDMA 的服务器环境下，DeepEP 展现出了令人印象深刻的性能指标：

普通模式：NVLink 速度可达 158GB/s，RDMA 速度在 43-47GB/s 之间。
低延迟模式：延迟可低至 163 微秒，RDMA 速度维持在 39-46GB/s。

这些数字表明，DeepEP 能够显著提高大规模 MoE 模型的训练和推理效率。

安装与使用指南

系统要求

为了充分利用 DeepEP 的性能优势，你需要具备以下硬件和软件环境：

GPU：Hopper 架构的 GPU（例如 H100 或 H800）
Python：版本 3.8 及以上
CUDA：版本 12.3 及以上
PyTorch：版本 2.1 及以上
网络设备：NVLink 和 RDMA 支持

安装步骤

首先，确保你已经安装了必要的依赖项，然后执行以下命令来安装 DeepEP：

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

使用示例

MoE 训练中的应用

以下是一个简单的代码示例，展示了如何在 MoE 训练过程中使用 DeepEP 来分配任务和合并结果：

import torch
import torch.distributed as dist
from deep_ep import Buffer

# 初始化通信 buffer
_buffer = None

def get_buffer(group, hidden_bytes):
    global _buffer
    if _buffer is None or _buffer.group != group:
        _buffer = Buffer(group, hidden_bytes, hidden_bytes)
    return _buffer

# MoE 分发任务
def dispatch_forward(x, topk_idx, topk_weights, num_experts):
    global _buffer
    recv_x, recv_topk_idx, recv_topk_weights, handle, event = _buffer.dispatch(x, topk_idx, topk_weights, num_experts)
    return recv_x, recv_topk_idx, recv_topk_weights, handle, event

这段代码创建了一个高效的“快速通道”，用于分发 MoE 任务并接收计算结果。

推理解码阶段的应用

在推理解码时，DeepEP 可以让数据在后台传输，不影响 GPU 的计算任务：

import torch
import torch.distributed as dist
from deep_ep import Buffer

_buffer = None

def get_buffer(group, num_max_tokens, hidden, num_experts):
    global _buffer
    if _buffer is None or _buffer.group != group:
        _buffer = Buffer(group, 0, num_max_tokens * hidden * num_experts, low_latency_mode=True)
    return _buffer

def low_latency_dispatch(hidden_states, topk_idx, num_max_tokens, num_experts):
    global _buffer
    recv_hidden_states, recv_expert_count, handle, event, hook = _buffer.low_latency_dispatch(hidden_states, topk_idx, num_max_tokens, num_experts)
    return recv_hidden_states, recv_expert_count, handle, event, hook