LLMs开源模型们的分布式训练和量化

前一篇博文整理了：

• 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=W0​x+ΔWx=W0​x+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 显存，单卡可推理，同时精度损失没有很大。