利用共享内存实现比NCCL更快的集合通信

01 简介

从 2080Ti 这一代显卡开始，所有的民用游戏卡都取消了 P2P copy，导致训练速度显著的变慢。针对这种情况下的单机多卡训练，MegEngine中实现了更快的集合通信算法，对多个不同的网络训练相对于NCCL有3%到10%的加速效果

MegEngine v1.5 版本，可以手动切换集合通信后端为 shm（默认是 nccl），只需要改一个参数。（由于 shm 模式对 CPU 有额外的占用，且只有在特定卡下才能提高效率，因此并没有默认打开）

gm = GradManager()
gm.attach(model.parameters(), callbacks=[dist.make_allreduce_cb("sum", backend="shm")])
# 目前只实现了单机版本，多机暂不支持

02 背景

在大规模训练中，数据并行是最简单最常见的训练方式，每张卡运行完全一样的网络结构，然后加上参数同步就可以了。

对于数据并行的参数同步，目前有两种常用的方法，Parameter Server 和 Gradient AllReduce：

Parameter Server 方案需要额外机器作为参数服务器来更新参数，而且中心式的通讯方式对带宽的压力很大，增加训练机器的同时通信量也线性增加；
而 AllReduce 方案只是参与训练的机器之间互相同步参数，不需要额外的机器，可扩展性好。

MegEngine 目前也是使用 AllReduce 方案作为数据并行的参数同步方案。而在AllReduce方案中，大家目前常用的是 NCCL，Nvidia 自家写的 GPU 集合通讯库，通信效率很高。

看到这里，可以得到一个结论，数据并行的情况，用 NCCL 通讯库能达到不错的效果。

到这里就结束了？当然不是，在 2080Ti 8卡训练的情况下，在多个网络下，我们相对 NCCL 有 3% 到 10% 的性能提升。（以 2080Ti 为例子是因为游戏卡不支持 P2P 通信，相对来说通信较慢，通信时间长，节省通信时间能获得的收益较大）

这是怎么做到的呢，我们一步一步来分析（以下数据都是用 megengine.utils.profiler 导出，相关文档在 profiler文档——https://megengine.org.cn/doc/...）。

通常我们是在 backward 阶段同时做 gradient 的同步，我们来看单卡的 backward耗时，只有 164ms：

图 1（单卡 ResNet50 训练）

再来看 8 卡训练时 backward 的耗时，增长到了 203ms，比单卡的情况下多了 39ms：

图 2（使用 NCCL 做 8卡 ResNet50 训练）

确实，backward 时间变长了不少，可是为什么？我们有没有可能消除它？

1）为什么

一句话：NCCL AllReduce 占用了 cuda 计算资源，所以计算变慢。

具体原因是 NCCL AllReduce 对应的 cuda kernel 需要既做通信又做计算，所以占用了计算对应需要的计算资源，但是大部分时间都花在了通信上，导致计算资源的利用率不够高。

2）能不能消除它

当然是可以的，cuda 计算和通信是可以并行的，让计算和通信运行在两个不同的 cuda stream 上，就可以并行起来，这一点在 cuda 开发者文档中有提到https://docs.nvidia.com/cuda/...。

3）实际测试计算和通信并行的例子

stream0 进行矩阵乘法运算，stream1 进行拷贝，前后矩阵乘法的速度没有受到影响：

图3 计算通信并行，stream0 做矩阵乘法，stream1 做 d2d copy （左右滑动）

03 实现思路

这里先介绍一下 AllReduce 的两种实现方法。

1）ReduceScatter + AllGather

图4（ReduceScatter算子图解）

图5（ AllGather 算子图解）

Ring AllReduce 就是用的 ReduceScatter + AllGather 的模式：首先第一轮通信在各个节点上计算出部分和，然后第二轮通信把部分和聚集一下得到最终结果。

2）Reduce + Broadcast

图6（ Reduce 算子图解）

图7（Broadcast算子图解）

Reduce + Broadcast 的方式像 Parameter Server，会先在一个节点计算出完整的累加和，然后再广播到各个节点上。

新的算法采用的是 Reduce + Broadcast 的方法，首先将数据全部拷贝到 cpu（使用 Shared Memory），然后由 cpu 进行累加，再拷贝回 gpu。

由于直接拷贝累加没有把计算和通信 overlap 起来，我们还采用了类似 Ring AllReduce 的分块策略，让通信和计算充分 overlap，如下图所示。

图8 Shared Memory AllReduce 中 Reduce 阶段图解

因为用到了 Shared Memory，所以把这个后端简称为 SHM。

04 实现效果

1）算子性能

相对 NCCL 来说，SHM 的延迟稍微高了一些，带宽低了一些，数据大一些的情况可以达到 NCCL 的 90% 左右的性能。

图9（ SHM 和 NCCL 算子性能对比）

SHM 性能需要在数据包大的情况发挥，和 ParamPack 策略搭配能最大发挥 SHM 的作用（在 MegEngine 中 distributed.make\_allreduce\_cb 中使用了 ParamPack 策略，默认打包大小为 10M）。ParamPack 策略是指将参数对应的梯度打包进行发送，减少小包发送，以减少通信延迟增加带宽利用率。

2）实际训练效果

继续使用 ResNet50 8卡训练的例子，SHM 与 NCCL 相比，backward 时间快了近 30ms（203ms->174ms）。