关于分布式:浅论分布式训练中的recompute机制

作者 | 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，并设置checkpointsdist_strategy.recompute = Truedist_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零碎