saas建站是什么意思,租用的网站空间的缺点,网站建设分几种编程语言,中国建筑股份有限公司 官网前言 虽然TPU的显存令人羡慕#xff0c;但是由于众所周知的原因#xff0c;绝大部分人还是很难日常化使用的。英伟达又一直在挤牙膏#xff0c;至今单卡的最大显存也仅仅到32G#xff08;参考V100、DGX-2#xff09;。然而#xff0c;训练一个24层的BERT Large模型的时候… 前言 虽然TPU的显存令人羡慕但是由于众所周知的原因绝大部分人还是很难日常化使用的。英伟达又一直在挤牙膏至今单卡的最大显存也仅仅到32G参考V100、DGX-2。然而训练一个24层的BERT Large模型的时候如果sequence length开满512那么batch size仅仅开到8有时候能到10就把这寥寥32G的显存打满了。如果想训练一个48层乃至100层的BERT Large那完全是土豪们的游戏了需要疯狂的模型并行分布式多机训练。 但是万能的小夕前不久在Daxiang Dong大佬的安利下发现了 陈天奇 大佬2016年的一篇宝藏paper 传送门https://arxiv.org/pdf/1604.06174.pdf 简单的划一下重点 这篇paper用时间换空间的思想在前向时只保存部分中间节点在反向时重新计算没保存的部分。论文通过这种机制在每个batch只多计算一次前向的情况下把n层网络的占用显存优化到了 ( √)O(n)。在极端情况下仍可用 ( )O(nlogn)的计算时间换取到 ( )O(logn)的显存占用。在论文的实验中他们成功将将1000层的残差网络从48G优化到了7G。且这种方法同样可以直接应用于RNN结构中。 看完摘要瞬间感觉在小破卡上训练BERT Large有救了 此外来快速过一遍paper中最重要的三点结论 梯度计算等价理论上没有精度损失 2. 可以节省4倍的显存开销 3. 训练速度仅仅会被拖慢30% 不过论文发表在2016年当时还没有BERT不过Baidu Paddle团队补了一个BERT的实验结果发现在BERT上面只用22.5%的训练速度损失就能换来5倍的显存开销节省相关实验在本文末尾不着急接下来我们先一起分析一下在训练阶段时显存为什么容易不足。 感谢Baidu Paddle团队提供本节图文素材和测试数据 训练阶段显存为何不足 深度学习中网络的一次训练包含前向计算、后向计算和优化三个步骤。 在这个过程中前向计算会输出大量的隐层变量Tensor当模型层数加深时Tensor数量可达成千上万个。如Bert Large模型单个Tensor可达到1GB这些Tensor在显存中累积显存很快就爆掉了 下图是Bert Large模型在一次训练过程中的显存使用情况可以明显看到在前向计算过程中显存累积趋势是一个陡峭的上升直线。而在反向计算过程中这些隐层Tensor又会很快地被消耗掉又是一个陡峭的下降曲线显存直接降到低位。 那么问题来了为什么不直接删除这些前向计算的Tensor呢 答案很简单因为这些隐层的Tensor在反向的时会被用到手动狗头 来个简单的证明。 假设前向计算中有一个矩阵乘法计算 Y W × X 对W求梯度 很容易发现对W求梯度的公式里有X而X就是那个巨能吃显存的隐层Tensor 那我们是否可以暂时扔掉这些隐层Tensor在反向计算时再把它们重新生成出来呢当然可以这正是上面这篇paper的思想。 重计算 顾名思义重计算就是让每个训练迭代过程做两次前向计算看起来有点奇怪实际上却非常有效对于刚刚那个吃显存的Bert Large支持重计算机制后显存占用直接从175GB降低到20GB陡峭的显存上升直线变成了缓慢增长的Z形曲线如下图所示。 核心思想是将前向计算分割成多个段将每个段的起始Tensor作为这个段的检查点checkpoints。前向计算时除了检查点以外的其他隐层Tensor占有的显存可以及时释放。反向计算用到这些隐层Tensor时从前一个检查点开始重新进行这个段的前向计算就可以重新获得隐层Tensor。 重计算机制有点像玩单机游戏。每过一个关卡就会保存一个检查点而隐层Tensor就相当于游戏中任何一个时刻的图像。普通的训练方式是打通一遍游戏并且将游戏中所有时刻的图像保存下来而重计算机制的思路是先把游戏通关保存检查点后面当收到某一时刻图像的请求时再重打一遍这一关卡就可以了。 如下图举一个简单的例子添加重计算机制前前向计算中需要存储的隐层是4个红点添加重计算机制后需要存储的隐层变为2个蓝点 从而节省了这部分内存。 虽然时间也是宝贵的但重计算方法的性价比很高。在论文的实验中作者用30%的计算时间换取了4倍的内存空间。并且重计算只是重复了一次前向的过程理论上精度没有任何损失。 这么宝藏的算法当然也少不了开源实现。 开源实现 调研了一波似乎TF没有原生支持但是生态里有OpenAI的第三方实现pytorch和paddlepaddle中都有原生API支持 Pytorchtorch.utils.checkpoint PaddlePaddleoptimizer.RecomputeOptimizer 不过pytorch的文档比较略也没有提供更细致的示例和相关数据有兴趣的小伙伴自行试一下。paddle框架中提供了详细到哭的文档甚至还有一个现成的BERT重计算的例子以及非常详细的实验测试结果。这里直接贴过来真香系列 Paddle中实现显存重计算大体分为三步 定义一个经典的优化器如SGD优化器在外面包一层重计算优化器设置检查点。以MLP为例只需要增加两行代码就可以进入重计算模式 import paddle.fluid as fluid
# 定义MLP
def mlp(input_x, input_y, hid_dim128, label_dim2):print(input_x)fc_1 fluid.layers.fc(inputinput_x, sizehid_dim)prediction fluid.layers.fc(input[fc_1], sizelabel_dim, actsoftmax)cost fluid.layers.cross_entropy(inputprediction, labelinput_y)sum_cost fluid.layers.reduce_mean(cost)return sum_cost, fc_1, predictioninput_x fluid.layers.data(namex, shape[32], dtypefloat32)
input_y fluid.layers.data(namey, shape[1], dtypeint64)
cost, fc_1, pred mlp(input_x, input_y)# 定义RecomputeOptimizersgd fluid.optimizer.SGD(learning_rate0.01)
recompute_optimizer fluid.optimizer.RecomputeOptimizer(sgd)
# 设置checkpoints
recompute_optimizer._set_checkpoints([fc_1, pred])
# 运行优化算法
recompute_optimizer.minimize(cost) 该示例github链接 https:// github.com/PaddlePaddle /examples/blob/master/community_examples/recompute/demo.py 此外官方还给出了一个BERT中做重计算的示例 github链接https://github.com/PaddlePaddle/Fleet/tree/develop/examples/recompute/bert BERT实验结论划重点 根据上面paddle官方提供的BERT示例和实验结果得出以下几个结论 结论一 在32GB显存的Tesla V100显卡上应用重计算机制可以训练更大、更深的深度学习模型。当num_tokens为4096(batch size32seqlen128时可以训练100层的Bert网络。 从Github的实验结果也可以看出显存上的收益比速度的损失要大很多 在batch_size上提升了5倍速度只降低了约1/5且精度没有损失。 结论二 模型训练的batch size最大可提升为原来的5倍且只有少量的速度损失。 重计算机制在Bert Large这一模型上收益最大最大batch size从93提升到562而在VGG-16这种比较浅的模型上重计算机制的收益则比较小。这充分符合重计算机制的设计理念为了训练更大、更深的模型。 结论三 在古董显卡Tesla K40显卡12G显存上训练BERT Large时batch size可以开到130 最后希望本文可以帮助大家在小破卡上尽情训练BERT Large
