NIPS 2018 | GradiVeQ: Vector Quantization for Bandwidth-Efficient Gradient Aggregation in Distributed CNN Training

数据并行可以提升卷积神经网络的训练速度,但同样会在梯度聚合时产生严重的通信开销。使用量化技术可以对梯度进行压缩,从而降低通信开销。但是,当这些量化技术与去中心化的聚合协议(如Ring All-Reduce,RAR)一起使用时,它们的性能并不好,这是因为它们无法直接聚合压缩后的梯度。本文发现了CNN卷积层梯度间的强线性相关性,并提出了一个名为GradiVeQ的梯度向量量化方法,以通过主成分分析来进行梯度降维。GradiVeQ可以直接聚合压缩后的梯度,因此我们可以通过并行GradiVeQ梯度压缩和RAR通信技术来构建一个分布式机器学习系统。实验表明,使用GradiVeQ训练CNN比使用原始的RAR通信方式训练快5倍,减少了50%的端到端训练时间,并且没有明显的精确度损失。GradiVeQ与其他量化技术(如QSGD)兼容,并且在相同压缩率下获得了更好的加速比。

1. Introduction

当前,因为数据量和模型都比较大,所以CNN一般都使用分布式训练方法,比如使用点对点通信的并行随机梯度下降(SGD)。其中每个节点计算本地梯度,然后将梯度广播给所有其他节点,再把梯度聚合,最终完成对参数的更新。然而,梯度聚合的方式会严重影响通信开销,使得通信成为训练过程中的瓶颈。

有两种方法可以提升梯度聚合的效率:梯度压缩并行聚合。梯度压缩主要是减少表示梯度的比特位,常见的方法有标量量化(有损)和稀疏编码(无损)。另一方面,并行聚合旨在通过从参数服务器的集中式聚合转换为分布式聚合(例如RAR)来最小化通信拥塞。
Figure-1

如图1所示,RAR将所有节点放到一个逻辑上的环路中,然后每个节点同时向下一个节点发送梯度向量的不同段。每个节点将它接收到的段加到本地梯度向量的相同段上,然后将结果发送到环路中的下一个节点。一旦某个段通过所有节点循环,它就变成一个完全聚合的梯度向量。然后,下一轮循环会把该段广播到所有节点上。

由于梯度压缩和并行聚合这两种方法具有很高的互补性,所以我们可以尝试将这两种方法结合起来。一个关键的问题是,压缩后的梯度无法直接进行聚合,除非解压缩它们。直接聚合压缩后的梯度还有可能导致算术溢出和梯度稠密化等问题。

Figure-2

在RAR的每一步,每个节点都会先解压下载的梯度段,然后再把它们与自己对应的梯度段相加。然后,每个节点压缩相加的结果并传输给环路中的下一个节点。显然,如图2(a)所示,下载与压缩梯度的过程不能并行化。但是,同样的梯度会再每个节点上被重复地压缩/解压缩。

为了将梯度压缩和并行聚合结合在一起,压缩函数和梯度聚合操作必须是可变换的。令$Q(·)$是压缩函数,那么它必须满足:

其中$N$是节点的数量,$\mathcal{g}_n$是节点n上的梯度向量。等式$\ref{1}$的左边表示直接对压缩梯度进行聚合。这一类压缩函数允许并行压缩和RAR通信,因此压缩时间和通信时间就可以重叠,如图2(b)所示。等式$\ref{1}$的右边表示解压缩函数$Q^{-1}(·)$只需要执行一次——在所有压缩梯度聚合之后。

Intuition and Motivational Data
GradiVeQ的核心是使用主成分分析(PCA)来探究梯度间的线性相关性并对梯度进行降维。GradiVeQ的主要思路来自于我们对CNN中梯度间的强线性相关性的观察结果。

Figure-3

  • Linear correlation 我们将梯度拉直成向量以便观察梯度间的线性相关性。如图3(a)所示,我们将同一卷积层F个卷积核相同位置的梯度放在一起。这F个梯度元素是由相同的输入数据得到的,因此这F个梯度之间会有强烈的线性相关性。这种线性相关性允许我们使用PCA对每个梯度向量计算一个线性压缩矩阵。
  • Spatial domain consistency 对于一个$(H,W,D,F)$的卷积层,较大的梯度向量可以分成许多粒度为$F\times D$的切片,这些切片之间也有线性相关性。通过使用一个切片的压缩矩阵来压缩其他切片的低压缩损失可以证明这种相关性。
  • Time domain invariance 最后,我们还注意到线性相关性在迭代过程中会缓慢变化,这可能是由于合理的学习速率下稳定的梯度方向与步长。因此,我们可以在未压缩聚合的训练迭代期间投入一小部分计算资源来执行PCA,并使用生成的压缩矩阵进行后续迭代的计算。

