DeepEP 是一个专为混合专家(Mixture-of-Experts, MoE)模型和专家并行(Expert Parallelism, EP)设计的通信库。它提供了高吞吐量、低延迟的全对全(all-to-all)GPU 内核,也就是 MoE 中的分发(dispatch)和合并(combine)操作,同时支持低精度运算,如 FP8。
相关前置知识
首先要理解 MoE 模型理论 DeepSeek MoE,配合 DeepEP 通信的机制 DeepEP 单节点通信 理解代码,以及需要一些通用知识。
- 分布式计算
- 分布式训练基础:理解数据并行、模型并行和专家并行的概念,以及它们在深度学习训练中的应用。DeepEP 主要聚焦于专家并行,需明白如何在多个设备或节点间分配和同步专家相关的计算与通信任务。
- MPI(Message Passing Interface):分布式通信标准,DeepEP 的部分测试用例使用 MPI 进行多进程通信,因此需要了解 MPI 的基本概念,如进程管理、消息传递(发送、接收)和集体通信操作(广播、归约等)。
- GPU 编程
- CUDA 编程:DeepEP 大量使用 CUDA 进行 GPU 加速计算,以及一部分 PTX,需掌握 CUDA 的基本概念,如线程块、网格、共享内存、核函数等,以及如何在 Python 中通过 PyTorch C++ 后端调用 CUDA 相关功能。
- GPU 体系结构: GPU 的硬件架构,如流多处理器(SM)、内存层次结构(全局内存、共享内存、寄存器等),有助于理解 DeepEP 中对 GPU 资源的管理和优化策略。
- NVSHMEM:DeepEP 依赖修改后的 NVSHMEM,需要了解 NVSHMEM 的基本概念和使用方法。
- 网络通信
- InfiniBand 和 RoCE:DeepEP 主要在 InfiniBand 网络上进行测试,理论上也支持 RDMA over Converged Ethernet (RoCE)。需要了解这两种网络技术的基本原理、特点以及如何在分布式系统中配置和使用它们。
- RDMA(Remote Direct Memory Access):DeepEP 的低延迟内核依赖 RDMA 技术,需掌握 RDMA 的基本概念和工作原理,以及它如何实现高效的远程内存访问。
- 深度学习框架
- PyTorch:DeepEP 提供 Python 接口,基于 PyTorch C++ 后端实现接口。为了在 PyTorch 项目中集成和使用 DeepEP 也需要熟悉 PyTorch 的张量操作、自动求导、分布式训练 API 等。
- 其他加速框架:DeepEP 也被用于如 Megatron 这类分布式框架中。
File Hierarchy
DeepEP/
├─csrc/ # 包含 C++ 和 CUDA 源代码,用于实现核心功能
│ │ CMakeLists.txt # CMake 配置文件,用于调试项目,正式设置使用 Torch 扩展
│ │ config.hpp # 定义配置相关的结构体和类,如 LowLatencyLayout 结构体,用于计算低延迟布局所需的内存大小
│ │ deep_ep.cpp # 实现 Buffer 类,负责初始化通信缓冲区,包括内存分配、IPC 句柄设置、计数器初始化等操作,提供缓冲区同步、销毁等功能
│ │ deep_ep.hpp # 声明 Buffer 类及其成员函数,定义缓冲区的各种属性和操作接口,是 C++ 代码的头文件
│ │ event.hpp # 可能包含事件处理相关的类和函数,用于管理 CUDA 事件,实现通信和计算的重叠
│ └─ kernels/ # 包含 CUDA 内核代码
# 配置、API、异常处理
│ api.cuh # 定义 CUDA 内核的 API 接口,供其他部分调用
│ buffer.cuh # 与缓冲区操作相关的 CUDA 内核代码,可能包括缓冲区的分配、释放、读写等操作
│ CMakeLists.txt # 用于构建 CUDA 内核库的 CMake 配置文件
│ configs.cuh # 包含 CUDA 内核的配置信息,如线程块大小、网格大小等,计算 NVLink 和 RDMA 缓冲区的大小,`LowLatencyLayout` 结构体用于计算低延迟布局所需的内存大小。
│ exception.cuh # 定义异常处理相关的类和函数,用于处理 CUDA 内核中的异常情况
# 通信
│ ibgda_device.cuh # 与 InfiniBand 设备相关的 CUDA 内核代码,可能用于实现 RDMA 通信
│ [internode.cu](http://internode.cu/) # 实现节点间通信的 CUDA 内核代码,用于在不同节点之间传输数据
│ internode_ll.cu # 实现低延迟节点间通信的 CUDA 内核代码,采用纯 RDMA 技术,减少延迟
│ [intranode.cu](http://intranode.cu/) # 实现节点内通信的 CUDA 内核代码,用于在同一节点内的不同设备之间传输数据
# 启动和相关计算
│ launch.cuh # 包含 CUDA 内核的启动函数,负责启动和管理 CUDA 内核的执行
│ [layout.cu](http://layout.cu/) # 实现布局计算的 CUDA 内核代码,用于计算数据在缓冲区中的布局
│ [runtime.cu](http://runtime.cu/) # 包含运行时相关的 CUDA 内核代码,如内存管理、事件处理等
│ utils.cuh # PTX 汇编实现的一些实用的 CUDA 内核函数,如数据转换、内存填充等
# Python 接口
├─deep_ep/ # Python 代码的主要目录,提供 Python 接口
│ [buffer.py](http://buffer.py/) # 实现 dispatch 函数,用于将令牌分发到不同 rank,支持节点内和节点间设置
│ [utils.py](http://utils.py/) # 定义 EventOverlap 类,用于管理 CUDA 事件,方便实现通信和计算的重叠;还提供了检查 NVLink 连接的函数
│ __init__.py # Python 包的初始化文件,导入必要的模块和类
├─tests/ # 包含测试相关的代码
│ test_internode.py # 测试节点间通信功能的脚本,验证节点间数据传输的正确性
│ test_intranode.py # 测试节点内通信功能的脚本,验证节点内数据传输的正确性
│ test_low_latency.py # 测试低延迟通信功能的脚本,验证低延迟内核的性能和正确性
│ [utils.py](http://utils.py/) # 定义 bench 函数,用于对函数进行基准测试,包括预热、测试和计算平均、最小、最大时间等操作
项目主要分为三个部分:Python 接口、C++ 封装层和 CUDA 内核层。
- Python 接口提供给用户使用,通过 Python 调用可用于构建专家并行(EP)通信缓冲区
Buffer,以支持混合专家(MoE)模型的通信操作; - C++ 封装层
deep_ep.cpp通过buffer负责处理 Python 和 CUDA 之间的交互,deep_ep.hpp提供了接口定义,; - CUDA 内核层实现具体的计算逻辑,承担起管理通信缓冲区、执行内核操作以及同步不同进程间信息的任务。
csrc/kernels 目录下包含了各种 CUDA 内核代码,如 internode.cu、intranode.cu 等,这些内核代码实现了具体的计算逻辑,如节点间和节点内的数据传输、布局计算等。
DeepEP 通信组件(软硬件协同)
依赖:
DeepEP Requirements CUDA >=12.3 NVLink for intranode communication RDMA network for internode communication NVSHMEM utils.cuh 定义了大量的 PTX 操作 (PTX),如 LD/ST 采用了 acquire/relaxed,在 kernel 中大量使用,进一步提高的处理效率
1. CUDA IPC 在 DeepEP 中的作用
技术背景
- CUDA IPC(Inter-Process Communication) 是 NVIDIA 提供的进程间通信机制,允许不同进程共享 GPU 显存。通过共享显存句柄(handle),进程可以直接访问远程 GPU 的内存,无需通过 CPU 中转。
- 应用场景:节点内多 GPU 之间的数据共享(如 NVLink 连接的 GPU)。
在 DeepEP 中的作用
- 节点内 GPU 通信优化:
在 MoE 模型的专家并行(EP)中,不同 GPU 上的专家需要频繁交换数据(如分发和合并操作)。CUDA IPC 可直接共享显存,避免显存拷贝到 CPU 或跨节点传输的开销。
- 示例:在 DeepEP 的
all-to-all通信中,专家间的数据分发(dispatch)和合并(combine)可利用 CUDA IPC 实现节点内 GPU 的高速通信。
- 低延迟分发与合并:
DeepEP 的 FP8 分发 和 BF16 合并 操作需要低延迟的显存访问。CUDA IPC 通过共享句柄直接传递数据,减少 CPU 干预,降低通信延迟。
- 资源利用率提升:
通过 CUDA IPC 避免 CPU 中转,GPU 的计算资源(SM)可以专注于模型计算,而非等待数据传输完成。
2. Nvshmem 在 DeepEP 中的作用
技术背景
- nvshmem 是 NVIDIA 基于 RDMA 技术开发的共享内存通信库,支持跨节点的高效数据传输。其核心特性包括:
- RDMA over InfiniBand:绕过操作系统内核,直接访问远程节点内存,实现低延迟、高带宽通信。
- All-to-All 通信:支持高效的点对点(P2P)和集合通信(如 AllReduce、AlltoAll)。
- 异步操作:允许计算与通信重叠,提升整体效率。
在 DeepEP 中的作用
- 跨节点通信优化:
在 MoE 模型的分布式训练中,专家可能分布在多个节点上。nvshmem 通过 RDMA 技术实现跨节点的高速数据传输,解决节点间通信瓶颈。
- 示例:DeepEP 的
非对称域带宽转发(如从 NVLink 域到 RDMA 域)依赖 nvshmem 的异构通信能力。
- 低延迟推理解码:
DeepEP 提供 纯 RDMA 的低延迟内核,适用于推理解码阶段。通过 nvshmem 的异步通信,解码任务的延迟可降低至 163 微秒(根据知识库数据)。
- 动态资源控制:
nvshmem 支持灵活的资源管理(如 SM 数量控制),DeepEP 利用这一特性优化 GPU 资源分配,适配不同规模的训练任务。
- 与 FP8 压缩结合:
DeepEP 的 FP8 分发操作通过 nvshmem 的 RDMA 传输,进一步压缩数据体积,减少带宽占用(节省 50% 带宽资源)。
3. CUDA IPC 与 Nvshmem 的协同作用
在 DeepEP 中,CUDA IPC 和 nvshmem 通常 分工协作,覆盖节点内和节点间的通信需求:
| 场景 | 技术 | 作用 |
|---|---|---|
| 节点内 GPU 通信 | CUDA IPC | 通过共享显存句柄,实现 NVLink 连接 GPU 的高速通信(单机带宽可达 158 GB/s)。 |
| 跨节点通信 | nvshmem | 利用 RDMA 技术,实现跨节点的低延迟、高带宽传输(单网卡带宽 43-47 GB/s)。 |
| 通信 - 计算重叠 | nvshmem | 异步通信与计算并行执行,减少 GPU 空闲时间(提升资源利用率)。 |
| 低延迟推理解码 | nvshmem | 通过纯 RDMA 内核,最小化延迟(<163 微秒),适合实时推理场景。 |
4. 技术解析:DeepEP 的通信优化
关键技术点
- All-to-All 通信内核:
- DeepEP 的
dispatch和combine操作基于 CUDA IPC 和 nvshmem 实现全对全通信,适配 MoE 的非对称数据分布。 - 示例代码片段:
from deep_ep import Buffer, EventOverlap
buffer = Buffer(…) # 初始化通信缓冲区
buffer.dispatch(…) # 分发操作
buffer.combine(…) # 合并操作
- 计算 - 通信重叠:
- DeepEP 通过 hook-based 机制 实现通信与计算的完全分离,无需占用 SM 资源。例如:
- 在等待 RDMA 数据传输时,GPU 可继续执行其他计算任务。
- 异构网络优化:
- DeepEP 支持 NVLink + RDMA 的混合网络架构,通过 nvshmem 的异构通信接口,灵活适配不同硬件配置。
总结
| 技术 | 核心作用 |
|---|---|
| CUDA IPC | 优化节点内 GPU 通信,实现 NVLink 高速数据共享,减少 CPU 中转。 |
| nvshmem | 优化跨节点通信,通过 RDMA 技术实现低延迟、高带宽传输,支持异步计算 - 通信重叠。 |
| DeepEP 协同 | CUDA IPC 与 nvshmem 结合,覆盖节点内和节点间通信,提升 MoE 模型的训练和推理效率。 |
DeepEP 用于 MoE 的流程概述
1. 硬件架构假设
假设系统由多台服务器组成,每台服务器包含:
- 8 个 GPU(通过 NVLink 互连)
- 每个 GPU 拥有 HBM 显存(高带宽内存)
- 服务器间通过 InfiniBand 网络连接
- CPU 内存(用于数据暂存和控制流)
2. 核心流程与软硬件分工
- 初始化通信缓冲区:在开始训练或推理之前,需要初始化通信缓冲区。这涉及到设置缓冲区的大小和其他参数,以确保能够有效地处理数据的分发和组合。
- 计算分发布局:在进行实际的分发操作之前,需要计算数据的分发布局。这包括确定每个 rank 接收的令牌数量、每个专家接收的令牌数量等信息。
- 执行分发操作:根据计算得到的分发布局,执行实际的分发操作。这涉及到将数据从一个 rank 发送到其他 rank,并接收来自其他 rank 的数据。
- 执行组合操作:在数据分发完成后,需要执行组合操作,将从其他 rank 接收到的数据组合成一个完整的张量。
- 反向传播:在反向传播阶段,分发和组合操作的过程与正向传播相反。分发操作的反向过程实际上是组合操作,而组合操作的反向过程实际上是分发操作。
# Python接口层 (CPU执行)
def moe_forward(input_tokens, router_logits):
# 1. 路由器计算 (CPU/GPU)
topk_experts, topk_weights = router(input_tokens, router_logits)
# 2. 创建通信Buffer (CPU初始化)
buffer = Buffer(group_size=8) # 对应8个GPU
# 3. 计算分发布局 (CPU/GPU)
layout = buffer.get_dispatch_layout(topk_experts)
# 4. 数据分发 (GPU执行,CPU同步)
# 调用C++封装层
dispatched_tokens, expert_indices = buffer.dispatch(
input_tokens, topk_experts, layout
)
# 5. 专家计算 (GPU执行)
expert_outputs = []
for i in range(num_gpus):
# 每个GPU执行本地专家计算
expert_outputs.append(expert_forward(
dispatched_tokens[i],
expert_weights[i] # 本地专家参数
))
# 6. 结果合并 (GPU执行,CPU同步)
combined_output = buffer.combine(expert_outputs, expert_indices)
return combined_output
3. 数据流向与存储位置
| 组件 | 存储位置 | 说明 |
|---|---|---|
| 原始 Tokens | GPU HBM | 输入序列,初始存储在发起计算的 GPU 显存中 |
| 分发布局 | GPU HBM + CPU 内存 | 布局元数据在 CPU 计算,最终传输到 GPU 用于指导数据分发 |
| 通信 Buffer | GPU HBM | 由 Buffer 类管理,用于暂存待发送 / 已接收的 Tokens,完全在 GPU 显存中 |
| 专家参数 | 各 GPU HBM | 每个 GPU 存储部分专家参数,例如 GPU0 存储专家 0-31,GPU1 存储专家 32-63 |
| 中间结果 | GPU HBM | 专家计算的输出结果,存储在执行计算的 GPU 中 |
4. 专家分布策略
DeepEP 支持两种主要的专家分布模式:
模式 1:节点内专家并行
- 场景:单机多 GPU 系统(如 8 卡服务器)
- 分布方式:
- 每个 GPU 存储一部分专家(例如每个 GPU 存储总专家数的 1/8)
- 专家 ID 与 GPU 的映射关系为:
gpu_id = expert_id % num_gpus - 数据流向:
- Tokens 根据路由结果被分发到对应的 GPU
- 例如,Token 被路由到专家 17,则该 Token 会被发送到 GPU1(17 % 8 = 1)
模式 2:节点间专家并行
- 场景:多机系统(例如 4 台服务器,每台 8 卡)
- 分布方式:
- 专家按层次分布:
- 第一层:跨服务器划分(例如服务器 0 存储专家 0-255,服务器 1 存储 256-511)
- 第二层:服务器内 GPU 划分(例如服务器 0 的 GPU0 存储专家 0-31)
- 映射关系:
server_id = expert_id // (experts_per_server) - 数据流向:
- Tokens 先通过 InfiniBand 发送到目标服务器
- 再通过 NVLink 发送到服务器内的目标 GPU
通信加速流程详解
假设一个 Token 被门控网络分配给 3 个专家:其中 2 个在 Machine 2(同机器),1 个在 Machine 5(跨机器)。DeepEP 的传输逻辑通过 “跨节点合并传输 + 节点内高效分发” 降低带宽压力:
1. 跨节点传输(Infiniband,目标是减少数据量)
- 合并同机器 Token:发送到 Machine 2 的 Token 数据只传 1 份(而非 2 份),发送到 Machine 5 的 Token 传 1 份。
- 通信载体:通过 NVSHMEM 机制(节点间单向通信)传输,跳过传统 Send/Receive 的双向握手,减少延迟。
Token传输第一步(跨节点):
源机器(如Machine 0)
├─ 打包发往Machine 2的Token → 1份数据(供Machine 2内2个专家共享)
└─ 打包发往Machine 5的Token → 1份数据(供Machine 5内1个专家使用)2. 节点内分发(NVLink,利用高带宽)
- 同机器内共享数据:Machine 2 收到 Token 后,通过 NVLink(160GB/s 带宽,远高于 Infiniband 的 50GB/s)将 1 份 Token 数据分发到 2 个专家所在的 GPU。
- 通信载体:使用 IPC 机制(节点内单向通信),结合显存一致性,无需 CPU 参与,直接在 GPU 间传输。
Token传输第二步(节点内):
Machine 2接收端
├─ 收到1份Token数据(存在本地显存)
├─ 通过NVLink+IPC → 发送到Expert A所在GPU
└─ 通过NVLink+IPC → 发送到Expert B所在GPU(同一份数据复用,无额外传输)3. DeepEP 的通信核心机制为什么高效?
1. 节点间:NVSHMEM(替代传统 NCCL/Send/Receive)
- 单向操作:无需等待接收方确认(传统 NCCL 是对称通信,需双向同步),发送方直接写入接收方的显存地址(基于显存一致性)。
- 无锁同步:通过显存一致性机制(如 GPU 显存的原子操作)实现同步,避免锁竞争,适合高并发场景。
2. 节点内:IPC(替代传统 PCIe 通信)
- 本地直连:同一机器内 GPU 通过 NVLink 物理直连,IPC 机制直接映射显存地址,数据无需拷贝,直接访问。
- 与 NVLink 协同:收到 RDMA 数据后,立即通过 NVLink 分发(如技术文档提到 “RDMA 收到后马上用 NVLink 发送到专家”),减少中间缓存延迟。
3. 对比传统通信(优势在哪里?)
| 通信场景 | 传统方式(NCCL/Send) | DeepEP(NVSHMEM+IPC) | 优势总结 |
|---|---|---|---|
| 节点间通信 | 双向握手,数据需多次拷贝 | 单向写入,直接显存访问 | 延迟降低 30%+,带宽利用率提升 |
| 节点内多专家 | 每个专家单独传输,重复数据 | 1 份数据复用,NVLink 分发 | 数据量减少(n→1),节省带宽 |
4. 关键优化点
这一节可以在理解整体流程和代码细节后再来理解通信加速设计。
| 优化维度 | 核心技术 | 具体实现 | 效果 / 优势 | 关联硬件 / 场景 |
|---|---|---|---|---|
| 通信与计算重叠 | CUDA 流与事件管理 | 在等待 Tokens 接收时执行本地计算 | 隐藏通信延迟,提升 GPU 利用率 | 所有计算密集型场景 |
| 低延迟设计 | 纯 RDMA 模式 | 避免 CPU 参与,GPU 直接访问远程内存(internode_ll.cu) | 减少数据拷贝路径,延迟降低 50%+ | 跨节点通信(Infiniband) |
| 内存优化 | GPU HBM 直接使用 + 动态 Buffer 调整 | Buffer 常驻 GPU 显存,根据 Token 分布动态分配大小 | 减少 CPU-GPU 数据传输,内存利用率提升 20%+ | 显存受限场景(如大模型推理) |
| 数据传输优化 | 同节点专家数据共享 | 跨节点传输量从 “专家数” 降至 “机器数” | 减少跨节点通信量,降低 Infiniband 带宽压力 | 多专家并行模式 |
| 通信机制优化 | 单向通信(NVSHMEM/IPC)+ 无锁同步 | 替代传统 Send/Receive,使用环形缓冲区和原子计数器 | 减少同步开销,提升并发能力 | 高并发通信场景 |
| 硬件特性利用 | NVLink 高带宽 + SM 专用分配 + PTX 指令优化 | 预留 20 个 SM 处理通信,使用 PTX 穿透 L2 缓存 | 充分发挥硬件性能,通信延迟降低 30%+ | NVIDIA GPU 集群 |
| 流水线处理 | RDMA 与 NVLink 衔接 | RDMA 接收未完成时即开始 NVLink 分发 | 隐藏传输延迟,提升整体吞吐量 | 多 GPU 节点间数据流动 |
对系统各层的优化在于:
- 硬件适配层
- SM 专用分配:将 GPU 的 Streaming Multiprocessors 分为计算核心和通信核心,避免资源竞争。
- PTX 指令优化:直接操作 GPU 底层硬件,绕过 L2 缓存,减少数据访问延迟。
- 通信协议层
- 纯 RDMA 模式:基于 InfiniBand 的远程直接内存访问,实现跨节点的高效数据传输。
- NVSHMEM/IPC 单向通信:无需双向握手,降低通信协议开销。
- 内存管理层
- HBM 直接使用:将关键数据结构驻留在 GPU 高带宽内存中,避免 CPU 参与。
- 动态 Buffer 调整:根据实际 Token 分布,弹性分配内存资源。
- 算法优化层
- 同节点数据共享:基于 MOE 模型的专家分配策略,减少冗余数据传输。
- 流水线处理:利用数据传输的时间窗口,提前进行后续处理。
5. 两种通信模式
DeepEP 针对不同场景设计了灵活的通信模式(低延迟模式与标准模式),并通过多种优化机制实现高效数据处理。
1. 低延迟模式(专为推理设计)
- 核心目标:最小化单次请求的响应时间(Latency)。
- 技术特点:
- 纯 RDMA 通信:通过网络适配器(如 InfiniBand)在 GPU 间传输数据。
- 同步请求 - 响应:采用类似 RPC 的模式,每个请求立即触发通信和计算,无需批量处理。
- 预分配资源:提前分配 GPU 内存和通信缓冲区,避免动态分配延迟。
- 适用场景:小批量请求(如单用户对话)、对延迟敏感的任务。
2. 标准模式(适合训练与批处理)
- 核心目标:最大化吞吐量(Throughput)。
- 技术特点:
- 批量操作:将多个 token 请求合并为一个通信批次,减少通信次数。
- 计算与通信重叠:利用 CUDA 流(CUDA Stream),在 GPU 计算时同时进行数据传输。
- 动态负载均衡:根据专家负载动态调整 token 分配,避免热点问题。
- 适用场景:大批量请求、对吞吐量要求高的任务。
多机多卡分布举例讲解
要理解 DeepEP 在 MOE 模型专家并行中的作用,核心逻辑是:将 “专家(Expert)” 分布在多卡 / 多机器上,通过 DeepEP 高效传递 “token(待处理数据)” 并聚合结果。先明确核心概念:
- MOE 模型:类似 “多个专家处理不同任务,门控决定谁来处理”—— 输入 token 先经过门控网络,得到每个 token 对专家的 “偏好分数”,再分配给 Top-K 个专家处理,最后聚合结果。
- 专家并行(Expert Parallelism):将大量专家(比如 128 个)分散到多卡 / 多机器上(而非单卡容纳所有专家),解决单卡内存不足问题。
- DeepEP:负责 token 从输入设备到专家所在设备的分发(dispatch),以及专家输出结果的聚合(combine),底层优化了通信效率。
下面讲解从 Expert 分配到 Token 传递的流程。
**第一步:Expert 如何分配到多卡 / 多机器?
假设场景:2 台机器(Machine A、Machine B),每台机器有 2 张 GPU 卡(Card 0、Card 1),共 16 个专家(Expert 0~15)。
| 设备层级 | 分配逻辑(DeepEP 支持灵活划分) | 可视化类比 |
|---|---|---|
| 机器(Machine) | 按 “就近通信” 原则划分专家组(比如 Machine A 负责前 8 个专家,Machine B 负责后 8 个) | 2 个工厂,各负责一半任务 |
| GPU 卡(Card) | 单台机器内的 GPU 平均分配本组专家(Machine A 的 Card 0 负责 Expert 0 | 每个工厂有 2 个车间,各分一半工人 |
| 专家(Expert) | 每个 GPU 卡上的专家固定绑定设备(如 Card 0 上的 Expert 0~3 独占该卡的计算资源) | 每个车间有 4 个工人(专家) |
分配结果图示:
Machine A(机器A)
├─ Card 0(GPU卡0):Expert 0、Expert 1、Expert 2、Expert 3
└─ Card 1(GPU卡1):Expert 4、Expert 5、Expert 6、Expert 7
Machine B(机器B)
├─ Card 0(GPU卡0):Expert 8、Expert 9、Expert 10、Expert 11
└─ Card 1(GPU卡1):Expert 12、Expert 13、Expert 14、Expert 15注:DeepEP 支持 “组限制门控”(Group-restricted Gating),可将专家按机器 / 卡分组,让 token 优先分配到同组专家(减少跨机器通信),比如限制 token 只能选择 Machine A 内的专家,降低跨机延迟。
第二步:Token 通过 DeepEP 传递
假设输入 10 个 token,门控网络为每个 token 选择 Top-2 个专家(比如 Token 0 选择 Expert 2 和 Expert 9),整个流程分 3 步:
**阶段 1:门控网络生成 “分配指令”
- 输入 token 先经过门控网络(Gating Network),计算每个 token 对所有专家的 “偏好分数”(如 Token 0 对 Expert 2 的分数 0.8,对 Expert 9 的分数 0.7)。
- 筛选 Top-K 个专家(比如 K=2),生成 “token - 专家” 分配表(谁该去哪个专家)。
门控输出(分配表):
Token 0 → Expert 2(Machine A, Card 0)、Expert 9(Machine B, Card 0)
Token 1 → Expert 3(Machine A, Card 0)、Expert 8(Machine B, Card 0)
…(共10个Token,每个对应2个专家)阶段 2:DeepEP 的 Dispatch
这是 DeepEP 的核心功能之一:根据分配表,将 token 从 “输入设备” 发送到 “专家所在设备”。
- 输入设备:假设所有 token 初始在 “输入卡”(比如 Machine A, Card 0,也可分散在多卡)。
- DeepEP 优化点:
- 用高效通信内核(如基于 NVLink/RDMA)直接传输,跳过 CPU 中转;
- 按 “设备分组” 批量发送(比如发往 Machine A 的 token 打包,发往 Machine B 的 token 打包),减少通信次数。
Dispatch流程(数据流向):
输入卡(Machine A, Card 0)
├─ 批量发送Token 0、Token 1中去Machine A, Card 0的部分 → Expert 2、Expert 3(本地卡内通信,最快)
└─ 通过RDMA发送Token 0、Token 1中去Machine B, Card 0的部分 → Expert 9、Expert 8(跨机器通信,DeepEP优化带宽)阶段 3:专家计算 + DeepEP 的 Combine(聚合专家输出)
- 专家计算:每个专家在本地卡上处理收到的 token(比如 Expert 2 处理 Token 0,Expert 9 也处理 Token 0),输出 “专家对 token 的处理结果”。
- DeepEP 的 Combine:将同一个 token 在不同专家的输出 “拉回原设备并聚合”(比如 Token 0 的两个专家输出,需传回原输入卡并加权求和)。
Combine流程(数据流向):
Expert 2(Machine A, Card 0)→ 输出Token 0的结果 → 通过本地通信传回输入卡
Expert 9(Machine B, Card 0)→ 输出Token 0的结果 → 通过RDMA传回输入卡
DeepEP在输入卡聚合:Token 0最终结果 = (Expert 2输出 × 0.8) + (Expert 9输出 × 0.7)总结:DeepEP 在通信中的作用
| 步骤 | 传统通信问题 | DeepEP 的优化(可视化理解) |
|---|---|---|
| Dispatch | 小批量传输多、跨机延迟高 | 像 “快递批量分拣”:按设备分组打包,用专用通道(RDMA)传输 |
| Combine | 聚合时数据乱序、带宽浪费 | 像 “包裹汇总”:按 token ID 整理结果,高效合并传输 |
| 整体 | 通信占比高,拖慢训练 / 推理 | 通信与计算重叠(专家计算时提前传输),隐藏延迟 |