DeepEP 是一个专为 混合专家模型(MoE) 和 专家并行(EP) 设计的通信库,旨在提供高效的 GPU 通信能力,支持高吞吐、低延迟的分布式计算,适用于训练和推理场景。
核心功能
-
高性能通信内核
- 正常内核:支持 NVLink(GPU 间高速互联)和 RDMA(远程直接内存访问)混合转发,适用于训练和推理预填充阶段。
- 低延迟内核:纯 RDMA 设计,专为推理解码阶段优化,减少通信延迟。
- 低精度支持:包括 FP8 等低精度操作,节省显存并提升效率。
-
性能表现
- 正常内核:在 H800 GPU 测试中,单节点内(Intranode)通信带宽可达 ~150 GB/s,跨节点(Internode)带宽约 ~45 GB/s。
- 低延迟内核:单次通信延迟最低 163 微秒,带宽稳定在 ~40 GB/s 以上,支持大规模并行(如 256 EP)。
快速上手
-
系统要求
- Hopper 架构 GPU(如 H800)、CUDA 12.3+、PyTorch 2.1+。
- 依赖修改版 NVSHMEM(需按指南安装)。
-
安装与测试
# 安装依赖并编译 NVSHMEM_DIR=/path/to/nvshmem python setup.py install # 运行测试(需根据集群配置调整参数) python tests/test_intranode.py python tests/test_internode.py
网络配置建议
- 流量隔离:使用 InfiniBand 的虚拟通道(VL)隔离不同任务(如正常内核、低延迟内核)。
- 自适应路由:低延迟内核支持该功能,可缓解网络拥堵,但可能增加延迟。
- 静态路由:轻负载环境下推荐,避免额外开销。
使用示例
-
训练/预填充阶段
# 初始化通信缓冲区 buffer = get_buffer(group, hidden_bytes) # 数据分发(Dispatch) recv_x, recv_idx, recv_weights, ... = buffer.dispatch(x, topk_idx, ...) # 数据合并(Combine) combined_x, event = buffer.combine(x, handle, ...)
-
推理解码阶段
# 低延迟模式初始化(需指定最大 token 数) buffer = get_buffer(group, max_tokens, hidden_size, num_experts) # 异步通信(支持 CUDA Graph) recv_states, expert_count, ... = buffer.low_latency_dispatch(...)
注意事项
- 性能优化:使用未公开的 PTX 指令(如
.nc
修饰符)提升性能,需在 Hopper GPU 上测试。 - 许可证:核心代码为 MIT 协议,但部分依赖 NVSHMEM 的代码受其 SLA 约束。
引用
若使用 DeepEP,可参考以下引用格式:
@misc{deepep2025,
title={DeepEP: 高效的专家并行通信库},
author={DeepSeek 团队},
year={2025},
howpublished={\url{https://github.com/deepseek-ai/DeepEP}}
}
总结:DeepEP 专注于优化 MoE 模型的分布式通信,通过灵活的配置和高效内核,显著提升训练和推理效率,适合大规模 AI 模型部署。