量化的8位LLM训练和推理使用bitsandbytes在AMD GPUs上

Quantized 8-bit LLM training and inference using bitsandbytes on AMD GPUs — ROCm Blogs

在这篇博客文章中，我们将介绍bitsandbytes的8位表示方式。正如你将看到的，bitsandbytes的8位表示方式显著地减少了微调和推理大语言模型（LLMs）所需的内存。虽然在这一领域有许多用于减小模型尺寸的量化技术，但bitsandbytes通过量化不仅减小了模型尺寸，还减小了优化器状态的大小。这篇文章将帮助你了解bitsandbytes 8位表示方式的基本原理，解释bitsandbytes 8位优化器和LLM.int8技术，并向你展示如何使用ROCm在AMD GPUs上实现这些技术。

要求

bitsandbytes 是一个Python封装库，提供快速高效的8位量化机器学习模型。它现在由ROCm 6.2支持，并且可以无缝部署在AMD GPU上。本文中使用的所有演示均通过以下详细设置进行。有关全面的安装详情，请参阅 ROCm文档。

硬件与操作系统:
- AMD Instinct 加速器
- Ubuntu 22.04.3 LTS

软件:

ROCm 6.2.0
Pytorch 2.3.0

或者，你可以使用相同设置启动Docker容器。

   docker run -it --group-add=video --ipc=host --cap-add=SYS_PTRACE --security-opt seccomp=unconfined --device /dev/kfd --device /dev/dri rocm/pytorch:rocm6.2_ubuntu22.04_py3.10_pytorch_release_2.3.0

安装所需的库。

pip install "transformers>=4.45.1" datasets accelerate "numpy==1.24.4" evaluate
pip uninstall apex #Hugging Face implements layer fusion without apex

注意： Apex 需要卸载，因为 Hugging Face 的一个 bug

如下所示安装bitsandbytes

如果你使用 AMD Instinct MI210/MI250/MI250X/MI300A/MI300X 或更高版本，直接从wheel包安装：

pip install --no-deps --force-reinstall 'https://github.com/bitsandbytes-foundation/bitsandbytes/releases/download/continuous-release_multi-backend-refactor/bitsandbytes-0.44.1.dev0-py3-none-manylinux_2_24_x86_64.whl'

如果你使用旧版 AMD Instinct 设备或希望直接从源代码编译，请按以下指示操作:

git clone --recurse https://github.com/ROCm/bitsandbytes
cd bitsandbytes
git checkout rocm_enabled
pip install -r requirements-dev.txt
cmake -DCOMPUTE_BACKEND=hip -S . #使用 -DBNB_ROCM_ARCH="gfx90a;gfx942" 以针对特定GPU架构
make
pip install .

量化基础知识

量化的一个主要优势是使得大模型在具有较小内存或较少GPU的设备上进行训练和推理变得可行。这是因为量化显著减少了内存占用，这依赖于量化模型使用的特定数据类型。例如，RoBERTa-Large（包含3.55亿参数）在半精度（FP16）下需要约720 MB的GPU内存，而在量化到8位精度（8-bit）后仅需要约360 MB。

