前一篇博文整理了:
- LLMs开源模型们和数据集简介
这篇博文主要整理一下目前流行的训练方法和量化。
(图自Towards a Unified View of Parameter-Efficient Transfer Learning)
Tuning Strategies
使通用LLMs适应下游任务的最常见方法是微调所有模型参数或微调尾层参数(Freeze)。然而这会导致每个任务都有一份单独的微调模型参数,训练成本高。
-
Adapter。冻结原有参数,添加adapter层用于微调。adapter层一般先向下投影,然后非线性激活函数,再使用向上投影,最后接残差连接。在Transformer中,一个放在MHA之后,一个在FFN之后。新增的参数量大致占原有模型的 0.5%-8%。
h = h + f ( h W d o w n ) W u p h=h+f(hW_{down})W_{up} h=h+f(hWdown)Wup -
P-Tuning。soft-prompt方法。P-tuning v1将learnable prompts插入input embedding中,导致可训练的参数被句子的长度限制,使得任务不通用和规模不通用,而在P-tuning v2中,prompts加在序列前端,并且每一层都加。
-
LoRA。原参数不变,仅训练额外增加的低秩参数以近似权重更新(即W0被冻结,A和B可训练,其中B初始化为0)。当“秩值”远小于原始参数维度时,新增的低秩矩阵参数量将很小。
h = W 0 x + Δ W x = W 0 x + B A x h=W_0x+ \Delta W x=W_0x+BA x h=W0x+ΔWx=W0x+BAx
当进行部署时,以显式的计算和存储低秩矩阵 ,并正常执行推理即可。当需要转换至另一个下游任务,可以通过减去低秩矩阵来恢复W0,然后添加其他的低秩矩阵。 -
BitFit 只对预训练模型中的 bias 向量进行微调,diff-pruning学习一个稀疏的参数更新向量等等。
Optimization
怎么以更低内存和时间训练大模型?
- 存储:将大模型参数和数据全部塞到单卡GPT的内存中进行训练很难。–>分布式并行和混合精度训练
- 计算:大模型参数的所有计算操作可能会导致长训练时间和慢推理时间。–>模型量化
Parallelism
- Data Parallelism。数据并行,复制多份模型,每个副本被放置在不同GPU上,并输入数据分片。然后模型并行完成forward和backward,在每个训练step结束时同步所有模型副本,即聚合梯度 + 下发梯度的AllReduce过程。
- Tensor Parallelism。张量并行,将矩阵乘法进行分块,从而将大矩阵拆分为更小的矩阵,这样就能把不同的矩阵放置在不同的GPU上。这种方式不把整个激活张量或者梯度张量放在单个GPU上,而是把这个张量的碎片放在单个GPU上。
- Pipeline Parallelism。流水线并行,其通过将模型的不同层放置在不同的GPU上,从而将单个大模型分拆至多张卡上,也称为垂直并行。为了提高利用率,引入块 (chunks)概念,即定义同一管级中按顺序输入数据块。例如,GPU0 在 chunk 0、1、2 和 3 (F0,0、F0,1、F0,2、F0,3) 上执行相同的前向路径,然后等待,等其他 GPU 完成工作后,GPU0 会再次开始工作,为块 3、2、1 和 0 (B0,3、B0,2、B0,1、B0,0) 执行后向路径。这与DeepSpeed的梯度累积 (gradient accumulation steps,GAS)类似。
当PP+DP结合时,如果batch size 1024 ,4卡,那么将拆分为 4 x 256,如果块或GAS的数量为 32,则最终得到的 micro batch size 为 8,即每个管级一次处理一个 micro batch。也称micro-batches (MBS) 。
当DP+PP+TP结合成3D Parallelism时:
实践中,广泛使用 Megatron-DeepSpeed (基于Megatron-Deepspeed的GPT-3 训练(Under Construction))进行训练,其中核心是ZeRO来优化DP的储存(类似FSDP)。
ZeRO
DP需要复制多份模型,然后AllReduce。主要存储三类数据,分别是模型参数P,模型梯度G、优化器参数O,即OPG状态,这些备份在每个GPU上都需要存储,存在冗余。当模型变大时,很容易爆显存。
- 参数只在做forward和backward时才用到
- gradients只在最后做AllReduce和updates时才用到
- Adam的optimizer states只在最终做update时才用到
但是实际上每个GPU上可以只存储部分必需的数据,需要其他数据的时候可以去其他GPU上做检索,即以通信时间换显存空间。
(
Φ
\Phi
Φ是模型参数W,K指parameter,momentum和variance占用字节数4+4+4,混合精度下p和g存2+2)
- forward。对W做一次All-Gather,取回分布在别的GPU上的W,得到一份完整的W。forward做完,立刻把不是自己维护的W抛弃。
- backward。对W做一次All-Gather,取回完整的W。backward做完,立刻把不是自己维护的W抛弃。
- gradients。对于完整的梯度G,要G做一次Reduce-Scatter,把其他GPT上的某块梯度进行聚合加总 。聚合后,再立刻把不是自己维护的G抛弃。
- Adam。用自己维护的O和G,更新部分参数W即可。但由于只维护部分W,因此无需其他AllReduce操作。
因此,ZeRO 有三个不同级别,分别对应不同程度的分割:
- ZeRO-1:分割Optimizer States。
- ZeRO-2:分割Optimizer States与Gradients。
- ZeRO-3:分割Optimizer States、Gradients与Parameters。
Mixed Precision Training
虽然希望参数越精准越好,但fp32来表示每个参数的计算开销很大,于是混合精度训练尝试引入fp16。但是fp16表值范围比fp32要狭窄很多,有以下问题:
- 全fp16会损失80%的精度。
- 很多累加操作如梯度累积、softmax容易上溢。–>累加用fp32,余下的操作则fp16。
- 训练后期梯度变小而下溢出。–>保存权重的副本(仅在前向和反向时使用fp16)、Loss Scaling。
但在很多大模型的report中,都显示出fp16在训练中的不稳定。因此它们都不约而同的选择了bf16!它与fp16大小相同,但表值范围和fp32一样。虽然精度更差,但是在结合权重副本的情况下,可以当作是随机梯度下降的一点点折损,再等待下一次迭代。
其他优化
- Gradient Accumulation。梯度累积,解决内存不足而无法训练大Batch Size,做法是累加N个mini-bacth之后,用累积的梯度进行更新,以达到和N*mini-batch一样的效果。
- Activation Checkpointing。只保留每层的输入和输出,丢弃中间结果,反向传递过程时再重新计算。
- Fused Kernels。最小化数据传输。GPU 主要写显存或读显存,并执行计算。当 GPU 忙于读写数据时, GPU 的计算单元就会空闲,因此核融合将多个操作的组合成一个GPU操作,中间结果留在寄存器而不是放回显存。
c = torch.add (a, b) #从显存读a,b,核函数计算,写回显存
e = torch.max ([c,d]) #从显存读c,d,核函数计算,写回显存
两个核函数融合后,只执行融合内核,从而c不会写到显存,这降低了GPU空闲。Megatron-LM提供了几个定制化融合CUDA核,如LayerNorm、缩放、掩码和softmax操作的各种组合等等。 - 硬件故障问题。对于长时间的训练,几乎每周有1-2个GPU故障,使用备份节点+每3h存一次checkpoint。
Quantization
量化是一种常见的模型压缩技术,核心思想是将模型参数从高精度转换为低精度,降低对推理内存的需求,从而可以使大模型在消费级显卡上运行。INT8量化是最流行的训练后量化方法,如下图所示。
-
对称量化symmetry和非对称量化asymmetric。
-
Outlier问题。
-
LLM.int8。自适应混合精度量化方法。分区域设置量化分辨率,并消除异常值对模型量化带来的负面影响。
-
INT4 量化。GLM-130B在INT4 量化级别下最低只需 6GB 显存,单卡可推理,同时精度损失没有很大。