Deepseek DualPipe 分析
一、DualPipe 特点
1、计算和通信 overlap（下图为由 forward chunk 和 backward chunk 组成的 chunk 对的计算和通信流程）
- 每个 chunk 均会经历四个过程：attention、all-to-all dispatch、MLP 以及 all-to-all combine。橙色表示 forward，绿色表示 input 的 backward，蓝色表示 weight 的 backward，紫色表示 PP 通信，红色表示 barrier
- 从下面第一张图可以看出，PP 和 alltoall 的通信过程被隐藏在 MLP 和 attention 的计算过程中
- 从下面第二张图可以看出，DualPipe 使用了 bidirectional pipeline scheduling，即同时向流水线的首位两端传入 micro-batch 从而充分利用硬件资源
- 不同 PP 策略对应的 bubbles 数量和 memory usage 的对比
2、all-to-all communication 优化
- 定制化 kernels：设计过程结合了 MoE gating algorithm 和集群的 network topology
- IB 和 NVLink 的 overlapping：每个 token 最多只会被分配到 4 个不同节点上来减少 IB 流量，同时在分配 token 时让 IB 和 NVLink 的带宽充分 overlap，从而在不引入 NVLink 通信 overhead 的情况下使得某个节点上的 token 有能力选择 3.2 个 experts（即最多支持 4x3.2=13 experts，DS-v3 限制在 8 experts）
- SM 动态分配：将 20 个 SMs 动态分配到 10 个 communication channels 上。具体来说
- Dispatch 过程：IB sending、IB-to-NVLink forwarding、NVLink receiving 三个过程分别使用不同的 warps
- Combine 过程：NVLink sending、 NVLink-to-IB forwarding&accumulation、IB receiving&accumulation 三个过程同样根据 workloads 动态分配 warps
- PTX 指令优化：通过 PTX 编程和 communication chunk size 自动调节机制减少 L2 cache 的访问和对其他 SM 的影响
3、 优化 memory footprint
- 对 RMSNorm 和 MLA Up-Projection 过程进行重计算
- Exponential Moving Average in CPU
- Multi-Token Prediction 情况下对 Embedding 和多头注意力的输出放在同个 PP rank 内
评论
发表评论