大语言模型（LLMs）训练技巧
作者： AINLP 来源： AINLP
作者：逃脱鱼子酱
文章地址：https://zhuanlan.zhihu.com/p/635321983
Trick 1：CPU offload (CPU卸载)
用额外的通讯开销换取显存。对于模型计算的中间结果（activation，优化器状态等），暂时放到内存（CPU）中，计算需要的时候再放回显存（GPU）中，需要占用传输带宽；
Trick 2：checkpointing (recompute，重计算)
用额外的计算换取显存。即在前向传播时，删除一些暂时用不到的中间激活结果以降低内存，在反向传播时，再根据需要临时进行前向计算恢复；
Trick 3：量化压缩
量化通过减少参数表示的位数来减小模型存储量和计算量，通常会带来一定的模型精度的损失；
Trick 4：通信算子
Trick 5：Ring ALL reduce
Scatter Reduce，每个服务器将参数分为N份，在相邻服务器传递，传递N-1次（Scatter），每个服务器将得到的参数累积起来（Reduce）；
All Gather，将每一份参数的累积同步到所有服务器上去；
总的来说ALLreduce就是对所有服务器上的数据做一个规约操作（如最大值、求和），再将数据写入服务器，如图：
Trick 6：混合精度
模型通常使用float32精度进行训练，随着模型越来越大，训练的硬件成本和时间成本急剧增加。采用float16精度进行训练可以解决这一问题。但是随着模型的训练，梯度值太小小，超出float16表示的精度，导致权重都不再更新，模型难以收敛。
因此模型训练采用混合精度，即训练中模型权重以及梯度使用float16，优化器参数使用float32。同时优化器保存一份float32的权重，以及两个参数状态（均值和方差）。具体的更新步骤如下图：模型使用float16进行前向传播，计算损失值。然后反向传播得到float16的梯度；通过优化器将float16的梯度转化成float32精度的权重更新量；更新float32的权重；将float32的权重转换成float16；
- 主要的显存占用分析
假设一个参数量为X的模型，模型的参数和梯度使用float16，则消耗的显存分别为2X和2X；优化器存储一个float32的模型副本，则会消耗4X的显存，同时存储两个状态参数，分别消耗4X和4X的显存，则总共需要16X的显存。对于GPT-2这样1.5B参数的模型，消耗至少24GB显存；
Trick 7：内存加速
零冗余优化器（Zero Redundancy Optimizer，Zero），一种高效的内存优化技术，可以克服数据并行和模型并行的缺点。在标准的数据并行中，每个GPU需要保存独立的权重、梯度和优化器状态。ZeRO是一种高效的数据并行策略，通过对模型状态(1优化器状态、2梯度和3权重)进行划分后存储在单个的GPU上，然后需要的时候通过动态通信调度来降低单卡显存占用。
7.1 优化器状态划分
将优化器状态划分成Nd，每一份存到不同的GPU上。每个GPU需要存储和更新总优化器状态的1/Nd。显存占用：假设在标准的数据并行中，优化器消耗的显存占用是KX。显存占用从原始的4X+KX降低至4X+KX/Nd。当Nd很大时，后者可以忽略不计。
7.2 梯度划分
在优化器状态划分的基础上，将梯度划分成Nd，每一份存到不同的GPU上。显存占用降低至2X+(2X+KX)/Nd。当Nd很大时，后者可以忽略不计。
7.3 参数划分
在前两者的基础上，将参数划分成Nd，每一份存到不同的GPU上。在前向和反向传播时，用broadcast获取完整参数。显存占用降低至(4X+KX)/Nd。只要有足够数量的显卡，ZeRO能够适应任意大的模型。
Trick 8：模型加速
8.1 数据并行
不同设备执行相同模型，不同数据；
8.2 朴素模型并行加速
当一个模型大到单个GPU无法训练时，最直接的想法是对模型层进行划分，然后将划分后的部分放置在不同的GPU上。整个朴素层并行前向传播和后向传播。假设有两个GPU，GPU1执行前向传播，并将激活(activations)缓存下来。然后将中间结果发送给GPU2，GPU2完成前向传播和loss计算后，开始反向传播。当GPU2完成反向传播后，会将中间的梯度返还给GPU1。GPU1完成最终的反向传播；如图：
- 缺点
低GPU利用率。在任意时刻，有且仅有一个GPU在工作，其他GPU都是空闲的。
计算和通信没有重叠。在发送前向传播的中间结果(FWD)或者反向传播的中间结果(BWD)时，GPU也是空闲的。
高显存占用。GPU1需要保存整个minibatch的所有激活，直至最后完成参数更新。如果batch size很大，这将对显存带来巨大的挑战。
8.3 GPipe
GPipe将minibatch划分为更小且相等尺寸的microbatch来提高效率。前一个计算设备计算出该microbatch对应的结果会马上传给下一个计算设备，同时接着对下一个microbatch进行计算。这样能同时进行通信和计算。Fi,j表示模型的第i+1层对第j+1个microbatch进行forward，Bi,j表示模型的第i+1个层对第j+1个microbatch进行backward。最后会以mini-batch为单位将各个设备上的梯度汇总在一起进行参数更新，即梯度累积。但是设备还是会有一段闲置时间，这被称为bubble。如图：
8.4 合并数据并行与模型并行
对于流水线并行。每个GPU需要与下个流水线阶段(前向传播)或者上个流水线阶段(反向传播)进行通信。
对于数据并行。每个GPU需要与分配了相同层的GPU进行通信。所有层的副本需要All Reduce后对梯度进行平均。
因此，任意给定的GPU都会有两部分的通信，一个是与包含所有相同层的GPU(数据并行)，另一个与不同层的GPU(流水线并行)。如下图，水平方向是完整的一个模型，垂直方向是相同层的不同副本：
8.5 张量并行，又称模型并行
张量并行的核心是将矩阵乘法进行拆分，分配到多个GPU上计算，降低对单个GPU的计算需求。张量并行(TP)需要大量通讯，因此不建议跨多个节点进行张量并行。实际中，若一个节点有4个GPU，最高的张量并行度为4。
一维张量并行，列并行将通信的结果进行拼接，行并行则是对通信结果相加。
8.6 Megatron-LM
针对Transformer的MLP和Attention结构提出了一种高效的张量并行方法。全连接层（MLP）和自注意力层（Sself-Attantion）的张量并行如下图，其中g指矩阵相加：
8.7 3D并行
基于流水线并行将模型按stage划分到不同的管道，每个管道处理一个模型的stage；基于张量并行将模型的每个stage按张量切分，划分成不同块；最后数据并行可以将这种2D组合扩展到任意数量的GPU上。
一个3D并行的例子：模型被分成4个stage，流水线并行度为4；每台机器有8张GPU，张量并行度为4，数据并行度为2；GPU分配如下图：
-
流水线并行：按照流水线并行度将模型分成submodel0~submodel4。
-
张量并行：在每个submodel内部按照张量并行度将submodel的权重切分成4份，每个机器内连续的8张GPU处理一个submodel。
-
数据并行：在每个stage内将submodel复制2份，分别处理2份不同的数据。
-
基于ZeRO的3D并行：允许每个GPU只保存训练步骤所需的一小部分数据(参数、梯度和优化器状态)；
已知Transformer encoder的参数为：embedding(E)，sequence(s)，attention head(ah)，vocabulary(v)，hidden size(h)，layer(n)
自注意力层（K，Q，V，线性变换）= h * h * 4
全连接层 = h * 4h * 2
词表 = v * h
输入 = s * h
设DP=8，TP=8，PP=16，使用基于ZeRO的3D并行，则单张GPU的模型参数量：
采用混合精度，单张GPU的存储开销：
8.8 优缺点分析
Trick 9: FLOPs计算
即floating point operations的缩写（s表复数），意指浮点运算数，理解为计算量。可以用来衡量算法/模型的复杂度。基于标准Transformer decoder结构的模型的FLOPs计算方法如下：
9,1 第一种方法
-
Embeddings
2×seq_len×vocab_size×d_model
-
Attention (Single Layer)
Key, query and value projections=2 × 3 × seq_len × d_model × (key_size × num_heads)
Key @ Query logits=2 × seq_len × seq_len × (key_size × num_heads)
Softmax=3 × num_heads × seq_len × seq_len
Softmax @ query reductions=2 × seq_len × seq_len × (key_size × num_heads)
Final Linear=2 × seq_len × (key_size × num_heads) × d_model
-
Dense Block (Single Layer)
2×seq_len×(d_model×ffw_size+d_model×ffw_size)
-
Final Logits
2×seq_len×d_model×vocab_size
Total forward pass FLOPs=embeddings+num_layers×(total_attention+dense_block) + logits
Total backward pass FLOPs=2*Total forward pass FLOPs
Total FLOPs=Total forward pass FLOPs+Total backward pass FLOPs
9,2 第二种方法
Total FLOPs≈6DN，其中D是总的训练tokens数，N是模型的参数量
**进技术交流群请添加AINLP小助手微信（id: ainlp2)**
**请备注具体方向+所用到的相关技术点**
**关于AINLP**
AINLP 是一个有趣有AI的自然语言处理社区，专注于 AI、NLP、机器学习、深度学习、推荐算法等相关技术的分享，主题包括LLM、预训练模型、自动生成、文本摘要、智能问答、聊天机器人、机器翻译、知识图谱、推荐系统、计算广告、招聘信息、求职经验分享等，欢迎关注！加技术交流群请添加AINLP小助手微信(id：ainlp2)，备注工作/研究方向+加群目的。
**阅读至此了，分享、点赞、在看三选一吧🙏**
更多AI工具，参考Github-AiBard123，国内AiBard123
