训练和推理使用那些大型语言模型，真是挺烧钱的，主要是因为它们太能吃计算资源和内存了。不过啊，我最近发现，用点儿小技巧，就是低精度格式，咱们可以大幅提升训练和推理的速度，快到三倍呢，而且一点儿都不影响模型的准确度。
咱们主要聊的虽然是大型语言模型，但这些技巧其实挺万能的，用在其他深度学习架构上也没啥问题。
咱们得先弄明白啥叫混合精度训练。这可是个牛技术，在现代GPU上训练的时候，能明显加快速度，有时候快个两三倍不是梦！咱们来瞅瞅这是咋回事。
先说32位精度。咱们在GPU上训练深度神经网络，一般用的都是32位浮点运算，这玩意儿比最高精度的64位浮点数要低。你看，PyTorch默认就是用32位浮点数。
可你要是搞传统的科学计算，那64位浮点数才是王道。位数越多，精度越高，计算过程中累积的误差就越少。所以啊，64位浮点数，也就是双精度，因为能表示的范围广，精度高，一直是科学计算的宠儿。
但在深度学习这一块儿，64位浮点运算就显得有点儿过了，成本也高，GPU硬件也不是为它优化的。所以呢，32位浮点运算，也就是单精度，就成了在GPU上训练深度神经网络的标配。
说到浮点数，咱们得知道“位”是啥意思。它就是计算机内存里用来表示数字的二进制位。位越多，表示的数字范围就越大，精度也就越高。浮点数里，数字是这么存的：符号、指数和小数部分。
浮点数的值，就是小数部分乘以底数的指数次幂，再加上符号。小数部分跟小数点后的数字差不多，但也不完全一样。你要是对具体的公式感兴趣，维基百科上有不少好资料。不过咱们简单来说，小数部分就当成一个“分数”或者“小数值”就行。
说回来，为啥要用32位浮点运算而不是64位呢？主要两个原因：
一是内存占用少了。32位浮点数只要64位浮点数一半的内存，这样GPU内存用得更高效，能训练更大的模型，批量大小也能更大。
二是计算能力和速度提上来了。32位浮点运算内存需求少，GPU处理起来更快，训练时间自然就短了。这在深度学习里特别重要，因为训练那些复杂的模型，花个几天甚至几周都是常事。
从32位到16位精度
咱们都知道32位浮点数的好处了，但能不能更进一步呢？当然可以！现在混合精度训练已经挺流行的了，有时候我们会用到16位精度去做浮点计算，这通常叫做“半”精度。
看上面那幅图，float16在指数上少了3位，在小数部分少了13位。
在深入混合精度训练的细节之前，咱们先来个直观的比较。看看PyTorch里的代码示例：
>>> import torch
>>> torch.set_printoptions(precision=60)
>>> torch.tensor(1/3, dtype=torch.float64)
tensor(0.333333333333333314829616256247390992939472198486328125000000, dtype=torch.float64)
>>> torch.tensor(1/3, dtype=torch.float32)
tensor(0.333333343267440795898437500000000000000000000000000000000000)
>>> torch.tensor(1/3, dtype=torch.float16)
tensor(0.333251953125000000000000000000000000000000000000000000000000, dtype=torch.float16)
这些代码示例告诉我们，精度越低，小数点后的准确数字就越少。
深度学习模型对精度低点儿的算术运算还是挺能扛的。大多数情况下，用32位浮点数代替64位的，对模型预测性能影响不大，这买卖挺划算的。但到了16位精度，事儿就有点儿复杂了。你可能会注意到，因为精度不够、数值溢出或下溢，损失可能就不稳定或者干脆不收敛了。
溢出和下溢就是说有些数字太大了，超出了精度能处理的范围。比如：
>>> torch.tensor(10**6, dtype=torch.float32)
tensor(1000000.)
>>> torch.tensor(10**6, dtype=torch.float16)
tensor(inf, dtype=torch.float16)
顺便提一下，上面的代码片段展示了不同精度类型的实际示例。当然，你也可以通过torch.finfo直接查看数值属性，就像这样：
>>> torch.finfo(torch.float32)
finfo(resolution=1e-06, min=-3.40282e+38, max=3.40282e+38,
eps=1.19209e-07, smallest_normal=1.17549e-38,
tiny=1.17549e-38, dtype=float32)
>>> torch.finfo(torch.float16)
finfo(resolution=0.001, min=-65504, max=65504,
eps=0.000976562, smallest_normal=6.10352e-05,
tiny=6.10352e-05, dtype=float16)
上面的代码显示了float32能表示的最大数字是340,282,000,000,000,000,000,000,000,000,000,000,000，而float16的最大值就不能超过65,504。
所以，咱们这里讨论了为什么在现代深度学习中，使用“混合精度”训练比单纯用16位精度训练要好。也可以看之前写的另一篇文章《您需要知道的：大模型中的算力精度FP16 vs. FP32》。但混合精度训练到底是怎么工作的呢？为啥它叫“混合”精度而不是单纯的16位精度呢？咱们接下来就解答这些问题。
混合精度训练机制
为啥叫“混合”呢？其实就是咱们在训练的时候，不是把所有东西都换成16位的浮点数，而是在32位和16位之间来回切换，就这么个意思。
就像下面这张图展示的，混合精度训练的步骤大概是这样的：先把权重弄成低精度的（FP16），这样计算起来快；然后算梯度，算完了再把这些梯度变回高精度（FP32），这样数值就稳定了；最后用这些调整过的梯度来更新原来的权重。
这么干的好处是，既保持了神经网络的准确性和稳定性，又能提高训练的效率。
具体步骤呢，有这么几条：
-
把权重转换成FP16：就是把神经网络的权重从FP32格式转换成FP16。这样内存用得少，计算速度还能提上来。
-
计算梯度：用低精度的FP16权重跑一遍前向和反向传播。这一步骤就是算算损失函数对网络权重的梯度，这些梯度在后面更新权重的时候用得上。
-
把梯度转回FP32：FP16的梯度算好了，再转回FP32。这一步特别关键，能保证数值稳定，避免因为精度低导致的梯度问题。
-
乘上学习率更新权重：现在梯度已经是FP32格式了，就乘上学习率（就是决定咱们优化步长的那个数）。
-
用第4步的结果来更新原来的FP32神经网络权重。学习率这个东东，对优化过程能不能顺利收敛，影响特别大。
听起来可能有点复杂，但其实做起来挺简单的。在后面的内容里，咱们会看到怎么就改一行代码，就能用混合精度训练来微调大型语言模型（LLM）。
混合精度编码示例
咱们来聊聊怎么用代码实现混合精度训练。用PyTorch的话，其实挺简单的，因为它有个自动cast的上下文管理器。再加上开源的Fabric库，切换到混合精度训练就是改一行代码的事儿，这通常也叫自动混合精度训练。
咱们先看个大语言模型（LLM），比如用DistilBERT来给电影评论分类，看看训练时间、准确度和内存需求怎么样。
然后，咱们再聊聊选不同的精度对LLaMA这种大模型有啥影响。
下面就是PyTorch混合精度训练的代码示例：
import torch
from torch.cuda.amp import autocast, GradScaler
# 先初始化模型、损失函数和优化器
model = YourModel() # 换成你的模型
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
# 创建一个GradScaler，用来在反向传播的时候缩放梯度
scaler = GradScaler()
# 训练循环开始了
for epoch in range(num_epochs):
for inputs, labels in dataloader: # dataloader是你的数据加载器
optimizer.zero_grad() # 清除之前的梯度
# 开始autocast上下文，自动切换到低精度计算
with autocast():
outputs = model(inputs)
loss = criterion(outputs, labels)
# 计算反向传播前，先用scaler缩放loss
scaler.scale(loss).backward()
# 应用缩放后的梯度来更新参数
scaler.step(optimizer)
# 更新scaler的缩放因子，为下一步的梯度缩放做准备
scaler.update()
# 每个epoch结束，可以调整学习率啥的
adjust_learning_rate(optimizer)
在这段代码里，autocast就是用来在计算时自动降低精度的，比如从FP32降到FP16。而GradScaler就是用来保证数值稳定的，防止因为精度低了导致的问题。这么一来，我们就能在不影响模型性能的情况下，提高训练的速度和减少内存的占用了。
微调基准测试
让我们从使用float32位精度以常规方式微调DistilBERT模型的代码开始，这是PyTorch中的默认设置：
from datasets import load_dataset
from lightning import Fabric
import torch
from torch.utils.data import DataLoader
import torchmetrics
from transformers import AutoTokenizer
from transformers import AutoModelForSequenceClassification
# ... 省略部分代码以节省空间
### 1 加载数据集
# ... 省略部分代码以节省空间
### 2 标记化和数值化
# ... 省略部分代码以节省空间
### 3 设置数据加载器
# ... 省略部分代码以节省空间
### 4 初始化模型
fabric = Fabric(accelerator="cuda", devices=1)
fabric.launch()
model = AutoModelForSequenceClassification.from_pretrained(
"distilbert-base-uncased", num_labels=2)
optimizer = torch.optim.Adam(model.parameters(), lr=5e-5)
model, optimizer = fabric.setup(model, optimizer)
train_loader, val_loader, test_loader = fabric.setup_dataloaders(
train_loader, val_loader, test_loader)
### 5 微调
start = time.time()
train(
num_epochs=3,
model=model,
optimizer=optimizer,
train_loader=train_loader,
val_loader=val_loader,
fabric=fabric
)
### 6 评估
# ... 省略部分代码以节省空间
print(f"Time elapsed {elapsed/60:.2f} min")
print(f"Memory used: {torch.cuda.max_memory_reserved() / 1e9:.02f} GB")
print(f"Test accuracy {test_acc.compute()*100:.2f}%")
在单个A100 GPU上训练的结果如下：
-
Python实现：CPython
-
Python版本：3.9.16
-
torch：2.0.0
-
lightning：2.0.2
-
transformers：4.28.1
-
Torch CUDA可用？是
...
-
训练准确率：97.28% | 验证准确率：89.88%
-
时间流逝：21.75分钟
-
内存使用：5.37 GB
-
测试准确率：89.92%
接下来，咱们试试float16混合精度训练，就改了一行代码，让Fabric知道咱们要用混合精度：
fabric = Fabric(accelerator="cuda", devices=1)
更改为：
fabric = Fabric(accelerator="cuda", devices=1, precision="16-mixed")
结果如下：
-
训练准确率：97.39% | 验证准确率：92.21%
-
时间流逝：7.25分钟
-
内存使用：4.31 GB
-
测试准确率：92.15%
内存需求下来了，训练速度直接快了三倍。而且，准确率还上去了，可能是因为精度低了，训练时引入的噪声有助于模型泛化。
咱们还试了试普通的float16训练：
fabric = Fabric(accelerator="cuda", devices=1, precision="16")
但发现损失函数不收敛，准确率就跟随机猜的一样，50%。
-
时代：0003/0003 | 批次2700/2916 | 损失：nan
-
时代：0003/0003 | 训练准确率：49.86% | 验证准确率：50.80%
-
时间流逝：5.23分钟
-
内存使用：2.87 GB
-
测试准确率：50.08%
上述结果汇总在以下图表中：
最后，看图表，float16混合精度训练速度跟纯float16差不多快，预测性能还超过了float32，可能是因为精度低带来的正则化效果。
这就是咱们的微调基准测试，用float16混合精度训练，效果确实不错。
张量核心和矩阵乘法精度
咱们再来聊聊张量核心和矩阵乘法精度的事儿。如果你用的GPU支持张量核心，跑之前的代码时，PyTorch可能会在终端提示你：
你正在使用的CUDA设备（'NVIDIA A100-SXM4-40GB'）有张量核心。
想要充分利用这些核心，你应该设置`torch.set_float32_matmul_precision('medium' | 'high')`
这样会牺牲一些精度，但能换来性能的提升。
更多细节，
可以查看 <https://pytorch.org/docs/stable/generated/torch.set_float32_matmul_precision.html#torch.set_float32_matmul_precision>。
默认情况下，PyTorch在做矩阵乘法时会用“最高”精度。但如果你想为了性能多牺牲点精度，也可以这么设置：
torch.set_float32_matmul_precision("high")
或者
torch.set_float32_matmul_precision("medium")
这样设置后，就会用bfloat16数据类型来做矩阵乘法，这货是float16的一个特殊版本，咱们后面会细说。用torch.set_float32_matmul_precision("high"/"medium")就等于是隐式地开启了一种混合精度训练，这得你的GPU支持张量核心才行。
这会对结果有啥影响呢？咱们来看看图：
从图上可以看出，对float32精度来说，降低矩阵乘法的精度能显著提升计算性能，性能能提升2.5倍，内存需求还能减半。而且，预测的准确率也上去了，可能是因为咱们之前说的，低精度能带来正则化效果。
实际上，用较低矩阵乘法精度的float32训练，在性能上几乎和float16混合精度训练一样。而且，对float16训练来说，再降低矩阵乘法精度也不会有性能上的提升了，因为float16混合精度训练本来就是用float16精度做的矩阵乘法。
大脑浮点数
大脑浮点数（Brain Floating Point，bfloat16），是最近流行起来的一种浮点数格式，是谷歌为了机器学习和深度学习应用特别开发的，尤其在他们的张量处理单元（TPUs）上。和传统的float16比起来，bfloat16虽然牺牲了一些精度，但动态范围更广。
这种格式的好处是能更好地表示特别大或者特别小的数，这对于深度学习来说挺重要的，因为可能会遇到各种各样的数值范围。不过，精度低了有时候会影响计算的准确性，或者在某些情况下会出现舍入误差。但在大多数深度学习应用中，这点精度的损失对模型性能影响不大。
bfloat16最初是为TPU开发的，但现在，一些NVIDIA GPU也开始支持这种格式了，比如A100张量核心GPU，它是NVIDIA Ampere架构的一部分。
想知道你的GPU是否支持bfloat16，可以用这段代码检查一下：
>>> torch.cuda.is_bf16_supported()
True
接下来，咱们看看bfloat16能不能给咱们带来额外的好处。咱们可以通过改一行代码，把bfloat16用在之前跑的DistilBERT代码上：
原来是这样的：
fabric = Fabric(accelerator="cuda", devices=1, precision="16-mixed")
现在改成：
fabric = Fabric(accelerator="cuda", devices=1, precision="bf16-mixed")
为了全面比较，我还试了试float64。另外，为了好玩，咱们也试试不用混合精度的bfloat16训练：
结果挺有意思，float64在这里比float32的准确率要高，这和咱们之前说的低精度可能有正则化效果的说法不太一样。不过，bfloat16混合精度训练在预测性能上比float16好一些，就是内存用得多了点。
总的来说，float16和bfloat16混合精度训练的表现差不多，这也在预料之中。
还有个有趣的点，bfloat16的大动态范围让咱们在不用混合精度的情况下也能训练模型，这在普通的float16训练里是做不到的。不过，这事儿也有偶然性，我的经验是，全bfloat16训练有时候效果不如bfloat16混合精度训练。
高效的低精度推理和LLaMA
咱们聊聊深度学习模型在推理时怎么用混合精度来提升效率、减少内存占用，还能让计算更快。不过得注意，用低精度推理可能会让模型准确度稍微下降，但在很多应用里，这点影响不大，算是为了速度和内存做出的可接受的妥协。
之前咱们微调模型的时候，用Fabric设置已经试过16位精度推理了。DistilBERT这模型不大，所以推理速度在整个训练时间里占的不多。
咱们再来看看Meta的LLaMA模型，这个用来生成文本的模型挺火的。咱们可以用Lit-LLaMA仓库（开源的，github 上自己搜索吧），它跟之前一样，用Fabric代码实现16位精度。
因为从头开始在好几TB的数据上预训练这么大的模型成本太高，咱们就用Meta现成的模型检查点来评估模型，生成新文本。
设置好仓库后，咱们用generate.py脚本根据提示来生成文本，默认是用bfloat16：
python generate.py --prompt "Large language models are" # 使用bfloat16
输出大概这样：
加载模型 ...
模型加载时间：24.84秒。
全局种子设置为1234
Large language models are an effective solution to the sequential inference task of natural language understanding, but are unfeasible for mobile applications. In this paper, we investigate a simple, yet effective approach to reduce the computational and memory demands of large language models by removing
第一次推理时间：2.99秒总计，每秒16.70个token
内存使用：13.52 GB想看看float32精度啥样，咱们得手动改改脚本，把Fabric设备类型从torch.bfloat16改成torch.float32。
python generate.py \\
--prompt "Large language models are" # 禁用bfloat16，使用float32
输出大概这样：
加载模型 ...
模型加载时间：17.93秒。
全局种子设置为1234
Large language models are an effective solution to the sequential
data modelling tasks, but the huge size of these models makes them
difficult to learn due to the large amount of parameters and the
time to train them. The high computational cost, as well as the
long training times
第一次推理时间：4.36秒总计，每秒11.47个token
内存使用：27.02 GB
可以看出来，用float32的时候，模型内存占用翻倍了，速度也慢了30%。这说明在像LLaMA这样的大型模型上，用低精度推理确实能减少内存占用，提高速度。当然，选精度的时候，还得根据模型性能和推理效率来权衡。
想让模型在推理的时候性能更上一层楼，除了用低精度浮点数，咱们还能用量化技术。量化就是把模型里的权重从浮点数变成低比特的整数，比如8位的整数，甚至有的能做到4位。不过呢，文章篇幅已经不短了，后面有机会再说吧。
结论
这篇文章里，咱们一块儿看了怎么用16位精度让LLM分类器的训练速度提上去，快了得有3倍。内存消耗也跟着减半了。
咱们还研究了生成式AI模型的推理速度，发现用混合精度技术，性能能提升30%，内存效率也上去了。
所以说，要是你的GPU支持混合精度训练，用这技术绝对超值，简单到改一行代码就搞定！
