为什么需要BMTrain?
PLM越来越大。为了使训练更高效和廉价。我们有必要
1.分析GPU的显存去哪了?
2.理解GPU间的合作模式是如何的?
显存都去了哪里?
CPU vs GPU
CPU适合复杂逻辑运算。GPU适合大量重复的数值运算。
显存成分
1.前向传播时,模型参数
2.反向传播时,模型梯度
3.中间过程的计算
4.优化器
GPU间的合作模式
数据并行-Data Parallel
参数服务器的参数会被复制到所有显卡上
数据切成三份,分别给每张显卡
最后聚合后的梯度传回参数服务器
PS:实际上参数服务器在0号显卡上
广播算子
规约
All Reduce
规约的结果广播
Reduce Scatter
All Gather
分布式数据并行
这样就不需要要参数服务器。
使用数据并行,使模型中间结果量降低了
模型并行-Model Parallel
一张GPU可能无法放下模型所以的参数和梯度和优化器。
由于矩阵乘法可以进行分解。将大矩阵切成小矩阵。
但是这样要保证每张GPU上的数据是一样的。
梯度需要拼接,所以需要All gather算子。
中间结果没有减少。但是模型参数和梯度和优化器参数减少了。
ZeRO Redundancy Optimizer
是基于数据并行的架构。
Zero阶段1
每张显卡更新一部分模型参数。
最后拼接每部分得到的模型参数。
Zero阶段2
继续优化。
阶段1需要在反向传播获得梯度后,进行reduce scatter。
由于优化器只需要用到gradient*。所以中间的gradient可以从显存中移除。
移除的时机是在进行反向传播的过程中。
Zero阶段3
继续优化。
前面数据并行并没有解决模型参数存储在GPU上的问题。
为此,前向传播过程中需要进行All Gather操作。
用完所有参数时,就进行释放。
PS:Zero3相比Zero2是用时间换空间的方法。
比较
流水线并行-Pipeline Parallel
模型不同的层分到不同的GPU。
弊端是:1号显卡工作时,后面的显卡处于空闲状态。
有一些方法是优化以解决资源浪费问题的。
优化技术
混合精度训练
FP16的数值表示范围更小,但是计算更快。
一般FP32。但是有时候也可以从FP32转到FP16。
将入梯度乘以学习率小于FP16的最小范围,那么可能产生下溢。它对参数的更新就会被忽略。
所以需要把参数更新量表示为FP32。
具体上,在混合精度训练中,为了加速前向和反向传播,会使用FP16的梯度。然后更新参数量用FP32进行累计。最后参数用FP16。
Offloading
优化器的参数可以放在CPU中。
通过把一张显卡绑定在多张CPU上,可以将每张CPU上的计算量降低。能够让CPU不会成为模型训练的瓶颈。
Overlapping
GPU中,memory操作一般是异步的。先给memory发送请求,然后进行其他计算,完了之后对memory请求进行接受。
Checkpointing
为了支持反向传播,需要保存中间结果。
通过设置检查点,只保留每个transformer层的输入。反向传播时,进行重计算,临时每个大层所有线性层的输入。
过了一层,就可以将检查点和临时重计算的中间结果从显存中扔掉。
BMTrain-使用介绍
表现:高效,便宜
使用时只需要进行简单替换。
BMCook
背景介绍
介绍大规模预训练模型压缩的相关技术。以及相关工具包BMCook。
下表是PLMs模型增长的趋势。
如何将大规模的计算量降下来,同时保留PLMs学习到的能力。
所以希望将大规模模型压缩。同时小模型基本上继承大模型的能力。
有效的方法可能包括:知识蒸馏;模型剪枝;模型量化;模型的专家化
知识蒸馏
假设:小模型只是去拟合大模型在输入空间的一个子空间的z映射。
teacher模型提供soft label,可能比直接提供金标准让student模型去学习,效果会更好。
第一篇关于PLMs的知识蒸馏的论文是PKD。
它的改进是student模型可以对teacher模型的中间层进行学习。
下面的工作进一步探索了老师模型中可以用做蒸馏计算的信号。
模型剪枝
剪枝可以分为非结构化剪枝和结构化剪枝。
研究发现,非结构化剪枝对计算的加速效果非常有限。
所以一般而言,对加速有用的是结构化剪枝:一次性将矩阵的一列或一行或一块删掉。
基于PLMs的剪枝工作和观察
对bert进行剪枝。
PLMs结构化剪枝的工作
注意力层剪枝。
层剪枝:随机dropout一些层
模型量化
神经网络并不需要很高的精度进行计算。所以考虑将浮点表示转化为定精度的表示。
这样可以把表示的位数和计算的位数降下来。
量化在研究方面的挑战
下图展示了不同的定精度表示对模型性能的影响。
其他模型压缩方法
参数共享-Weight Sharing
低秩分解-Low-rank Approximation
结构搜索-Architecture Search
BMCook-使用介绍
现在的PLMs是十分过参数化的。有一些方法被用于提高模型效率。
BMCook是一个工具包。它的目的是结合已有的有效的模型压缩方法,加速现有大规模模型。
BMInf
背景介绍
在部署大模型CPM2的demo时,遇到了一些问题:
1.对于单个应用的实力,就需要4块A100的GPU来推理。
2.单词请求需要10秒才能处理,即qps=0.1。
3.成本很高。4块A100一天费用为1200。
于是考虑能否在用户自己的电脑上把大模型跑起来。
市面上常见的GPU是GTX1060。
面临的困难
1.大模型有很高的显存占用
2.大模型推理需要的算力要求很高。
深入理解Transformer
分析Transformer可以发现。模型推理过程中,大量时间在进行线性层运算。
所以要考虑如何优化线性运算。
考虑如何在允许一些精度损失的前提下,来优化整个线性层的运算效率。
量化-Quantization
将浮点数缩放到-127到127的范围内,这样就能用INT8来近似表示。
这种方法在模型尺寸小的时候还行,但是在Transformer上效果不好。原因可能是矩阵中有好几百万数字,只用一个缩放因子效果不好。因为相当于原本有好几万个值,但是直接变成只能是-127到127的值,表达能力下降。
更精细缩放
进行细粒度缩放,如行矩阵进行缩放。
优化后。
内存调度-Memory Scheduling
参考虚拟内存。将CPU也利用起来。
可以将暂时不用的参数放在CPU上。
实际测试发现:传输一层的时间往往超过计算一层需要的时间
所以考虑2层用于调度,其他n-2层固定
尽量扩大加载2层的间隔。