作者|王佳宁
组织|NewBeeNLP
ChatGPT于2022年12月上旬发布，震惊世界。发布后的这段时间，国内外一系列大型模型训练开源项目相继涌现，如Alpaca、BOOLM、LLaMA、ChatGLM、DeepSpeedChat、ColossalChat等。无论是在学术界还是工业界，都有需要训练大型模型来优化下游任务。
但大量实验证明，在高质量训练语料进行指令调优的前提下， 超过百亿参数量的模型才具备一定的涌现能力 ，尤其是在一些复杂的推理任务上，比如下图：
图片来自论文《Chain-of-Thought Prompting Elicits Reasoning in Large Language Models》
也就是说，如果我们需要通过大规模模型技术来提升业务指标，就必须要求我们训练一个百亿规模的模型。
不过一般情况下，我们没有这么大规模的计算资源，尤其是对于学校一般的科研团队来说，可能只有少量的V100（32G），运气好的话可能有几个A100。因此，在计算能力有限的情况下训练或推理数百亿量级的大型模型是不现实的。因此，无疑有必要在训练和推理阶段都采取一些优化策略来解决此类问题。
本博文主要梳理了大型模型在训练和推理两个阶段的一系列优化技巧，以满足我们在计算资源有限的情况下训练自己的大型模型。本文主要介绍的一些优化技术如下：
混合精度训练 ：FP16+FP32或BF16+FP32； DeepSpeed分布式训练 ：ZeRO-1、ZeRO-2、ZeRO-3； Torch FSDP + CPU Offloading ； 3D并行 ； INT8模型量化 ：对称/非对称量化、量化感知训练； 参数有效性学习（Parameter-Efficient Learning） ：LoRA、适配器、BitFit、P-tuning； 混合专家训练（Mixed-of Experts，MoE） ：每次只训练一些参数； 梯度累积（Gradient Accumulation） : 时间换空间 梯度检查点（Gradient checkpointing） : 时间换空间 Flash Attention
在介绍优化技术之前，先介绍一下如何评估大模型的计算能力。众所周知，如今的预训练语言模型都是基于Transformer结构实现的，因此大模型的参数主要来自于Transformer的Self-Attention部分。 EleutherAI团队最近发表了一篇博客，介绍如何估算大型模型的算力成本。公式如下：
其中： 表示Transformer需要的计算量，以FLOPs为单位；表示Transformer模型包含的参数数量；表示训练数据的规模，以Token为单位；表示吞吐量，单位为FLOPs 表示训练时间；
该公式的原理如下：代表训练过程中的前向传播和后向传播；前向传播的计算成本约等于参数量乘以数据大小的两倍；反向传播的计算成本约等于参数量乘以数据规模的四倍；
它是定量计算成本的单位，通常用FLOP来表示，我们也可以使用一些新的单位来表示：
FLOP/s-s：表示每秒浮点运算的次数； PetaFLOP/s-days：表示实际情况下每秒的浮点运算次数。下图展示了不同规模的预训练语言模型的算力成本：
可见，随着规模的增大，其算力成本将会呈指数级增长。
原文见：Transformer Math 101[1]
混合精度训练是一种非常常见的内存优化技术，适用于单机单卡或多卡并行场景。一般来说，计算机在进行浮点运算时采用FP32（单精度），其中8位用于存储整数部分，23位用于存储小数部分，因此可以存储高精度浮点数字。
因此，在内存优化场景下，牺牲浮点运算的精度可以减少存储容量。例如，当使用FP16进行浮点运算时，只需要一半的存储空间，因此就变成了半精度浮点运算。但FP16的整数最多只能到65536，很容易出现溢出问题。因此，BF16是另一种半精度浮点算术表示形式。与FP16相比，增加了整数部分的存储位，避免计算溢出问题，但也牺牲了一定的精度。
在实际训练中，浮点运算通常采用单精度和半精度混合的方式来实现。典型代表为 动态混合精度法（Automatic Mixed Precision，AMP） ，如下图所示：
O0：表示最原始的FP32浮点运算； O1：除了优化器部分为FP32外，其余均使用FP16； O2：在O1的基础上，在FP32中额外保存一个参数，用于参数更新； O3：所有参数均为半精度； AMP采用混合FP32+FP16，动态指定不同训练阶段将哪些部分转换为半精度进行训练。 AMP通常使用上图的O2部分，即使用混合精度训练不仅可以提高乘法过程的效率，还可以有效避免累加时的舍入误差问题。
Pytorch 1.5版本之后，继承了AMP的实现。调用AMP进行混合精度训练的示例如下：
从torch.cuda.amp 导入自动转换，GradScaler
# FP32 型号
模型=Net().cuda()
优化器=optim.SGD(model.parameters(),)
缩放器=GradScaler()
对于epochs: 中的纪元
对于输入，目标在data: 中
优化器。零_梯度（）
使用自动施法():
输出=模型（输入）
损失=loss_fn(输出， 目标)
scaler.scale(loss).backward()
缩放器。步骤（优化器）
scale.update()
32G GPU 可能无法拟合100 亿个模型的权重、梯度、优化器等参数，但我们或许可以将这些参数按照一定的规则拆分到多张卡上。这就是分布式并行优化的思想。
DeepSpeed是微软的一个开源分布式训练加速框架，它使用了一种称为零冗余（ZeRO）的内存优化技术。本质上是一个 数据并行 分布式训练策略，重点优化数据并行中的内存使用问题。在ZeRO数据并行中，虽然每个GPU都有完整的网络，但每个GPU只保存一部分权重、梯度和优化器状态信息，使得权重、梯度和优化器状态信息可以均匀分布到多个GPU上。
下图显示了DeepSpeed 的三个ZeRO 阶段。假设待训练的模型占用显存120G，集群中有GPU：
Baseline ：传统的数据并行策略，每个GPU存储所有模型的权重、梯度、优化器等参数，在每张卡上并行训练不同的数据，并实现参数聚合和更新。这样的话，每张卡仍然需要加载120G的参数，这在一般机器上显然无法实现； ZeRO Stage1 —— 优化器并行：在训练过程中，优化器状态参数占用的内存空间非常大，因此将优化器状态参数分布到不同的GPU上，此时单卡上的内存占用会大大减少； ZeRO Stage2 —— 梯度+优化器并行：在ZeRO Stage1的基础上，额外分布式存储梯度，可以发现120G的内存占用直接减少到16G； ZeRO Stage3 —— 权重+梯度+优化器并行度：模型所有参数分布式存储，此时一张卡仅占用1.9G。基于ZeRO的原理在训练过程中，有博主分享了更多精美图片，来自【多图秒懂】如何训练“万亿大模型”？ [2]。假设有2张卡，训练2层Transformer模型：
每张卡完整存储了模型的所有参数（橙色部分），包括权重、梯度和优化器。在前向传播过程中，每张卡独立计算馈送的数据，逐层得到激活值（Transformer模型中FeedForward模块的输出）：
计算梯度时，每张卡上的每个模型和每个参数分别计算梯度并存储（紫色部分）：
在梯度更新阶段，对所有卡上的梯度进行平均，然后每张卡独立进行梯度更新，并保存当前优化器状态信息（浅蓝色部分）：
DeepSpeed基于数据并行性，还执行权重、梯度和优化器状态的分布式存储。下图展示了ZeRO Stage3的情况。一开始，假设两张卡分别只存储一层Transformer。当某张卡进行前向传播时，如果此时该参数不存在，则需要向有该参数的卡借用该参数进行前向计算。例如，在GPU1上计算第二层Transformer时，需要将GPU2上的参数复制到GPU1上，以实现第二层Transformer的计算。
这就是为什么在使用ZeRO时，GPU的显存会不断变化。
前向传播结束后，需要进行梯度计算。例如GPU2需要保存w2对应的梯度g2，因此所有其他GPU都将g2梯度发送给GPU2。在GPU2上得到各个GPU的g2梯度后，进行归约运算并保存，得到g2~。其他GPU会删除w2、g2。然后重复该过程，直到所有层都完成反向传播计算：
更新参数时，直接更新即可：
目前HuggingFace的Transformers库已经集成了DeepSpeed框架。您只需要配置ZeRO 文件即可。下面列出一些博主常用的配置：
(1) ZeRO 级1:
{
'train_micro_batch_size_per_gpu' : '自动',
'零优化' : {
‘舞台’: 1,
“cpu_offload”: 假
},
'fp16' : {
'启用' : '自动'
},
“每打印步数”: 1000
(2) ZeRO 阶段2:
{
'train_micro_batch_size_per_gpu' : '自动',
'零优化' : {
‘舞台’: 2
},
'fp16' : {
'启用' : '自动',
“损失规模”: 0，
'loss_scale_window' : 1000,
'初始规模功率' : 16,
“迟滞”: 2，
'最小损失规模' : 1
},
'每打印步数' : 1000,
'优化器' : {
'类型': 'AdamW',
'参数' : {
'lr' : '自动',
'测试版' : '自动'，
'eps' : '自动'，
'weight_decay' : '自动'
}
}
(3) ZeRO 阶段3:
{
'fp16' : {
'启用' : '自动',
“损失规模”: 0，
'loss_scale_window' : 1000,
'初始规模功率' : 16,
“迟滞”: 2，
'最小损失规模' : 1
},
'bf16' : {
'启用' : '自动'
},
'优化器' : {
'类型': 'AdamW',
'参数' : {
'lr' : '自动',
'测试版' : '自动'，
'eps' : '自动'，
'weight_decay' : '自动'
}
},
'零优化' : {
“舞台”: 3，
'卸载优化器': {
'设备' : 'CPU',
‘pin_memory’: 是
},
'offload_param': {
'设备' : 'CPU',
‘pin_memory’: 是
},
‘overlap_comm’: 正确，
‘连续渐变’第： 章真，
'子组大小' : 1e9,
'reduce_bucket_size' : '自动',
'stage3_prefetch_bucket_size' : '自动',
'stage3_param_persistence_threshold' : '自动',
'stage3_max_live_parameters' : 1e9,
'stage3_max_reuse_distance' : 1e9,
'stage3_gather_16bit_weights_on_model_save' : 是
},
'gradient_accumulation_steps' : '自动',
'gradient_clipping' : '自动',
“每打印步数”: 2000，
'train_batch_size' : '自动',
'train_micro_batch_size_per_gpu' : '自动',
“wall_clock_breakdown”: 假
}使用基于HuggingFace的Transformer库时，可以直接指定配置文件，例如：
--deepspeed=./ds_config_fp16_z1.json\
给定一个数组，首先找到该数组中的最大值5.4，然后计算127/5.4=23.5，因此量化因子则为23.5（相当于当前浮点数中最大值放大至-127～127区间内的最大值）。数组中的数乘以量化因子得到的值进行四舍五入估计，即可得到整型数组。解码时，则将整型数除以量化因子即可。由于期间进行了四舍五入估计，因此量化时会有损失。
（2）基于threshold的量化（量化裁剪） 在浮点数范围内，设置两个阈值，记作 和 （），因此当给定一个浮点数 时，可以定义一个裁剪函数：
只保留在区间 范围内的浮点数，其余的则抛弃。该方法又称为饱和量化，由于通过阈值去掉了一些不重要的元素，可以有效解决不均匀问题。
当浮点数的分布均匀时，absmax量化精度损失较小。但当浮点数分布不均匀时，按照最大最小值映射，则实际有效的int8动态范围就更小了，精度损失变大。因此，如果将最大值换为阈值，即超出阈值的部分舍去，在阈值范围内的进行量化，可以降低精度误差。
因此核心的问题是 如何寻找最优的阀值T使得精度的损失最小 。通过实验发现，在range和precision之间的trade-off关系如下图所示：
NVIDIA选择的是 KL-divergence 实现量化校准，其实就是相对熵，那为什么要选择相对熵呢？而不是其他的别的什么呢？因为 相对熵表述的就是两个分布的差异程度 ，放到我们的情境里面来就是 INT8量化前后两个分布的差异程度 ，差异最小就是最好的了。因此问题转换为求相对熵的最小值！
NVIDIA的量化校准流程如下：
收集激活值的直方图；基于不同的阀址产生不同的量化分布；然后计算每个分布与原分布的相对熵，然后选择KL散度最小的一个。
伪量化 是指将模拟量化操作引入训练过程中，如上图（b)，在每个weight的输入后与output的输出前进行伪量化，将浮点量化到定点整型数，再反量化成浮点，用round过程中所产生的误差的浮点值进行前向运算。
伪量化的操作可以使权值、激活值的分布更加均匀，也就是方差更小；相比直接进行后量化的精度损失能更小；能够控制每层的输出在一定范围内，对溢出处理更有帮助；值得注意的是，量化训练中都是采用浮点运算来模拟定点运算，所以训练过程中的量化结果与真实量化结果是有差异的。相关文献：
量化 | 深度学习Int8的部署推理原理和经验验证[8]
Int8量化-介绍（一）[9]
其他常用的量化方法：
PACT：https://arxiv.org/abs/1805.06085v2Dorefa：(PDF) DoReFa-Net: Training Low Bitwidth Convolutional Neural Networks with Low Bitwidth Gradients (researchgate.netLLSQ：Learned Step Size QuantizationLSQ+：LSQ+: Improving low-bit quantization through learnable offsets and better initialization类ChatGPT模型量化：
GPTQ算法GPTQ-for-LLaMa
如上图所示，蓝色部分为原始的模型参数，其将输入 通过一个FC层映射到 。然而矩阵 的训练参数量为 。通过添加一个LORA层（红色部分），将输入 先映射到低纬度空间，再映射回 维度，此时需要的参数量只有 ，其中 为LORA的秩。在训练时，只需要对红色部分的参数进行训练和梯度计算保存，因此大大降低了训练过程中的开销。引入LORA部分的参数，并不会在推理阶段加速，因为在前向计算的时候，红色部分的参数还是需要参与计算的，因此推理阶段应该比原来的计算量增大一点。
接下来给出采用LoRA进行训练的案例，例如选择OPT-6.7B模型进行参数有效性训练时，可以借助HuggingFace PEFT 库实现：
原文Finetune_opt_bnb_peft[10]
使用PEFT库进行训练代码：
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
import torch
import torch.nn as nn
import bitsandbytes as bnb
import transformers
from datasets import load_dataset
from transformers import AutoTokenizer, AutoConfig, AutoModelForCausalLM
from peft import prepare_model_for_int8_training, LoraConfig, get_peft_model
# 正常地加载大模型参数
model = AutoModelForCausalLM.from_pretrained(
"facebook/opt-6.7b",
load_in_8bit=True,
device_map="auto",
)
# 加载tokenizer
tokenizer = AutoTokenizer.from_pretrained("facebook/opt-6.7b")
# 将大模型参数进行INT8量化
model = prepare_model_for_int8_training(model)
# 配置Parameter-efficient LORA
config = LoraConfig(
r=16,
lora_alpha=32,
target_modules=["q_proj", "v_proj"],
lora_dropout=0.05, bias="none",
task_type="CAUSAL_LM"
)
# 获得增加LORA的新模型
model = get_peft_model(model, config)
# 加载数据
data = load_dataset("Abirate/english_quotes")
data = data.map(lambda samples: tokenizer(samples["quote"]), batched=True)
# 获得Trainer
trainer = transformers.Trainer(
model=model,
train_dataset=data["train"],
args=transformers.TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=4,
warmup_steps=100,
max_steps=200,
learning_rate=2e-4,
fp16=True,
logging_steps=1,
output_dir="outputs",
),
data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False),
)
# 模型训练
model.config.use_cache = False # silence the warnings. Please re-enable for inference!
trainer.train()LoRA涉及到如下一些配置：
在推理阶段，只需要加载LoRA的参数，并集成到原始的OPT-6.7B模型中即可，实现如下：
import torch
from peft import PeftModel, PeftConfig
from transformers import AutoModelForCausalLM, AutoTokenizer
peft_model_id = "ybelkada/opt-6.7b-lora" # 他人针对OPT-6.7B训练好的LORA参数
config = PeftConfig.from_pretrained(peft_model_id)
model = AutoModelForCausalLM.from_pretrained(
config.base_model_name_or_path, return_dict=True, load_in_8bit=True, device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
# 将原始的OPT模型与LORA参数合并
model = PeftModel.from_pretrained(model, peft_model_id)
混合专家是一种比较古老的专家系统方法，对于一个决策问题，交给众多专家进行决策投票，根据投票的结果来进行加权求和实现最终决策。在预训练中，则采用了这种思想。
上图中展示了MoE的单层结构，其中包括一个router和若干个expert。router负责决定给每个expert的权重，并制定权重最高的expert作为当前数据进行前后向传播的路由。例如上图中的FeedForward参数有4个，分别指定了FFN2和FFN1作为当前Batch的路由，此时只会对FFN2和FFN1进行参数更新，而其余的参数则固定不变。
因此可以发现，MoE是一种变相的参数有效性训练方法，只不过不同于LoRA等方法，MoE所引入的参数只是控制路由的，且在推理阶段不再使用router，因此对具体的模型推理能力并不起作用。
我们知道Self-attention的计算公式是：
算子主要由“matmul + div + masking + softmax + matmul”几个组成。当Sequence length比较大的时候（例如2048，甚至是GPT-4中的32K），Attention矩阵会是 的空间复杂度，如果在单卡上进行计算，会大量占用显存。
Flash Attention则是让Attention的几个算子能够通过分块并行地进行计算，如果这4个算子都能够分块处理，那么就可以实现这一目的，因此下面一一介绍各个算子的分块处理过程。
（1）：矩阵乘积算子，可以采用分块矩阵的方法进行并行计算。如下图所示，两个矩阵相乘，可以用分块矩阵分别在矩阵 的行、 的列（即 的行）上进行滑动，并将滑动的每个分块结果累加即可：
（2）：这一步需要进行除法操作，因为除法是element-wise的操作，所以非常容易进行分块处理；
（3）：这一步是关键之处，因为涉及到Softmax和乘法操作。特别地，Softmax既不是乘法，也不是element-wise操作，而是对矩阵 的每一行进行归一化，因此需要对该算子单独设计并行处理策略。斯坦福大学团队提出Softmax Tiling策略实现Softmax和乘法算子的并行合并处理。
譬如我们要计算数组x的softmax。然后我们每次只能算2个数，我们先算第1、2个数的softmax，即 cur_sum = exp(x[0]) + exp(x[1]) y[0:2] = x[0:2] / cur_sum pre_sum = cur_sum 然后我们算第3、4个数的softmax，这时候cur_sum会被更新，之前的sum在变量pre_sum里，这个时候我们可以通过把之前前两个数的softmax结果除以cur_sum/pre_sum来得到正确的结果。如果softmax后面还跟一个matmul的话，上次softmax的结果会和D的一个块乘在一起，然后累积起来，这样我们只需要scale这个累积的值就行。依次类推，在每轮循环都把累积的值scale一下，就能incrementally计算softmax或者softmax + matmul的结果。
整个Flash Attention的详细算法流程如下所示：
HuggingFace新版本集成了OpenLLaMA库，其中采用了Flash Attention的训练方法，代码可参考：modeling_open_llama.py[13]
更多技术干货详见：
https://github.com/wjn1996/LLMs-NLP-Algo
