文@小P维修员
前言
深度学习算法研究如火如荼。为了更好的算法表现，深度学习模型变得更大、更复杂，计算复杂度也越来越高。
这里的计算复杂度主要是指深度学习模型的计算量大小。日益增加的算力需求将模型的训练推向了多 GPU 训练甚至更大规模的分布式训练。而分布式训练的加速效果和可扩展性，很大程度上受限于节点之间的通信。
正文分为四个章节，分别介绍：
- 在深度学习模型训练中常见的通信方式
- 深度学习框架中通信的实现
- 通信的实用优化策略
- 在模型搭建中进行通信优化：分布式数据并行 DDP
模型训练中的通信
分布式训练通常在计算集群上进行，集群的每个节点（计算机）分别执行一部分计算。不同节点的计算之间有数据依赖和共享，需要将数据在不同节点间传输，这就是通信。
分布式的通信一般有两大类：
- 集合通信： 在一组节点（进程）内进行通信。
- 点对点通信： 在两个节点（进程）之间进行通信。
深度学习训练中用到的大多是集合通信。
以最常用的数据并行为例，在一个迭代步骤内，每个 GPU 先对各自的数据进行前向和后向的计算得出模型梯度，再将各个 GPU 的模型梯度进行平均，最后将梯度更新到各自的模型参数上。
梯度平均的常见实现方式是先归约求和，再除以 GPU 数量做算术平均。这里的梯度求和是跨 GPU 的，它需要用到一个集合通信操作 Allreduce。
Allreduce 将数组的每一个元素分别跨 GPU 归约，这里的归约可以是求和、求乘积、求最大值等。
图中的 MPI_SUM 是 MPI 接口中定义的求和运算的标识，后文会详细介绍 MPI 接口。
Allreduce 是深度神经网络分布式训练中最常见的通信操作，本文主要对 Allreduce 在训练过程中的性能进行探讨。
框架中的通信实现
集群的网络硬件多种多样，可以是 Ethernet、InfiniBand 等，深度学习框架通常不直接操作硬件，而是使用通信库。
通信库屏蔽了底层硬件细节，提供了统一封装的通信接口。
下面是两种最常用的通信库（接口）：
Message Passing Interface (MPI)
MPI 信息传递接口，是一个用于编写并行计算程序的编程接口。它提供了丰富全面的通信功能。
MPI 常用于在计算集群、超算上编写程序，比如很多传统科学计算的并行程序。MPI 接口的兼容性好，通信功能丰富，在深度学习框架中主要用于 CPU 数据的通信。
MPI 是一个开放接口，有多种实现的库，一种广泛使用的开源实现是 Open MPI。一些硬件厂商也提供针对硬件优化的实现。
NVIDIA Collective Communication Library (NCCL)
NCCL 英伟达集合通信库，是一个专用于多个 GPU 乃至多个节点间通信的实现。它专为英伟达的计算卡和网络优化，能带来更低的延迟和更高的带宽。
NCCL 也提供了较丰富的通信功能，接口形式上与 MPI 相似，可满足大多数深度学习任务的通信需求。它在深度学习框架中专用于 GPU 数据的通信。
目前成熟的深度学习框架都集成了包括上述两者的多种通信库。SenseParrots 亦是如此，并且针对这两种最常用的库进行了优化。
在深度学习训练中，大部分的计算和数据都在 GPU 上，所以更多地使用 GPU 数据通信，即使用 NCCL。
本文主要基于 NCCL 来进行通信性能分析和优化介绍，大部分结论和方法在 MPI 上也适用。
实用通信优化策略
通信融合
在训练中，大部分通信是对参数的梯度进行平均。最朴素的做法是对模型的每个参数梯度都调用一次 Allreduce 进行通信。
然而，对于大多数通信库的实现，通信调用本身具有一定开销，这样就会导致通信时间并不是与通信量成严格的正比关系。
比如，分别做两次 4MiB 的通信会比做一次 8MiB 的通信耗时更多，其中会增加一次通信调用开销。在通信量较小时，调用本身带来的开销就会很明显。
所以，通信量一定时，通信次数越少，效率越高。
一般的，每个参数的梯度是互相独立的，每个梯度唯一依赖的是产生它的反向计算。所以可以将整个模型的梯度统一用一次 Allreduce 进行通信。
具体的做法是：
- 将模型的梯度按顺序拷贝到一个连续的 buffer；
- 对这个 buffer 做 Allreduce；
- 按顺序将 buffer 中的数据拷贝回模型的梯度变量。
上面的步骤假设梯度的数据类型是相同的。如果有多种数据类型，可以先对梯度按照数据类型分类，再对每种类型执行上面的步骤。
通信融合能提高梯度通信的效率，在模型参数梯度的数据量较小而个数较多的情况下效果显著。
通信与计算并行
通信融合优化在模型训练上有很好的优化效果，但还是不够的。
当模型梯度数据量很大，即使所有的梯度只用一次通信，通信本身的时间也会很长，占训练时间的比重很大。通信本身的时间难以再压缩，总的训练时间中还有很大一部分是反向计算的时间。
在朴素的做法中，总是先进行反向计算，再进行梯度平均，计算时间和通信时间是累加的。类比于网页一边下载一边渲染的思路，如果通信和计算也并行进行，则可以节省一部分时间。
对于 NCCL，其通信主要使用网络硬件资源，只会使用少量 GPU 计算资源。可以认为通信和反向计算的资源竞争很小。关键在于通信和计算之间存在数据依赖，通信使用的梯度数据是反向计算产生的，所以一个梯度的通信必须在其对应的反向计算之后发生。因此，并行对通信的顺序有一定要求，最直接的做法是在每一个梯度被反向计算出来之后，立即对其进行通信。
注意图中通信的顺序是反过来的， 是因为要跟反向计算保持一致。反向计算的顺序是由自动求导机制决定的。
事实上，反向计算和通信的时间都没有减少（实际可能有少量增加），但是因为两者在时序上的并行重叠（图中 Overlap 部分），使得一次训练的总时间减少了。
并行加上融合
即使通信与计算有了并行重叠，总的时间中还是有很大部分是单独的通信耗时。这是因为通信时间比计算时间更长，在计算结束后有较长的剩余通信耗时。
在【通信融合】一节中，介绍了使用通信融合的策略来缩短通信的总时间。以上两种优化策略在原理上是不冲突的，所以很容易想到将两种策略结合起来使用。
思路很直观，就是在并行通信的基础上，将若干个梯度的数据合并在一起，一并通信。
需要注意的是，不能简单地将所有通信融合为一个。原因是全部梯度的通信依赖于全部梯度的计算的完成，融合的通信实际发生在全部反向计算完成后，得不到并行重叠的效果。
所以上文表述的“若干”，是一个影响并行重叠效果的关键数值， 若太多，则要等待的计算太多，通信开始太晚；若太少，则通信的次数太多，通信总开销大。
一般定义通信融合的数据量阈值为一个可调的参数。可以人为选定一个合适的值，也可以根据模型实际的计算和通信量自适应调整。
从图中可以看到，通信次数变少，重叠部分之后的剩余通信的耗时会减少，计算和通信的总时间也相应减少。
Distributed Data Parallel (DDP)
DDP 是 PyTorch 提供的用于模型数据并行训练的接口，它不仅实现了常用的梯度同步，还封装了一整套分布式并行训练的功能。
DDP 的用法例子在基于 Python API 搭建深度学习模型中数据并行部分有介绍，只需稍微修改代码，就能将一个模型改造成分布式并行训练的模型，这里不再赘述。
SenseParrots 提供了兼容的 DDP 接口，并在内部实现了上述的通信优化策略，使得现有的模型可以轻易地获得分布式训练的加速。
结语
感谢阅读，关于深度学习中的通信优化欢迎大家一起留言讨论哦~
更期待有相同兴趣的同学加入我们，一同探索解决 AI 框架面临的问题和挑战！
大家如果有感兴趣的技术内容和难点欢迎随时指出，可以直接在评论区给我们留言~
关注 公众号「SenseParrots」，获取人工智能框架前沿业界动态与技术思考。