DeepEP 是一个专为 混合专家模型(MoE)专家并行(EP) 设计的通信库,旨在提供高效的 GPU 通信能力,支持高吞吐、低延迟的分布式计算,适用于训练和推理场景。


核心功能

  1. 高性能通信内核

    • 正常内核:支持 NVLink(GPU 间高速互联)和 RDMA(远程直接内存访问)混合转发,适用于训练和推理预填充阶段。
    • 低延迟内核:纯 RDMA 设计,专为推理解码阶段优化,减少通信延迟。
    • 低精度支持:包括 FP8 等低精度操作,节省显存并提升效率。
  2. 性能表现

    • 正常内核:在 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 模型部署。