Figure-4

2. Description of GradiVeQ

考虑使用含有N个工作节点的分布式系统在某一数据集上训练卷积层含有M个参数的CNN,w表示他们的权重。在第t此迭代时,每个节点n使用数据集的不同子集进行训练,得到一个长为M的向量$\mathcal{g}_n[t]$。这些梯度向量通过下式进行聚合:

然后用$w[t+1] = w[t] -\eta[t]\mathcal{g}[t]$来更新每个节点的权重,其中$\eta[t]$是学习率。

在训练阶段,每个卷积层使用相同的卷积核在训练数据上滑动。这就启发我们通过将梯度$\mathcal{g}_n[t]$表示成向量形式以观察相邻梯度间的线性相关性。

GradiVeQ的目标是对每K个相邻的梯度进行单独的压缩,其中K是切片的大小。令$\mathcal{g}_{n,m}[t]$是第m($m\in[0,M/K-1]$)个梯度切片,GradiVeQ会使用函数Q按照以下方式对该切片进行压缩:

这里U是一个kd的矩阵且k小于d,$\mu_m$是一个K维的白化向量。压缩之后,每个切片的维度变为d,也就是说压缩率为K/d。

不同节点压缩后的梯度切片可以直接按照下式聚合:

这表明GradiVeQ的压缩与聚合操作是可变换的。之后,我们就可以通过下式来对$\mathcal{g}’[t]$进行解压缩:

Figure-5

GradiVeQ使用PCA来计算压缩矩阵和白化向量。对于每个切片m,所有节点周期性地使用Lt次迭代进行未压缩梯度的聚合。这之后,每个节点会有Lt个切片m的样本,即$\mathcal{g}_m[t],\cdots,\mathcal{g}_m[t+L_t-1]$。白化向量是这Lt个样本的平均值。每个节点会使用这Lt个样本计算K个梯度的协方差矩阵Cm,利用奇异值分解SVD来计算Cm的特征矩阵Um和特征向量sm。压缩矩阵Udm是特征矩阵Um的前d列,对应于特征向量sm中的前d个最重要的特征值。接着,Udm和mum在节点n的$L-L_t$次迭代中被用来压缩$\mathcal{g}_{n,m}[t+L_t],\cdots,\mathcal{g}_{n,m}[t+L-1]$

Algorithm-1

由于良好的可变性,GradiVeQ梯度压缩可以在RAR上被并行化,如算法1所示。因此,压缩时间可以和通信时间重合。我们将N个节点放在一个逻辑环拓扑中,每个节点只能给它的直接后继发送数据。我们将梯度向量分割成N个相等的段,每段都包含一些梯度切片。聚合操作包括两轮:压缩轮与解压缩轮。初始化时,每个节点n都只压缩第n个梯度段中的梯度切片,并将它们发送到后继节点。然后,在压缩轮的每一步,所有的节点同时做以下两件事情:1)从前驱节点下载压缩后的梯度;2)压缩第n段的梯度切片。一旦这两件事情都做完了,节点就会将两个压缩后的梯度段相加并将结果发送给后继节点。N-1步之后,压缩轮执行完毕,每个节点都拥有完整的聚合后的梯度。接下来,在解压缩轮的每一步,每个节点同时做以下两件事:1)从前驱节点下载新的压缩梯度段;2)解压最新下载的压缩梯度段。在解压缩完成后,必须保留压缩梯度段以供后续节点下载。原始的RAR算法是算法1的特例,只不过没有压缩和解压缩操作。

3. Implementation Details of GradiVeQ

3.1. System Overview

在训练初始阶段,使用一个短暂的预热阶段来稳固模型。然后,对每个梯度切片进行Lt次的无压缩聚合以及Lc次的压缩聚合,进行无压缩聚合的主要原因是PCA需要一些原始数据来计算压缩矩阵Ud。同时,我们通过周期性地更新Ud的值来减少压缩误差。根据一个型的梯度向量计算压缩矩阵Udm需要耗费大量的计算时间,为了减少PCA的计算量,GradiVeQ利用了梯度相关性的空域一致性。对同一卷积层中连续的s(可称之为压缩重用因子)个切片,GradiVeQ使用PCA计算一次压缩矩阵Ud,然后使用这个Ud来压缩剩余的s-1个切片。我们注意到,尽管计算Ud时不需要其他的s-1个切片中的值,但是它们仍然像第一个切片一样,以一种无压缩方式进行聚合,因此所有切片的参数可以是相同的。

Figure-6

