TOWARDS SCALABLE DISTRIBUTED TRAINING OF DEEP LEARNING ON PUBLIC CLOUD CLUSTERS（论文笔记）

问题

top-k稀疏化是一种减少分布式机器学习通讯量的一种梯度稀疏算法，它能够只发送梯度的一小部分，对于收敛性有很小的影响。但是，使用top-k稀疏化实现真正的性能提升并非易事，主要原因有两个: (1) top-k的选择过程在GPU上效率不高，(2) 稀疏通信通常需要all-gather操作，其在公共GPU集群上性能很差。

此外，网络文件系统(NFS)的I/O也应该仔细设计，以在大规模训练中实现更高的吞吐量，因为在每个训练迭代中都需要大量的数据访问。

这篇文章中，提出了一种基于top-k稀疏化的高效通信库，首先提出了一种对GPU架构友好的近似top-k算子，然后提出了一种层次化的top-k通信，以更好地利用连接带宽资源。为了进一步提高系统的可扩展性，提出了一种简单有效的I/O读取机制，支持多级数据缓存和并行梯度处理。

设计

MSTopK: an Approximate Top-k Operator

首先介绍典型的top-k。对于任何输入向量$\mathbf{x}\in {\mathbb{R}}^d$，

上标(i)表示第i个元素，thres是$|\mathbf{x}|$的第k个最大的值。在GPU上实现这种算法，效率不高。

而MSTopK算法的核心思想就是使用二分法求出近似值，这里的二分法不需要排序，而是根据平均值加上最大值与平均值的差值的一定比例来计算边界。最后返回估计的k个最大的值以及它们的索引。

Hierarchical Top-k Communication

由于节点内部GPU之间的通信比节点之间的通信快得多，本文提出了一种层次化的通信方式，分为以下四步：

1. 在节点内部，所有的GPU执行reduce-scatter操作（类似于ring-based all-reduce中的reduce-scatter），完成后每个GPU有1/n的聚合后的数据（n是GPU的数量）
2. 对于聚合后的数据，每个GPU进行top-k选择
3. 在节点之间，调用n个通信流。在$j^{th}$个通信流中，所有节点的$j^{th}$个GPU对k个最大的值和其索引执行all-gather操作。当GPU得到其他机器的GPU的梯度后，根据索引进行聚合。
4. 在节点内部，所有的节点执行all-gather操作，得到聚合后的梯度。

因为节点之间的通信比节点内的通信速度慢得多，所以时间耗费最长的是第三步的all-gather操作。

数据缓存

一方面，网络文件系统的读取性能很差，另一方面，训练数据需要被预处理。为了提高数据读取和预处理的性能，提出了以下方法：

在第一个训练回合，数据应该从NFS（网络文件系统）读取到本地机器的本地文件系统缓存。数据应该被进一步解码为可以作为深度模型的输入的格式，这可能会消耗许多CPU资源。我们使用键值存储将预处理的数据进一步缓存到内存中，其中键是样本索引，值是预处理的数据。这样一来，内存缓存可以大大减少训练过程中的I/O时间。为了减少内存消耗，完整的数据集被分割成多个部分，分别存储在多个节点上。从第二个训练回合开始，所有的数据都被存储在内存缓存中。通过数据读取和GPU计算之间的流水线操作，从内存缓存中读取数据的时间成本几乎可以完全被GPU计算所覆盖。总之，一方面，启用本地文件系统缓存来存储数据，使得调整超参数的多次运行在数据读取上更加高效。另一方面，使用内存缓存使得每次运行时从第二个回合读取数据的速度很快。

并行张量处理（PTO）

在梯度聚合后，所有的GPU都有相同的梯度。在更新模型参数之前，这些梯度需要进一步的处理（比如优化器的处理）。这个时间可能会比较慢，此时就可以使用一种并行的处理方式。假如有P个worker，那么就把梯度分成P份，每个GPU处理一份，然后对处理后的部分调用all-gather，将其分发到其他worker上。

在实际应用中，如果all-gather操作的时间花费比计算时间的减少更小的化，PTO就可以加速梯度的处理了。