作者 | FesianXu
导读
咱们在进行比照学习训练时候，常常须要设置大的 batch size，而显卡的显存大小是限度 batch size 大小的最次要因素，在实际过程中咱们常常采纳 recompute 机制，通过用计算换空间的形式，缩小模型的内存耗费。然，在动态图训练时候，recompute 机制须要进行手动的进行同步和梯度交融，本文纪录下这个问题。
全文 7095 字，预计浏览工夫 18 分钟。
在比照学习场景，或者其余须要大 batch size 的场景中，因为显卡显存的限度，常常会受限 batch size 的进一步增大，此时能够采纳“以计算换空间”的形式缩小模型的显存占用，得而进一步增大 batch size。目前支流框架都对这个机制提供了反对，个别称之为 recompute 或者 checkpoint 机制，比方 pytorch 提供在[1]，paddle（动态图）提供在[2]，tensorflow（动态图）提供在[3]；而在动态图框架中，比方 tensorflow（动态图）提供在[4]，而 paddle（动态图）的这个能力由 fleet- x 提供[5]。为了了解 recompute 机制在分布式场景会导致的问题和解决方案，咱们首先须要理解 recompute 机制，咱们先简略介绍下。
一般来说深度学习网络的一次训练由三局部形成：
前向计算（forward）：在该阶段会对模型的算子进行前向计算，对算子的输出计算失去输入，并传给下一层作为输出，直至计算失去最初一层的后果地位（通常是损失）。
反向计算（backward）：在该阶段，会通过反向求导和链式法则对每一层的参数的梯度进行计算。
梯度更新（优化，optimization）：在该阶段，通过反向计算失去的梯度对参数进行更新，也称之为学习，参数优化。
在之前反向求导公式的推导过程中 [6]，咱们晓得进行反向求导链式传递的时候，须要前一层的激活输入 \(\sigma^{\prime}(z_{j}^{l})\) 作为输出参加本层的梯度计算，如式子 (1-1) 所示（既是 [6] 中的公式(4.1)）：
△(1-1)
公式看起来让人头大，咱们以代码为例子。在个别深度学习框架中，提供对自定义层的梯度定义，如博文 [7] 中介绍的。个别这类型的自定义都会提供两种输出，op 和 grad，如下代码：
# 应用润饰器，建设梯度反向流传函数。其中 op.input 蕴含输出值、输入值，grad 蕴含下层传来的梯度
@tf.RegisterGradient("QuantizeGrad")
def sign_grad(op, grad):
input = op.inputs[0] # 取出以后的输出
cond = (input>=-1)&(input<=1) # 大于 1 或者小于 - 1 的值的地位
zeros = tf.zeros_like(grad) # 定义出 0 矩阵用于掩膜
return tf.where(cond, grad, zeros)
# 将大于 1 或者小于 - 1 的上一层的梯度置为 0 其中的 op 示意以后的算子操作符，而 op.inputs 即是该算子的输出列表，当然如果该算子是中间层算子，那么其输出就是上一层的输入了，而 grad 就是累积的梯度，个别咱们都会对 op 和 grad 进行操作，以计算以后层的梯度。绝对应的一些代码例子，读者有趣味可移步到[8]，笔者实现了一个很简略的主动梯度求导试验例子。
如同有点跑题了，然而笔者以这个例子次要是想通知诸位读者，在模型的训练过程中为了反向梯度计算的不便会贮存很多两头变量，比方前向计算过程中的激活输入值，梯度值等等。有些两头值会被框架主动回收，比方非叶子节点的梯度值是会被主动回收的，见[9]，然而有些两头变量不会，比方此时的中间层的输入值，这些两头变量占据了整个训练过程的大量内存。对于这些两头变量，如果心愿采纳更大的 batch size 进行训练，那么就须要缩小这些两头变量以换取更大的内存两头，recompute 就是依据这个思路设计的。
recompute 将深度网络切分为若干个局部（segment），对于每个局部而言，前向计算的时候，除了小局部必须贮存的变量外，其余两头变量都将被删除；在反向计算的时候，首先从新计算一遍前向算子，以取得须要的两头后果，再失常地运行反向算子。因而，recompute 比照惯例的网络迭代而言，多计算了一遍前向计算，是典型的以计算换空间的“斗争”技术。整个过程如 Fig 1. 所示。
△Fig 1. 前向计算，反向计算和重计算的图示，其中重计算会将除了 checkpoints 之外的非必要两头变量删除，在进行反向梯度计算时候再从新进行前向计算失去。
通常会把切分网络的变量称之为 checkpoints，有大量学者在钻研如何抉择适合的 checkpoints 能力更好地平衡计算性能和内存，通常以 ERNIE，BERT 等为例子，在其每个 Transformer 模块的两头变量作为切分就比拟适合。留神到无论在动态图还是在动态图中，都须要对 checkpoints 进行定义，比方 paddle fleet 中的 recompute 应用如下所示：
dist_strategy = fleet.DistributedStrategy()
# 应用 Recompute，并设置 checkpoints
dist_strategy.recompute = True
dist_strategy.recompute_configs = {"checkpoints": model.checkpoints}
# 定义 checkpoints 作为切分点
optimizer = fluid.optimizer.Adam(learning_rate=configs.lr)
optimizer = fleet.distributed_optimizer(optimizer, dist_strategy) # 设置分布式优化器
optimizer.minimize(model.loss)然而，问题来了。在动态图中应用分布式的 recompute 机制可能并不会有问题，因为动态图的分布式应用暗藏了一些细节，然而在动态图中应用 recompute 机制时候（以 paddle 为例子），则会产生报错如 Fig 2. 所示，类似的报错信息同样在 pytorch 上也会遇到，见[10]。
△Fig 2. 在 paddle 动态图分布式场景中，采纳 recompute 机制将会产生这个报错。
在了解这个报错之前，咱们须要了解数据分布式并行（Data Distributed Parallel，DDP）的逻辑。数据并行指的是将数据程度划分，并给不同的节点（不同的过程或者卡，甚至是分布式节点）进行计算，而后将各个节点的梯度更新后果进行汇总后更新（这个步骤称之为规约，reduce），使得最终每个节点的梯度更新后果是保持一致的。一般来说 DDP 能够分为几个步骤[12]:
1、构建 ：DDP 会将 rank 0 节点上的本地模型参数 state\_dict() 播送到其余节点上，以确保每个节点都是有着同样的模型正本进行初始化的。而后，每个节点上的 DDP 过程将会创立一个本地规约器（reducer），这个规约器用于负责后续反向流传过程中的多节点梯度同步。为了进步通信效率，通常会将多个梯度打包到一个“桶（bucket）”中，并且对整个桶进行规约，由此缩小通信老本，如 Fig 3. 所示。如果某个桶的某些梯度因为某些起因还没有筹备好，那么就须要期待这个梯度准备就绪能力同步，这通常都会影响训练效率。除了装桶外，规约器还须要在构建过程中对每个参数进行主动求导钩子函数（hook）的注册。在反向求导阶段，这些钩子函数在梯度就绪的时候将会被触发。
△Fig 3. 梯度同步以桶为单位进行。
2、前向流传：DDP 拿到输出后就传递给本地模型，如果 find\_unused\_parameters 设置为 True，那么就会持续分析模型的输入。这个模式容许对模型的子图进行反向计算，DDP 会遍历模型的主动求导图，从中找出参加反向计算的参数，并且将所有未应用的参数（也即是不须要加入规约的参数）标识为 ready 状态。在反向过程中，规约器只会期待 unready 状态的参数进行同步，然而规约器同样会规约所有参数，而仅是不会期待这些未应用的参数而已。
3、反向流传 ：反向的 backward() 函数间接蕴含在损失 Tensor 中，而这脱离了 DDP 的管制，因而 DDP 利用在构建阶段注册好的主动梯度钩子进行梯度同步的触发。当一个梯度 ready 后，其对应的 DDP 钩子函数会被触发，DDP 因而会将其参数的梯度标识为 ready 状态，意味着曾经筹备好被规约了。当一个桶中所有梯度都曾经就绪后，规约器就对该桶触发 allreduce 操作，对所有节点该桶的值进行汇总求均匀。
4、优化阶段：对于优化器而言，它优化的是本地的模型。因为每个节点的初始状态和参数更新都是统一的，因而最初的多节点的模型参数也是统一的。
让咱们回到原来的问题上，理解了 DDP 的运行逻辑后，咱们就能读懂这个报错信息了。
Error happened, when parameter[385] [xxxxx@GRAD] has been ready before. Please set fine\_unused\_parameters=True to traverse backward graph in each step to prepare reduce in advance. If you have set , xxxx
当采纳了 recompute 机制后，将会有 K 个 Transformer 模块的 checkpoints 重叠在一起，在进行 loss.backward()的时候，将会对同样的模型触发产生 K 个前向 - 反向过程，这意味着对于同一个参数将会有 K 个主动求导钩子函数进行绑定，一旦某一个钩子函数将参数设置为 ready 后，其余钩子函数就会导致这个报错。因而报错中会显示某个 GRAD 梯度参数曾经被标识为 ready 了，让你关上 fine\_unused\_parameters = True 以遍历反向图进行提前规约，然而你即使设置了同样也会报这个错的，因为实质起因在于 recompute 导致了参数被多个钩子函数所绑定了。[11]
那么怎么解决这个问题呢？一个简略的办法就是将 DDP 的前向 - 反向过程用 no\_sync()上下文进行包裹，此时能够避免 DDP 进行多节点的梯度规约，并且在本地会集所有的本地模型梯度。在退出了 no\_sync()上下文后，手动触发 DDP 的前向 - 反向，进行梯度规约。这个 no\_sync 上下文曾经在 pytorch 和 paddle 中实现了，咱们以 paddle 为例子（pytorch 也是一样的，和 paddle 差异极其小，留神须要 paddle 2.2 以上才反对 no\_sync 上下文）：
# required: distributed
import numpy
import paddle
import paddle.distributed as dist
from paddle.autograd import PyLayer
from paddle.distributed.fleet.utils.hybrid_parallel_util import fused_allreduce_gradients
class cus_tanh(PyLayer):
@staticmethod
def forward(ctx, x):
y = paddle.tanh(x)
ctx.save_for_backward(y)
return y
@staticmethod
def backward(ctx, dy):
y, = ctx.saved_tensor()
grad = dy * (1 - paddle.square(y))
return grad
class SimpleNet(paddle.nn.Layer):
def __init__(self):
super(SimpleNet, self).__init__()
self.linear = paddle.nn.Linear(2, 2)
def forward(self, inputs):
inputs = cus_tanh.apply(inputs)
return self.linear(inputs)
if __name__ == '__main__':
dist.init_parallel_env()
model = SimpleNet()
model = paddle.DataParallel(model)
opt = paddle.optimizer.SGD(learning_rate=0.01, parameters=model.parameters())
for step in range(10):
x_data = numpy.random.randn(2, 2).astype(numpy.float32)
x = paddle.to_tensor(x_data)
x.stop_gradient = False
# step 1 : skip gradient synchronization by 'no_sync'
with model.no_sync():
y_pred = model(x)
loss = y_pred.mean()
loss.backward()
# step 2 : fuse + allreduce manually before optimization
fused_allreduce_gradients(list(model.parameters()), None)
opt.step()
opt.clear_grad()代码中的 model.no\_sync()进入 no\_sync 上下文，在进行本地梯度计算完后，采纳 fused\_allreduce\_gradients 进行多节点的手动梯度规约。当然，将梯度规约全副放到了模型梯度计算完后，这样显然会比一边计算一边同时装桶进行多节点梯度规约来的慢，因为后者能够暗藏一些通信工夫，而前者则齐全是串行的过程。不过这也没方法，目前没有其余解决办法，权且先对付吧。
——END——
参考资料：
[1] https://pytorch.org/docs/stable/checkpoint.html，TORCH.UTILS.CHECKPOINT
[2]https://www.paddlepaddle.org.cn/documentation/docs/zh/2.2/api…\_cn.html#recompute, paddle recompute
[3]https://www.tensorflow.org/api\_docs/python/tf/recompute\_grad, tf.recompute\_grad
[4]https://www.tensorflow.org/versions/r1.15/api\_docs/python/tf/contrib, Module: tf.contrib
[5]https://fleet-x.readthedocs.io/en/stable/paddle\_fleet\_rst/fleet\_large\_batch\_training\_techniques\_cn.html#forward-recomputation-backpropagation，Forward Recomputation Backpropagation
[6]https://blog.csdn.net/LoseInVain/article/details/78092613,《深度学习系列》反向流传算法的公式推导
[7]https://blog.csdn.net/LoseInVain/article/details/83108001，在 TensorFlow 中自定义梯度的两种办法
[8]https://github.com/FesianXu/ToyAutoDiff, Toy Automatic Differentiation on computation graph
[9]https://blog.csdn.net/LoseInVain/article/details/99172594，在 pytorch 中对非叶节点的变量计算梯度
[10]https://github.com/pytorch/pytorch/issues/24005，Using torch.utils.checkpoint.checkpoint\_sequential and torch.autograd.grad breaks when used in combination with DistributedDataParallel
[11]https://github.com/pytorch/pytorch/issues/24005#issuecomment-519719412
[12]https://pytorch.org/docs/stable/notes/ddp.html, DISTRIBUTED DATA PARALLEL
[13]https://www.paddlepaddle.org.cn/documentation/docs/zh/2.2/api…\_cn.html#dataparallel, paddle DataParallel
举荐浏览：
分析多利熊业务如何基于分布式架构实际稳定性建设
百度工程师的软件品质与测试随笔[](http://mp.weixin.qq.com/s?__biz=Mzg5MjU0NTI5OQ==&mid=22475598…)
百度 APP iOS 端包体积 50M 优化实际 (一) 总览
基于 FFmpeg 和 Wasm 的 Web 端视频截帧计划
百度研发效力从度量到数字化变质之路
百度内容了解推理服务 FaaS 实战——Punica 零碎