为了提高没有使用GradiVeQ的前Lt次迭代的带宽,本文会使用标量量化技术如QSGD来进行这些次迭代。

3.2. Parameter Selection

Slice size K
对于一个卷积层来说,K一般是FD的倍数,其中F是卷积核的数量,D是深度。此外,可以通过控制K的选择来平衡压缩率和PCA计算复杂性。增加K会增加计算和内存开销,但是可以捕捉到更强的线性相关性。

Compressor resue factor s
一个较大的s意味着PCA根据一个切片计算得到的压缩矩阵会被永东到许多连续的切片上,与此同时压缩损失会增大。通过实验我们发现,当s=正无穷时,也就是说每个卷积层只有一个压缩矩阵时,精度相比于benchmark仅仅下降了不到1%。

Compressor dimension d

d的值取决于我们可以忍受多大的压缩损失,使用特征向量s可以投影该损失。令$0\leq \lambda < 1$为损失阈值,最小的d需要满足下面的不等式:

此时的Ud将会保证压缩损失至多为lambda。
Number of iterations with uncompressed aggregation $L_t$
这个值可以被当作一个超参数,也可以由无压缩的RAR决定:即在收集的梯度样本切片上执行PCA,以找出需要多少样本才能获得稳定的d值。本文按照第二种方式处理,最终大概为100。
Number of compressed iterations $L_c$
该值也可以视为超参数,或在运行时让节点根据本地压缩损失来确定,定义为:

当大多数节点的压缩损失都比较大时,我们就停止压缩,然后进行无压缩迭代以计算新的压缩矩阵。在后面的实验中,我们使用第二种方式确定Lc的值,约为400。
Computation Complexity
GradiVeQ在每个梯度切片上有两个主要的操作:(a)每Lt+Lc次迭代会有一次在聚合切片的Lt个样本上的SVD计算(b)每个节点每次迭代会有的两次低维矩阵乘法。梯度拉直和切片允许在一个切片上计算的压缩矩阵可以重复地用在同一卷积层地不同切片上,这就大大降低了SVD的计算时间。另一方面,操作(b)是低维矩阵相乘,它需要的时间小于RAR通信时间。因此,GradiVeQ不会增加训练时间。

4. Experiments

实验环境配置

  • 网络模型:ResNet-32
  • 数据集:CIFAR-100
  • 节点数量:6
  • 本地CPU集群:每个节点包括2个Ten-Core Intel Xeon 处理器(E5-2630v4)
  • 云端GPU集群:每个节点包括1块NVIDIA Tesla K80 GPU

参数设置

  • 损失阈值为 $\lambda = 0.01$
  • 对每个卷积层,令$s = \infty$

Figure-7

4.1. Model Convergence

本文比较了无压缩算法、4-bit-QSGD和GradiVeQ在给定精度下的收敛速度,3种算法均在45000次迭代后收敛。如图7(a)所示,两种压缩方法都大大降低了模型收敛的墙钟时间。在CPU环境下,无压缩算法的模型收敛时间为135000秒,4-bit-QSGD需要90000秒才能收敛,而GradiVeQ只需要76000秒就可以收敛。而在GPU环境下,无压缩算法、4-bit-QSGD和GradiVeQ的收敛时间分别为75000秒、30000秒和24000秒。在测试集准确率方面,无压缩算法的top-1准确率为0.676,4-bit-QSGD的top-1准确率为0.667,GradiVeQ的top-1准确率为0.666。可以看到,GradiVeQ以略微降低准确率为代价,大大地加速了网络模型的训练过程。

4.2. End-to-End Training Time Reduction Breakdown

端到端训练时间由计算时间和梯度聚合时间组成。计算时间包括反向传播算法计算梯度的时间以及更新模型参数的时间。梯度聚合时间主要是聚合所有工作节点上的梯度向量需要花费的时间。由于GradiVeQ和4-bit-QSGD没有对计算过程进行处理,所以它们和无压缩的训练拥有相同的计算时间。在CPU环境下,每进行$500(L_t+L_c)$次迭代的计算时间大约为650秒。对于聚合时间来说,无压缩训练每500次迭代需要850秒,大概占端到端训练时间的60%。GradiVeQ可以将聚合时间降低到162秒每500次迭代,比无压缩训练少了5.25倍,同时减少了46%的端到端训练时间。虽然4-bit-QSGD拥有与GradiVeQ相同的压缩率,但是它的聚合时间大概是GradiVeQ的2倍。在GPU环境下,无压缩训练的计算时间可以缩短6倍,因此聚合时间成为了瓶颈,占到了总训练时间的88%。通过降低聚合时间,GradiVeQ的端到端训练时间比无压缩训练减少了4倍,比4-bit-QSGD减少了1.4倍。