背景
在分布式机器学习中,同步的数据并行方法是最常用的一种方法。在此情形下,大模型带来的网络通信开销通常会成为瓶颈。下图展示了一个简单的场景。
在此场景下一共有三个计算节点和一个参数服务器。现实中参数服务器可以在任意一个计算节点上。对于同步训练方法,每一轮训练各个节点之间都需要进行模型的同步,以保持模型数据一致,从而让模型可以快速收敛。计算节点间通过网络连接进行通信,可能在一个局域网中,也可能通过互联网连接,这取决于具体的使用场景。对于一些大的模型来说,即使是在局域网范围内,网络的数据传输延迟也是不能接受的。因此,此时就需要降低每一轮的数据传输量,这就需要用到梯度压缩算法。
当前已经有许多方法用于梯度压缩。
梯度压缩算法
符号压缩法
符号压缩法只仅使用一比特来表示一个浮点数值,比如使用1表示正数,使用零表示负数。通常情况下,对于一个矩阵数据来说,还可以传递一个矩阵正数以及非正数的均值用于降低在恢复梯度时的偏差。一些变体可能还会对每个浮点数值使用额外的1比特用于指示浮点数是否为0。
量化方法
量化方法也包含很多,常用的方法比如将32位浮点数转化为16/8位浮/定点数。也有一些方法会将浮点数缩放到0~255(8比特)/0~15(4比特)之间;还有一些方法会只取浮点数的指数部分来进行压缩,比如0.001=1e-3就用-3替代,300=3e2用2替代。
降秩
此方法会对梯度矩阵进行矩阵分解,比如SVD分解。根据具体精度要求,可以仅传输最大的几个奇异值以及相应的左右奇异向量即可。该法有一个缺点就是SVD分解耗时极大。本文的方法类似此方法,但是对于耗时的SVD分解的部分进行了优化。
还有一些方法,这里不一一介绍。这些梯度压缩算法存在的一些问题是:
- 可能存在精度损失。显然这些方法都是有损的,多多少少会对精度产生一些影响
- 聚合时需要对于每个计算服务器发送过来的梯度信息进行解压缩后才能进行聚合(求均值)。这就增加了时间损耗(有一些方法不需要,这里不在赘述。本文提出的方法无需先解压即可进行聚合(求和与解压缩顺序可互换),这就是“线形”性质)
本文对于现有的这些问题,结合之前一些方法的优点,提出了一种改进方法,此方法:
- 引入线性:优点是可以直接使用现有分布式机器学习库中的提供的一些方法比如 gather/all reduce 进行操作优化,而无需进行先解压缩再进行聚合、发送。这简化了操作,节省了时间
- 使用错误反馈机制:本文的方法也是对矩阵进行了分解,将重要的分量进行传输聚合。此种方法仍然是有损压缩方法,即无法完全恢复梯度原来的值,因此必然会引入误差。错误反馈机制,会将每一轮的误差保存下来,加到下一轮的梯度中去,此方法在一些实验中被说明在模型收敛性以及准确率保证方面有优势
方法
本文为了避免SVD计算的复杂性,使用了另外一种高效的方法来计算一个近似的矩阵。该方法就是1975年提出的 subspace iteration。
矩阵压缩方法
对于一个深度神经网络模型来讲,一面的参数基本都是由张量构成的,另外还包括一些偏置(bias)量。对于偏置量,在本文的方法中没有进行压缩,毕竟bias的数据量相对于张量来说数据量小很多。对于张量来说,所有的张量都被转化为二维的,也就是n行m列的张量以方便进行矩阵分解。
具体算法如下(python 伪代码):
import numpy as np
def compress_and_aggregate(grad_matrix, previous_q):
"""压缩聚合算法。
grad_matrix: 为 n * m 的梯度矩阵
prefious_q: 为前一轮的 q 矩阵,维度为 m * r,r为需要取的特征值的个数,q矩阵最初用标准
正态分布进行初始化
"""
# 0. 计算 p (n * m X m * r --> n * r)
# 该操作应该由工作节点完成,工作节点将p而不是grad_matrix
# 发送给聚合节点, grad_matrix 自身并没有被发送给聚合节点
# 这里为了和原文的伪代码保持一致,所以直接写在
# 这个函数里
p = np.matmul(grad_matrix, previous_q)
# 1. 此处需要对所有节点的p求和然后发送给所有节点
p = all_reduce(p)
# 2. 使用 Gram-Schmidt 进行正交化
p = orthogonalize(p)
# 3. 计算 q (m * n X n * r = m * r)
q = np.matmul(grad_matrix.T, p)
# 4. 此处需要对多有节点的q求均值然后发送给所有节点
q = all_reduce_mean(q)
return p, q
def decompress(p, q):
"""解压缩算法"""
return np.matmul(p, q.T)
笔注:上述步骤1在论文中为 all reduce mean,根据我的理解,以及对其代码实现的考察,此处应为求和操作,下面步骤4中才需要进行求均值。
从代码中可以看出,需要通信的地方为第1、4处,其中1处p为$n r$维矩阵,4处q为$m r$维矩阵,由于r较小,相比于传递原来的$m n$维矩阵来说,传递$(m+n)r$维矩阵的数据量要小得多。
Gram-Schemidt正交化代码可以参考powersgd的实现[2]。
算法的基本思想就是使用一个高效的算法得到数据量较小的p和q,并且$pq^T$可以很好的近似原来的数据量较大的矩阵$grad\_matrix$。
文中指出,使用前一轮的q值来做热启动(warm-start),而不是每次都使用新的q,可以让上述算法对于矩阵的分解的方法达到与SVD分解相当的效果,相比之下,上述方法的时间消耗远小于SVD分解。
高效聚合
上述方法的带来额外的一个好处是方便并行计算,而无需重复进行压缩解压缩。之前的一些方法,我们需要将压缩后的梯度矩阵恢复为原始的32位浮点数矩阵之后才方便求和或均值操作,而后在将聚合之后的梯度矩阵再次进行压缩,再发给所有计算节点。而本文提出的方法直接可以对压缩后的数据(矩阵)直接进行相加操作,聚合操作和解压操作是可以相互交换的,这样更便于使用并行计算进行优化。如下图所示,我们很容易进行下图所示的reduce操作,该操作可以使用多个worker并行。
错误反馈
本文的SGD算法融入了动量(momentum)和错误反馈,算法python伪代码如下:
import numpy as np
# 工作节点个数,假设为100
W = 100
# 矩阵大,假设大小为 10000维
d = 10000
# 错误反馈缓存
error_feedback = np.zeros((W, d))
qs = np.random.normal(0, 1, (W, m*r))
def power_sgd(model_parameter, learning_rate, momentum_decay, momentum=0):
"""
model_parameter: 模型参数 d 维
learning_rate: 学习率
momentum_rate: 动量衰减率
init_momentum: 初始动量
"""
# 每一轮中,所有worker执行下列操作(并行)
for w in range(W):
# 0. 计算当前轮梯度
grad = compute_grad(w)
# 1. 加上错误反馈
grad_with_feedback = grad + error_feedback[w]
# 2. 计算误差
q = qs[w]
# 2.1 计算近似值
grad_approximate = np.matmul(
grad_with_feedback,
np.matmul(q, q.T))
# 2.2 更新误差矩阵
error_feedback[w] = grad_with_feedback - approximate
# 3. 使用前面的算法进行压缩+聚合
# 注意这里调用只是形式,grad_with_feedback 并没有真的
# 发送给参数服务器去进行聚合,实际发送的应该是
# grad_with_feedback X q (X: 矩阵乘法)
p, q = compress_and_aggregate(grad_with_feedback, q[w])
# 4. 计算聚合后的梯度
epoch_grad = decompress(p, q)
# 5. 计算动量
momentum = momentum_decay * momentum + epoch_grad
# 6. 计算最新权重值
model_parameter = model_parameter - learning_rate * (epoch_grad + momentum)
由于对power iteration没有完全理解,上述步骤2的计算(也就是计算计算近似矩阵)可能有点问题,如果写错了,还请大家留言告知。
如果是SVD分解这里就很好理解了,这里算法按照SVD分解那一套写的,但是没有进行实验,不知道对不对。
整个算法流程还是很直观的,这里不在赘述。
实验结果
本文的实验结果表明 PowerSGD:
- 在测试集上的表现可以达到SGD的水平,当然速度更加快
- 即使在网络状况不太好的情况下,可扩展性也较强
- 对于大模型来说可以有效降低训练时间
不同压缩算法之间的对比
其中加粗的Rank 7和Rank 2,分别表示r取7和2时PowerSGD的性能表现,可以看到准确率表现甚至超过了不压缩的情况。这里文中也提及了,使用压缩后,这种引入的一些聚合的误差作用有点类似于正则化,所以性能表现会稍微好点。
时间占比
可以看到在worker数量不断增多的情况下,PowerSGD(上图中的Rank 2)的梯度交换的时间消耗优势还是挺明显的。
加速比
可以看到rank 2相比于其它两种压缩算法的加速比有显著优势,加速比超过了0.5。这里的加速比是相对于单个机器进行训练而言的。
引用
[1] Vogels, Thijs, Sai Praneeth Karimireddy, and Martin Jaggi. "PowerSGD: Practical low-rank gradient compression for distributed optimization." arXiv preprint arXiv:1905.13727 (2019).
[2] PowerSGD源码:https://github.com/epfml/powersgd
更多推荐