在机器学习中，量化涉及将一种数值表示映射到一种更紧凑的表示，例如从FP32映射到INT8。这种映射总会导致信息损失，因为FP32涵盖的范围和精度超过INT8。请注意这两个术语“范围”和“精度”。范围指的是涵盖的数值范围（IEEE的INT8范围是[-127, 127]），而精度指的是小数部分中的数字数量，如[这里](https://sites.radford.edu/~mhtay/CPSC352/Lecture_PDF/Ch02POCA.pdf)所定义。较大的数值数据类型通常具有更广的范围和更高的精度。

例如，如下图所示，原始值`3.1415927`在IEEE-754的FP32格式中存储为`3.1415927410`，而在FP16中存储为`3.140625`。这种表示中的不同逼近是由于两种数据类型分配给指数（粉色单元格）和尾数（蓝色单元格）组件的位数不同。因此，当我们将值从一种表示转换为低精度表示时，会有信息损失。

precision

示例图来源于Maarten Grootendorst的文章。

这种信息损失会随着时间的推移在深度学习模型中累积，导致准确度下降。当我们将FP32权重量化为INT8时，这种情况会变得更严重，因为后者只能存储整数，从而导致更高的量化损失。为了解决这个问题，提出了几种量化技术，如零点量化和最大值量化。这些技术大多涉及通过统计变换（例如，最大绝对值、最小值-最大值、零点和尺度）将模型权重归一化为INT8，然后进行四舍五入。这为大型语言模型提供了“足够好”的推理准确度，不同于全精度模型。您可以参考[这篇文章](A Visual Guide to Quantization - Maarten Grootendorst)以深入了解量化的数学原理。

bitsandbytes的8-bit优化器

量化在推理过程中减少模型的内存消耗，从而允许像OPT-66B这样的超大模型能够适应单个AMD MI210 GPU。然而，当你想要训练模型时，情况可能并非如此。Adam和带动量的SGD等有状态优化器在模型的训练过程中，为模型的每一个参数维护状态值，并用于梯度更新。这导致了内存需求增加至模型大小的两到三倍。例如，Adam在全精度（32位）下每个参数的两个状态消耗8字节内存，对于一个拥有10亿参数的模型，需要8GB的专用内存。

为了降低这种内存需求，bitsandbytes提出了一种称为块式动态量化(Blockwise Dynamic Quantization)的8位量化方法，用于训练期间的优化器状态。这将把Adam状态的内存消耗从32位减少到每个状态仅8位。然而，使用如此有限的8位范围会因为以下三方面的挑战而变得困难：量化精度、计算效率和大规模稳定性。在讨论这些问题如何解决之前，我们先通过一个例子来理解块式动态量化的过程。

量化过程可以描述如下：

将张量分成大小为4096或更大的块（块式）。
找到每个块的局部绝对最大值。
使用绝对最大值进行归一化，将范围减小到[-1,1]。
使用“动态数据类型”找到最近的8位值。
存储该值的索引。

下面的图片展示了量化和反量化的过程。

quantization

图片来源：块状动态量化论文

块式量化减少了通过将长张量划分为小块来计算归一化的成本，这意味着这一过程是特定块内局部化的。这减少了GPU内的多核同步化加速归一化过程。而且，块分割还帮助孤立异常值，从而在训练中实现稳定性和效率。

动态量化提出了一种动态8位数据类型，来适应在[-1,1]范围内的高精度和高幅值张量值。这种数据类型，与常规FP16/FP32/INT8数据类型不同，没有固定数量的位数用于指数和小数部分。这种位数分配的非线性有助于捕捉由优化器状态展示的极值。如下图所示，第一位作为符号位，直到指示位（即下一个'1'）的连续零位用于指数部分。其余位数分配给小数部分。通过这种方式，我们可以用8位表示高达10^(-7)的指数。与其他量化技术相比，动态量化对非均匀分布（如优化器状态中可能出现比大多数状态大几个数量级的离群值）具有更小的量化误差。

dynamic-datatype

图像来源是Tim Dettmers的视频：高效深度学习的8-bit方法

最后，Blockwise Dynamic Quantization（块状动态量化）还建议在大型语言模型中使用稳定嵌入层代替常规嵌入层。这包括Xavier均匀初始化、位置嵌入之前的层归一化和32位优化器状态。基于实验数据，作者发现嵌入层的梯度较高。因此，为了避免传播量化误差，他们仅对该层使用32位优化器状态。

总而言之，Blockwise Dynamic Quantization有效地解决了8-bit量化问题，具体如下：

块状量化隔离异常值，实现了在训练大型模型时的稳定性和效率。
动态量化以动态8-bit表示存储浮点值。它提议使用一种新的数据类型，可以适应较高幅度值和高精度值，从而产生低量化误差。
稳定嵌入层在大型语言模型中基于实验数据提高了训练稳定性。

使用8-bit优化器进行模型训练显著减少了内存消耗而不降低准确性。这反映在在相同硬件约束下量化和非量化训练模型大小的差异。

results

图像来源是Blockwise Dynamic Quantization论文

我们可以使用bitsandbytes和HuggingFace的官方集成在AMD的Instinct GPU上无缝运行bitsandbytes的Blockwise Dynamic Quantization。以下是使用Adam 8-bit优化器在单个GLUE任务‘cola’上训练Google的t5-11B模型的代码。该代码基于HuggingFace上关于GLUE任务的文本分类示例。

from datasets import load_dataset
from evaluate import load
from transformers import AutoTokenizer
from transformers import T5ForSequenceClassification
from transformers import TrainingArguments, Trainer, TrainerCallback
import numpy as np
import bitsandbytes as bnb
import torch
from sys import argvdataset = load_dataset("glue", data_dir="cola")
metric = load("glue", "cola")
model_checkpoint = "google-t5/t5-3b"tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, use_fast=True)def preprocess_fun(examples):return tokenizer(examples["sentence"], truncation=True)dataset = dataset.map(preprocess_fun, batched=True)model = T5ForSequenceClassification.from_pretrained(model_checkpoint, device_map='cuda', torch_dtype=torch.float16)train_args = TrainingArguments(f"T5-finetuned-cola",evaluation_strategy = "epoch",save_strategy = "epoch",learning_rate=2e-5,per_device_train_batch_size=16,per_device_eval_batch_size=16,num_train_epochs=1,load_best_model_at_end=True
)def compute_metrics(eval_pred):predictions, labels = eval_pred# print(predictions)predictions = np.argmax(predictions[0], axis=1)return metric.compute(predictions=predictions, references=labels)if argv[-1]=='1':print("Using bnb's 8-bit Adam optimizer")adam = bnb.optim.Adam8bit(model.parameters())
else:adam = None # 默认使用常规Adamtrainer = Trainer(model,train_args,train_dataset=dataset["validation"],eval_dataset=dataset["validation"],tokenizer=tokenizer,compute_metrics=compute_metrics,optimizers = (adam,None)
)
trainer.train()

