第3章：算子融合策略：通信算子融合的动机与收益分析
聊到多卡训练，很多人第一反应就是“通信开销大”。没错，通信确实是瓶颈。但怎么解决？我个人的经验是——别总想着优化通信本身，而是想办法让它“消失”。
说白了，就是让通信和计算重叠起来。你算你的，我传我的，谁也不等谁。这就是通信算子融合的核心思想。
3.1 为什么要做通信算子融合？
先看一个典型场景。你在做数据并行训练，每张卡算完梯度后，需要做一次 AllReduce 把梯度同步一下。传统做法是：
- 前向传播完成
- 反向传播完成，梯度就绪
- 启动 AllReduce 通信
- 等待通信结束
- 更新参数
你看，第3步和第4步之间，GPU 在干嘛？啥也没干，干等着。这就是典型的“通信空洞”。
核心收益：通信算子融合可以把通信时间“藏”到计算时间里。理想情况下，通信开销降为0。
我在项目中遇到过这样一个案例：一个 64 卡的大模型训练任务，原始方案里通信占了 40% 的时间。做了算子融合后，通信时间几乎看不到了。嗯，效果就是这么直接。
3.2 常见的融合模式
融合模式其实就几种，我按实战中遇到的频率排个序：
3.2.1 AllReduce 与反向计算 Kernel 融合
这是最常用的一种。反向传播是一层一层算的，每层算完梯度后，其实可以立刻开始通信，不用等所有层都算完。
具体怎么做？
- 分桶策略：把梯度分成多个桶（bucket），每个桶算完就发
- 异步启动：计算下一个桶的同时，上一个桶的 AllReduce 已经在跑了
- 流水线化：形成计算-通信-计算-通信的流水线
你想想看，这就像工厂流水线。一个工人在组装，另一个在打包，互不耽误。
我的建议：桶的大小很关键。太小了通信次数多，太大了等待时间长。我一般从 1MB 开始调，根据实际带宽和计算量做微调。
3.2.2 AllGather 与前向计算融合
这个在模型并行场景下很常见。比如张量并行中，你需要把各卡的输出拼起来（AllGather），然后继续下一层的计算。
融合思路：
- 在计算当前层的同时，提前发起下一层的 AllGather
- 利用 NVLink 或 RDMA 的异步特性
- 注意内存管理，别把显存撑爆了
我曾经踩过一个坑：AllGather 和计算融合时，没注意显存复用，结果 OOM 了。后来加了个内存池，才搞定。
3.2.3 ReduceScatter 与梯度裁剪融合
这个稍微冷门一点，但在大模型训练中很实用。梯度裁剪（gradient clipping）需要计算全局梯度范数，而 ReduceScatter 正好可以顺便把这个活儿干了。
怎么做？
- 每卡算完局部梯度
- 在 ReduceScatter 的 reduce 阶段，顺便算一下平方和
- 最后拿到全局范数，直接做裁剪
省了一次额外的 AllReduce，收益很可观。
3.3 收益到底有多大？
我直接给一组实测数据吧。这是我在一个 8 卡 A100 集群上测的：
| 模型规模 | 未融合（s/step） | 融合后（s/step） | 加速比 |
|---|---|---|---|
| 1B 参数 | 0.85 | 0.62 | 1.37x |
| 7B 参数 | 2.34 | 1.58 | 1.48x |
| 13B 参数 | 4.12 | 2.67 | 1.54x |
| 70B 参数 | 18.76 | 11.23 | 1.67x |
看到了吧？模型越大，收益越明显。原因很简单：大模型的计算密度高，更容易把通信藏进去。
3.4 核心逻辑流程图
下面这张图，是我自己画的。它展示了通信算子融合的核心逻辑：
这张图很直观。上面是串行，下面是流水线。你会发现，融合后通信完全被计算“盖住”了。嗯，这就是我们想要的效果。
3.5 避坑指南
融合虽好，但也不是银弹。我总结几个常见坑：
我曾经踩过的坑：
- 过度融合：把所有通信都塞进计算里，结果显存爆了。记住，融合是有代价的——需要额外的 buffer 来存中间结果。
- 忽略带宽争抢：计算和通信同时跑，如果共享同一根 PCIe 总线，反而会互相拖慢。我建议用 NVLink 或专门的通信通道。
- 桶大小调不好：太小了通信次数多，太大了等待时间长。我一般从 1MB 起步，用 profiler 看实际效果再调。
一个小技巧：用 NVIDIA Nsight Systems 看一下时间轴。如果通信和计算没有重叠，说明融合没生效。这时候检查一下是不是用了同步的通信 API。
3.6 总结
通信算子融合，说白了就是让 GPU 别闲着。计算和通信并行跑，总时间就降下来了。
我个人觉得，这是多卡训练里性价比最高的优化手段之一。不需要改模型结构，不需要换硬件，纯软件层面就能拿到 30%~50% 的加速。
当然，具体怎么做，还得看你的场景。数据并行用 AllReduce 融合，模型并行用 AllGather 融合，混合并行就组合着来。嗯，多试试，总能找到最优解。
想要更多资料，联系可可