该脚本可以在这里找到。在运行模型之前，打开一个不同的终端并使用AMD的官方amd-smi工具监控内存消耗。

amd-smi monitor -vGPU  VRAM_USED  VRAM_TOTAL0     10 MB    65501 MB

在全精度下训练模型一个epoch而不进行任何量化，并观察在训练结束时模型的内存占用情况，这是在第二个终端上。显示60.5%的内存被占用。

python optim8bit.py

amd-smi monitor -vGPU  VRAM_USED  VRAM_TOTAL0   39671 MB    65501 MB

现在，使用INT8量化训练模型一个epoch，并观察在训练结束时模型的内存占用情况，这是在第二个终端上。此时仅占用了36%的内存。

python optim8bit.py 1

amd-smi monitor -vGPU  VRAM_USED  VRAM_TOTAL0   24190 MB    65501 MB

如上所述的指标显示，8-bit优化器运行高效，使得内存需求在AMD Instinct GPU上减少了约41%。

LLM.int8() by bitsandbytes

Bitsandbytes 还提出了 ‘LLM.int8()’，这是首个针对 transformer 模型进行 INT8 量化的多亿参数规模的推理过程，没有任何性能下降。这是第一个针对高达 175B 参数的模型提出量化的研究。作者观察到，因缺乏较高的量化精度和存在的**突现特征 (emergent features)**，使得大于 1B 参数的模型的 INT8 量化扩展失败。他们在 LLM.int8 paper 中提出了两个调整——向量式量化和混合精度分解，以解决这两个问题。

如前一部分所述，块式量化使用多个归一化常量处理同一个张量，从而将异常点的影响限制在块内，提升了稳定性和效率。类似地，对于矩阵乘法，本文采用向量式量化，通过采用多个归一化常量来提高效率。

块式量化是独立地对矩阵的行/列进行量化，而不是对块进行量化。这样可以减少矩阵乘法输出的反量化费用。例如，假设有两个矩阵 A 和 B，大小分别为 m×n 和 n×p。矩阵乘积 AxB 的每个元素是 A 的一行和 B 的一列的内积。在向量式量化中，我们通过常量 a 对行向量归一化，通过常量 b 对列向量归一化，然后执行内积。在反量化过程中，我们可以通过 a×b 轻松恢复值。

LLM.int8() 的第二个方面，也是确保 INT8 量化大模型的核心组件，便是混合精度分解。作者在定量分析中注意到，大模型中的量化误差在矩阵乘法过程中可能累积，导致某些隐藏维度的值高达正常值的 20 倍之多。作者称这些极端异常点为 **突现特征 (emergent features)**。这些突现特征是导致量化模型性能下降的原因。

作者还注意到这些突现特征的出现是稀疏而系统性的，因此仅仅通过向量式量化并不能修复参数超过 6.7B 的模型（如下图所示）。然而，LLM.int8() 维持了与 16 位基准相同的准确率。

feature-emergence

Image source is LLM.int8 paper

要支持有效量化，作者开发了混合精度分解技术。首先，他们开发了一套依据实验数据的标准来识别这些突现特征。根据实验数据，他们定义突现点为具有至少 6.0 值的特征（正常特征维度范围为 [-3.5,3.5]，见附录 B 表 4），影响至少 25% 的层，并且影响至少 6% 的序列维度。使用这些标准，他们在所有模型的 25% 的 transformer 层中找到了突现特征。在 6.7B 参数之后发生相位变化，突现特征出现在 100% 的层中。

Image source is from Finding Outlier Features, Section 4.1, LLM paper

他们发现这些特征发生在隐藏维度（列）。将这些特征置零会进一步降低性能，表明这些特征的重要性。将其量化为 INT8 也无济于事，如下图中的橙色曲线所示。因此，他们将矩阵乘法分解成两种形式。对于异常列使用 16 位矩阵乘法，其他元素使用 8 位矩阵乘法，然后将两者相加（结果为 16 位）。向量式量化和混合精度分解的结合如下图所示：

llm-int8-example

Image source is LLM.int8 paper

这种分解不仅能处理异常点，还在推理大型语言模型时实现了稳定性和更好的性能。由于 99.9% 的 transformer 列具有正常特征，少于 1% 的列含有异常点，因此大部分元素存储为 INT8。这减少了模型的内存占用并加快了推理速度。

在下面的代码中，我们在 AMD GPU 上加载 Facebook 的 66B 参数的 OPT 预训练模型，并使用 bitsandbytes 配置将其量化为 INT8。我们看到 INT8 模型完美地适配 GPU 内存，并成功进行推理。

import torch
from transformers import AutoModelForCausalLM, BitsAndBytesConfig, AutoTokenizerquantization_config = BitsAndBytesConfig(load_in_8bit=True)
tokenizer = AutoTokenizer.from_pretrained('facebook/opt-66b')
model_8bit = AutoModelForCausalLM.from_pretrained("facebook/opt-66b", quantization_config=quantization_config, torch_dtype=torch.float16 # 将模型加载为 float16 并量化为 INT8 并移动到 GPU 设备
)
print(model_8bit.device)
inp = torch.randint(100,(1,20))
print(model_8bit.generate(inp, max_new_tokens=1))

 warnings.warn(
`low_cpu_mem_usage` was None, now set to True since model is quantized.
Loading checkpoint shards: 100%|████████████████████████████████████████████████████████| 14/14 [01:29<00:00,  6.41s/it]
device(type='cuda', index=0)
tensor([[92, 82, 37, 79, 95, 76, 84, 42, 60, 62, 21, 48, 18,  0, 72, 31, 48, 99,69, 69,  5]])

你可以通过以下命令来验证内存消耗：

amd-smi monitor -v

模型占用了 63.48 GB 的 64 GB 内存，如预期所示。

GPU  VRAM_USED  VRAM_TOTAL0   64908 MB    65501 MB

而如果你在 GPU 上以 half 精度或全精度加载模型，则会遇到 OOM 错误。

import torch
from transformers import AutoModelForCausalLM, BitsAndBytesConfigmodel_8bit = AutoModelForCausalLM.from_pretrained("facebook/opt-66b",device_map='cuda'
)

OutOfMemoryError: HIP out of memory. Tried to allocate 324.00 MiB. GPU

总结

在这篇博客中，我们讨论了 bitsandbytes 提出的快速高效的 8 位表示，这些表示在训练和推理大语言模型时显著减少了内存占用。bitsandbytes 提出的 8 位优化器和 LLM.int8 技术不仅在内存使用上表现出色，还在训练过程中提供了稳定性和效率，使得大语言模型（LLM）更易于访问。这些技术包括块式量化、动态量化、稳定嵌入层、向量式量化和混合精度分解。在 AMD GPU 上使用 Hugging Face 的这些技术可以大幅减少内存消耗，约 50%，这使得 AMD 的 Instinct GPU 在现代生成式 AI 工作负载中更具优势。