本章涵盖：
- 使用 FlexGen 技术，将 LLM 的一部分卸载到内存或磁盘
- 使用 SmoothQuant 这种高级训练后量化技术，减少内存占用并加速推理
- 使用 BitNet 这种可扩展的 1-bit Transformer 架构，减少内存占用和能耗
通常来说，领域专用语言模型比较小：根据我的专业经验，它们通常不超过 70 亿或 80 亿参数，不包括 LoRA 或 QLoRA 部分调优带来的额外参数，而这些额外参数本来就很小。前面章节中解释的技术，已经使这类模型能够在计算资源受限的环境中实际部署和使用。不过，如果专门化训练数据规模更大，非结构化训练数据的表示比自然语言更大，或者使用了更适合任务但更大的基线模型，那么即使经过量化，最终模型大小也可能超过可用推理硬件的容量。
本章介绍高级 LLM 量化技术。这些技术可以帮助你在目标环境中高效运行这类专门化模型，并在质量和速度之间取得平衡。下面各节会描述这些策略，并比较它们的优缺点。你还会看到详细的代码示例和 benchmark。
注意
本章介绍的是高级 LLM 量化技术，适用于第 6 章中的方法无法满足性能或计算约束的情况。在继续阅读之前，请确保你已经理解第 6 章。
9.1 FlexGen
FlexGen 是一个高吞吐推理引擎，用于在 GPU 内存受限的情况下运行 LLM。它在 2023 年论文《FlexGen: High-Throughput Generative Inference of Large Language Models with a Single GPU》中提出。FlexGen 可以根据不同硬件约束进行配置，通过在 GPU、CPU 和磁盘之间汇聚内存与计算资源来运行模型。它通过求解一个线性规划问题，搜索高效存储和访问 tensor 的方式，并进一步将权重和 attention cache 压缩到 4 bit，同时几乎不损失准确率。这扩大了可行 batch size 的范围，并提高了模型最大吞吐量。论文报告称，它相比传统 offloading 方法能实现显著更高的吞吐。
为了让你理解 FlexGen 能带来什么，可以考虑一个 GPT-3 规模的模型，拥有 1750 亿参数。使用第 4.2 节中的公式，加载它所需的 VRAM 在 FP32 下为 840 GB，在 FP16 下为 420 GB。这意味着，如果要把模型加载并用于推理，你需要例如 6 块 NVIDIA A100 Tensor Core GPU，见图 9.1。
图 9.1 —— 将一个 1750 亿参数语言模型加载到内存所需的硬件
考虑到硬件价格，以这种方式运行单个模型可能非常昂贵——往往超出中小型组织，甚至一些大型非科技公司的预算。假设有一个节点，只有一块消费级 GPU，也就是 16 GB DRAM，但有更多 RAM 和磁盘空间，而 RAM 和磁盘比 GPU 更便宜，也更容易获得。FlexGen 线性规划优化器会找到最佳方式，将模型权重、激活和 attention KV cache 存储和访问在 GPU、CPU 与磁盘之间。结果就是模型的一部分会被卸载，如图 9.2 所示，这里假设模型以 FP16 格式加载。这个方法显然比购买或租用 6 块 NVIDIA A100 GPU 更便宜。
图 9.2 —— 使用 FlexGen 跨 GPU、CPU 和磁盘加载 1750 亿参数语言模型
这种 offloading 可以让我们在硬件受限环境中部署非常大的模型，但它会增加延迟并降低吞吐。这是预期之中的：CPU 和磁盘比 GPU 慢，而 offloading 会增加额外 I/O。
FlexGen 不只是基础 offloading：在生成式推理期间，它会决定哪些 tensor，也就是权重、激活和 KV cache，需要 offload，放在哪里，以及什么时候移动它们，以便在保持 I/O 高效的同时提升吞吐。它还通过将权重和 KV cache 压缩到 4 bit 来进一步提升推理性能，并且几乎没有准确率损失。这里使用的是非常大模型的极端示例——当你为领域专用任务微调模型时，通常不会达到这个规模——但它以简单方式说明了 FlexGen 方法的好处。
注意
FlexGen 只通过 CUDA 支持 NVIDIA GPU，并且没有计划支持其他厂商的硬件加速。
我们来看看 FlexGen 在实践中如何工作，以及它是否实现了论文中的主张。本节代码包含在配套 Colab notebook 中。
FlexGen 是开源的，可在 GitHub 上获得，仓库名为 FlexLLMGen；你也可以直接从 PyPI 安装，而不必从源码构建。它包含 Facebook 的 OPT 和 Galactica 模型，以及 BigScience 的 BLOOM 模型的预构建实现，也支持从这些 baseline 微调而来的模型。本节聚焦 OPT，这是一组 decoder-only 预训练 Transformer，规模从 1.25 亿参数到 1750 亿参数不等。论文报告称，OPT-175B 与 GPT-3 相当，但开发所需碳足迹只有后者的七分之一。所有 OPT 模型都托管在 Hugging Face Hub 上。
配套 notebook 中的这个示例需要硬件加速，并在 Colab 免费层 GPU VM 上运行。我们会从 OPT-1.3B 模型开始，先了解 FlexGen。安装 FlexGen 代码包后，具体如 notebook 所示，可以从 CLI 运行 flex_opt Python 脚本，在一个受支持模型上评估 FlexGen：
python -m flexgen.flex_opt --model facebook/opt-1.3b
在最简单形式下，flex_opt 需要 model 参数，也就是 Hugging Face Hub 上目标模型的 ID。下载模型后，脚本会根据所选模型配置将权重转换为 FlexGen 格式，然后把权重、attention cache 和 activations 加载或卸载到目标存储中。默认情况下，如果你在运行时没有特别指定，所有内容都会加载到 GPU。最后，它会运行一个小 benchmark。该命令输出如下：
TorchDevice: cuda:0
cur_mem: 2.6505 GB, peak_mem: 3.2478 GB
TorchDevice: cpu
cur_mem: 0.0000 GB, peak_mem: 0.0000 GB
model size: 2.443 GB cache size: 0.398 GB hidden size (p): 0.008 GB
peak gpu mem: 3.248 GB projected: False
prefill latency: 0.350 s prefill throughput: 5846.976 token/s
decode latency: 0.932 s decode throughput: 132.989 token/s
total latency: 1.283 s total throughput: 99.791 token/s
可以看到，所有内容都运行在 GPU 上，因为该模型能够装入我们执行环境中的设备。现在尝试一个更大的模型，例如 OPT-6.7B，它无法装进我们的 GPU。flex_opt 脚本的另一个关键参数是 percent。它允许指定权重、attention cache 和 activations 中有多少百分比加载到 GPU，有多少百分比 offload 到 CPU。它需要一个由 6 个整数组成的字符串：
- GPU 上权重的百分比
- CPU 上权重的百分比
- GPU 上 attention cache 的百分比
- CPU 上 attention cache 的百分比
- GPU 上 activations 的百分比
- CPU 上 activations 的百分比
默认值是 100 0 100 0 100 0，即所有内容都保留在 GPU。如果更大的模型装不进 GPU，可以运行脚本并指定如下 offloading 策略：
python -m flexgen.flex_opt --model facebook/opt-6.7b
➥--percent 50 50 50 50 50 50
CLI 是快速评估 FlexGen 的方式。不过，对于作为更大生态系统一部分运行推理的模型，通常更好的方式是通过程序调用 FlexGen API，并与 Hugging Face Transformers API 一起使用。例如，要使用一个无法装入 GPU 的模型运行文本补全任务，可以先导入所需 FlexGen 类，后面会在使用它们时进行说明：
from flexgen.flex_opt import (Policy, OptLM, ExecutionEnv,
➥CompressionConfig, str2bool)
接下来，从 Hugging Face Hub 下载目标模型的 tokenizer：
from transformers import AutoTokenizer
model_id = "facebook/opt-6.7b"
tokenizer = AutoTokenizer.from_pretrained(model_id, padding_side="left")
tokenizer.add_bos_token = False
stop = tokenizer("\n").input_ids[0]
在用 FlexGen 初始化模型之前，需要先设置几件事。首先，使用 ExecutionEnv 类及其 create 方法，为转换后的模型创建执行环境。这个方法接收一个目录名，FlexGen 可以把不应放进 GPU 内存或 RAM 的数据 offload 到该目录：
offload_dir = './flexgen_offload'
env = ExecutionEnv.create(offload_dir)
我们还需要设置 FlexGen 的 offloading 和压缩策略。Policy 类负责处理这两者。下面展示如何使用它来设置所需内容：
policy = Policy(len(prompts), 1,
50, 50, 50, 50, 100, 0,
overlap=True, sep_layer=True, pin_weight=True,
cpu_cache_compute=True, attn_sparsity=1.0,
compress_weight=True,
comp_weight_config=CompressionConfig(
num_bits=4, group_size=64,
group_dim=0, symmetric=False),
compress_cache=True,
comp_cache_config=CompressionConfig(
num_bits=4, group_size=64,
group_dim=2, symmetric=False)
)
创建 Policy 类实例时，我们传入 offloading 百分比，方式与 CLI 命令相同，并指定是否压缩权重及其压缩设置，是否压缩 cache 及其压缩设置。
接下来，创建目标模型的 OptLM 实例。我们会传入 Hugging Face Hub 模型 ID、执行环境实例、本地 cache 路径，也就是模型从 Hub 下载后的位置，以及配置好的策略。OptLM 类会处理下载和转换目标模型所需的所有步骤：
path = '~/opt_weights'
model = OptLM(model_id, env, path, policy)
经过 FlexGen 准备后的模型现在可以用于 few-shot 文本补全任务。设置好示例 prompt 列表后，可以按如下方式运行补全：
inputs = tokenizer(prompts, padding="max_length", max_length=128)
output_ids = model.generate(
inputs.input_ids,
do_sample=True,
temperature=0.7,
max_new_tokens=32,
stop=stop)
outputs = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
print("Outputs:\n" + 70 * '-')
for i in [0, len(outputs)-1]:
print(f"{i}: {outputs[i]}")
print("-" * 70)
为 FlexGen 准备好的 LLM 完全受 Hugging Face Transformers 库支持。当你完成文本补全生成后，可以关闭执行环境：
env.close_copy_threads()
FlexGen 也支持多 GPU，这在需要运行的模型无法装入现有硬件 VRAM 时很有帮助。
为了支持其他开源 LLM 家族，你可以创建一个新的 Python 脚本，遵循 flexgen.opt_config.py 的结构。在其中指定目标模型架构配置，并实现从 Hub 下载其权重的工具。
FlexGen 开箱即用地提供高级 offloading 和压缩能力，并且在大 batch 上提供比本书到目前为止介绍的其他方法，例如 Accelerate 和 DeepSpeed，更高的吞吐。但它是针对吞吐优先场景优化的，也就是推理批次包含数百万 token 的 LLM 应用。如果延迟是你的主要 KPI，那么 FlexGen 并不是在预算硬件上运行更大模型的最佳方式；你需要考虑其他优化选项。
9.2 SmoothQuant
SmoothQuant 是一种训练后量化方法，可以在保持准确率的同时提高 LLM 效率。传统量化方法通常聚焦于权重，但可能难以处理 activation outlier，尤其是在更大的模型中。这些 outlier 会在量化后损害性能。SmoothQuant 在同名论文中提出，它通过无需额外训练即可实现 8-bit 权重和 activation 量化来解决这个问题。它通过数学变换平滑 activation outlier，实际上是把量化困难从 activations 转移到 weights 上。这使得一系列 LLM 架构都可以被成功量化。根据论文实验，SmoothQuant 可以带来最高 1.56 倍的速度提升，以及 50% 的内存使用降低，同时准确率损失很小。这可能使极大模型可以被部署，甚至在单节点上部署，具体取决于节点硬件规格。SmoothQuant 与 PyTorch 框架兼容。
在动手实践 SmoothQuant 之前，先看看它要解决的主要量化挑战。第 6 章描述的技术对大约 60 亿或 70 亿参数以内的 LLM 效果很好。超过这个阈值后，通常会出现系统性的 activation outlier，从而导致量化后准确率下降。在 6.2.2 节中，我介绍了 LLM.int8() 技术，它通过混合精度分解缓解这个问题，也就是将 outlier 保持在 FP16，同时对其余 activations 使用 INT8。该技术已经集成到流行 Python 库中，包括本书中使用的一些库。但对于更大模型，这种分解很难在硬件加速器上高效实现。在这些情况下，更好的选择是使用另一种无需训练的量化方案，使所有计算密集型操作都使用 INT8。SmoothQuant 通过应用数学等价的逐通道缩放变换，将量化难点从 activations 转移到 weights 上，从而平滑各通道幅度，使模型更加量化友好。图 9.3 来自原始 SmoothQuant 论文，说明了这个思想。
图 9.3 —— SmoothQuant 的高层解释
图 9.3 上半部分展示了一个具有巨大 activation outlier 的 LLM 中 activation 左侧和 weight 右侧的分布。这个 outlier 拉伸了量化范围，使量化变得困难。图 9.3 下半部分展示了同一模型在 SmoothQuant 离线，也就是推理前，将尺度方差从 activations 转移到 weights 之后的情况。activations 被平滑，weights 被调整，使两者都更容易量化。SmoothQuant 可以在量化后保持准确率，同时降低 serving 成本。
来看一个例子；完整代码在配套 Colab notebook 中。我们会看到，对于给定 LLM，SmoothQuant 可以同时对 activations 和 weights 使用 INT8，并在同一模型上达到与 FP16 相当的准确率。SmoothQuant 已在多个 LLM 家族上测试，包括 Llama、OPT、Mistral、Mixtral 和 Falcon。对于 PyTorch，SmoothQuant 使用 CUTLASS INT8 GEMM kernel 提供 INT8 推理，这些 kernel 被包装为 torch-int Python 包中的 PyTorch 模块。SmoothQuant 也已经集成到一些推理引擎中，例如 Intel Neural Compressor、NVIDIA TensorRT-LLM 和 Amazon SageMaker。本书第 11 章会介绍 LLM 推理引擎。为了聚焦 SmoothQuant 在代码中的工作方式，我们会使用 FP16 模拟权重和 activations 的 8-bit 动态逐 tensor 量化。本节所有示例都需要硬件加速。
第一个目标模型是 Facebook 的 OPT-2.7B。我们先从一个小于 60 亿参数的模型开始，看看 SmoothQuant 策略是否对小模型也有帮助。
首先从源码安装 SmoothQuant：
pip install git+https://github.com/mit-han-lab/smoothquant.git
接下来，从 Hugging Face Hub 下载模型的 FP16 权重和 tokenizer：
import torch
from transformers.models.opt.modeling_opt import OPTForCausalLM
from transformers import GPT2Tokenizer
model_id = 'facebook/opt-2.7b'
tokenizer = GPT2Tokenizer.from_pretrained(model_id)
model_fp16 = OPTForCausalLM.from_pretrained(model_id,
torch_dtype=torch.float16,
device_map='auto',
)
现在加载一个数据集，用于评估目标模型的不同版本：
from datasets import load_dataset
dataset = load_dataset('cimec/lambada', split='validation[:1000]')
在这个示例中，我们使用 Hugging Face Hub 上 LAMBADA 数据集的前 1000 个样本。LAMBADA 是一组叙事段落，用于通过词预测任务评估文本理解。该数据集有两列：text，叙事段落；domain，叙事类型标签。在该数据集上评估 FP16 模型之后，例如使用第 4.1.4 节介绍的 LM-Eval 或其他工具，可以确定模型准确率：
Original model (fp16) accuracy: 0.759
现在对这个模型进行量化。我们会使用从 SmoothQuant repo 改写而来的自定义函数。
清单 9.1 一个将 LLM 权重和 activations 量化为 INT8 的自定义函数
from transformers.models.opt.modeling_opt import OPTAttention,
➥OPTDecoderLayer, OPTForCausalLM
from smoothquant.fake_quant import W8A8Linear
def quantize_model(model, weight_quant='per_tensor',
➥act_quant='per_tensor', quantize_bmm_input=True):
for name, m in model.model.named_modules():
if isinstance(m, OPTDecoderLayer):
m.fc1 = W8A8Linear.from_float(m.fc1, weight_quant=weight_quant,
➥act_quant=act_quant)
m.fc2 = W8A8Linear.from_float(m.fc2, weight_quant=weight_quant,
➥act_quant=act_quant)
elif isinstance(m, OPTAttention):
m.q_proj = W8A8Linear.from_float(
m.q_proj, weight_quant=weight_quant, act_quant=act_quant,
➥quantize_output=quantize_bmm_input)
m.k_proj = W8A8Linear.from_float(
m.k_proj, weight_quant=weight_quant, act_quant=act_quant,
➥quantize_output=quantize_bmm_input)
m.v_proj = W8A8Linear.from_float(
m.v_proj, weight_quant=weight_quant, act_quant=act_quant,
➥quantize_output=quantize_bmm_input)
m.out_proj = W8A8Linear.from_float(m.out_proj,
➥weight_quant=weight_quant, act_quant=act_quant)
return model
这个函数会把 self-attention 和 activation 中的线性层替换为 SmoothQuant 的 W8A8 实现，也就是 8-bit weight、8-bit activation。
现在可以将前面的函数应用到原始模型：
model_w8a8 = quantize_model(model_fp16)
将这里展示的原始 FP16 模型结构：
OPTForCausalLM(
(model): OPTModel(
(decoder): OPTDecoder(
(embed_tokens): Embedding(50272, 2560, padding_idx=1)
(embed_positions): OPTLearnedPositionalEmbedding(2050, 2560)
(final_layer_norm): LayerNorm((2560,), eps=1e-05,
➥elementwise_affine=True)
(layers): ModuleList(
(0-31): 32 x OPTDecoderLayer(
(self_attn): OPTSdpaAttention(
(k_proj): Linear(in_features=2560, out_features=2560, bias=True)
(v_proj): Linear(in_features=2560, out_features=2560, bias=True)
(q_proj): Linear(in_features=2560, out_features=2560, bias=True)
(out_proj): Linear(in_features=2560, out_features=2560,
➥bias=True)
)
(activation_fn): ReLU()
(self_attn_layer_norm): LayerNorm((2560,), eps=1e-05,
➥elementwise_affine=True)
(fc1): Linear(in_features=2560, out_features=10240, bias=True)
(fc2): Linear(in_features=10240, out_features=2560, bias=True)
(final_layer_norm): LayerNorm((2560,), eps=1e-05,
➥elementwise_affine=True)
)
)
)
)
(lm_head): Linear(in_features=2560, out_features=50272, bias=False)
)
与 8-bit 量化版本进行比较，可以验证 self-attention 和 activation 层是否被正确替换：
OPTForCausalLM(
(model): OPTModel(
(decoder): OPTDecoder(
(embed_tokens): Embedding(50272, 2560, padding_idx=1)
(embed_positions): OPTLearnedPositionalEmbedding(2050, 2560)
(final_layer_norm): LayerNorm((2560,), eps=1e-05,
➥ elementwise_affine=True)
(layers): ModuleList(
(0-31): 32 x OPTDecoderLayer(
(self_attn): OPTSdpaAttention(
(k_proj): W8A8Linear(2560, 2560, bias=True,
➥weight_quant=per_tensor, act_quant=per_tensor,
➥output_quant=per_tensor)
(v_proj): W8A8Linear(2560, 2560, bias=True,
➥weight_quant=per_tensor, act_quant=per_tensor,
voutput_quant=per_tensor)
(q_proj): W8A8Linear(2560, 2560, bias=True,
➥weight_quant=per_tensor, act_quant=per_tensor,
➥output_quant=per_tensor)
(out_proj): W8A8Linear(2560, 2560, bias=True,
➥weight_quant=per_tensor, act_quant=per_tensor,
voutput_quant=None)
)
(activation_fn): ReLU()
(self_attn_layer_norm): LayerNorm((2560,), eps=1e-05,
➥elementwise_affine=True)
(fc1): W8A8Linear(2560, 10240, bias=True,
➥weight_quant=per_tensor, act_quant=per_tensor,
➥output_quant=None)
(fc2): W8A8Linear(10240, 2560, bias=True,
➥weight_quant=per_tensor, act_quant=per_tensor,
➥output_quant=None)
(final_layer_norm): LayerNorm((2560,), eps=1e-05,
➥elementwise_affine=True)
)
)
)
)
(lm_head): Linear(in_features=2560, out_features=50272, bias=False)
)
如果在同样的 1000 个 LAMBADA 样本上对量化模型重复评估，可以看到准确率仍然与原始 FP16 版本相当：
Naive W8A8 quantized model accuracy: 0.758
这是预期结果：模型足够大，但其 activations 中还没有包含大的 outlier。
到目前为止，我们在没有 smoothing 模型 activations 的情况下进行了量化。现在先进行 smoothing，再量化：
act_scales = torch.load('./act_scales/opt-2.7b.pt')
smooth_lm(model_fp16, act_scales, 0.5)
model_smoothquant_w8a8 = quantize_model(model_fp16)
smooth_lm 函数接收原始模型、模型 scales，以及 alpha 作为输入。模型 scales 可在 Hugging Face Hub 上获取，适用于该模型以及许多其他模型；alpha 是一个超参数，用于控制迁移强度，也就是把多少量化困难从 activations 转移到 weights 上。SmoothQuant 论文将给定 activation Xj 和 weight Wj 的 smoothing factor sj 定义为 alpha 的函数：
图：SmoothQuant smoothing factor 公式
对于这个示例，以及一般小于 1000 亿参数的模型，我们可以保持 alpha = 0.5。
如果在与前面版本相同的数据集上重新评估这个模型，可以看到准确率保持不变，并且与原始 FP16 模型相当：
SmoothQuant W8A8 quantized model accuracy: 0.759
看看更大模型会怎样——尝试 OPT-6.7B。对其 FP16 版本重复同样评估，可以观察到如下准确率：
Original model (fp16) accuracy: 0.798
在没有 smoothing 的情况下进行 8-bit 量化后，该模型在同一评估数据集上的准确率急剧下降：
Naive W8A8 quantized model accuracy: 0.423
这证实该模型 activations 中至少存在一个大的 outlier。在量化前应用 activation smoothing 后，同一评估下的准确率与 FP16 相当：
SmoothQuant W8A8 quantized model accuracy: 0.799
正如你所看到的，SmoothQuant 是一种强大的策略，能够在不影响模型性能的情况下降低计算成本，尤其是在专门化模型超过 60 亿到 70 亿参数之后。如前所述，多个 LLM 家族开箱即用地受到支持；其他模型则需要一些代码定制，以便将 smoothing 和 quantization 步骤适配到对应架构上。
9.3 BitNet
BitNet 是一种可扩展的 1-bit Transformer 架构，旨在降低大型语言模型的成本和能耗。它用一个名为 BitLinear 的组件替代传统线性层，从而允许从零开始训练 1-bit 权重。实验结果表明，BitNet 首次在论文《BitNet: Scaling 1-bit Transformers for Large Language Models》中提出，与领先的 8-bit 量化方法和 FP16 Transformer benchmark 相比，它在语言建模性能上具有竞争力，同时内存使用和能耗显著更低。BitNet 为了保持训练稳定性，会让 optimizer states 和 gradients 保持高精度，并且可以扩展到更大模型，遵循与全精度 Transformer 类似的 scaling law。论文还描述了用于模型并行的 group quantization，可以在没有额外通信开销的情况下提升效率。总体来看，BitNet 有潜力降低推理成本并提升性能，是其他量化方法的有前景替代方案。
图 9.4 来自原始 BitNet 论文，展示了 BitLinear 组件左侧，以及使用 BitLinear 作为 attention 层的 BitNet 架构右侧。可以看到，BitNet 架构与 Transformer 具有相同布局。主要区别是，它在 attention 层中使用 BitLinear，而不是矩阵乘法。这会将 attention weights 二值化，也就是 1-bit，而模型其余部分使用更高精度，也就是 8-bit。这样可以降低计算成本，同时保留 input/output embeddings 的精度，因为模型仍然使用高精度概率进行采样。
图 9.4 —— BitLinear 组件左侧和 BitNet 架构右侧
BitNet 论文还报告了多项对比实验。图 9.5 比较了 BitNet 和 FP16 Transformer 的 scaling curve。
图 9.5 —— BitNet 与 FP16 Transformer 的 scaling curve 比较
对于较小模型，传统 FP16 Transformer 的 loss 低于 BitNet。随着模型规模增加，两者差距缩小：在约 300 亿参数时 loss 差距已经很小，在接近 1000 亿参数时降至 0.09。
训练低 bit Transformer 的最大挑战是优化稳定性。论文作者通过使用不同峰值学习率训练一系列模型，对 BitNet 和 FP16 baseline 进行了测试。图 9.6 展示了结果。
图 9.6 左侧显示，BitNet 在高学习率下可以收敛，而 FP16 Transformer 无法收敛，说明 BitNet 训练稳定性更好。右侧显示，BitNet 在不同学习率下训练都可以收敛，而且随着学习率增加，perplexity 会改善。
图 9.6 —— BitNet 和 FP16 Transformer 稳定性测试结果
BitNet 论文还将该架构与其他训练后量化，也就是 PTQ 方法进行了比较，包括 Absmax、SmoothQuant、GPTQ 和 QuIP。QuIP 是一种 2-bit 方法，见论文《QuIP: 2-Bit Quantization of Large Language Models With Guarantees》。图 9.7 来自 BitNet 论文，对比了 FP16 Transformer 与同一模型在使用不同技术进行 PTQ 后，包括 BitNet，在 zero-shot 左侧和 few-shot 右侧下游任务中的准确率。
图 9.7 —— vanilla FP16 Transformer、BitNet 及多种 PTQ 方法在 zero-shot a 和 few-shot b 设置下的下游任务平均准确率
这些不同技术已经在多个数据集上评估；图 9.8 中的表格来自 BitNet 论文，总结了 zero-shot 结果。
图 9.7 和图 9.8 中的图表与表格显示，BitNet 的平均表现，包括准确率和 perplexity，略低于 FP16 Transformer baseline，但明显优于任何其他 PTQ 场景。
图 9.8 —— BitNet 和 baseline Transformer，FP16，以及应用多种量化技术后的 zero-shot 结果
9.4 BitNet 与 Python
BitNet 框架基于第 6 章介绍过的 llama.cpp，其核心由 C++ 实现。它为 CPU 上的 1-bit 模型提供优化 kernel，用于快速、节能且无损的推理。它目前还不支持 GPU，但已有计划支持。你可以从源码安装 BitNet；请按照 Microsoft 官方 GitHub 仓库中针对你操作系统的说明操作。
BitNet 也提供 Python binding，因此你可以在 Python 中运行托管在 Hugging Face Hub 上的 BitNet 量化模型推理。主 Python 脚本 run_inference.py 使用如下语法：
python run_inference.py -m models/
➥Falcon3-7B-Instruct-1.58bit/ggml-model-i2_s.gguf -p "You are a helpful
➥assistant" -cnv
在基础形式下，它接收目标模型在 Hugging Face Hub 中的 ID、prompt，以及用于启用或禁用 chat mode 的标志作为参数，chat mode 只对 instructed model 有效。你也可以传入其他参数，用于设置 temperature、最大生成 token 数、要使用的线程数，以及 prompt 上下文大小。截至本文写作时，BitNet 支持其原生模型，0.7B 和 3.3B 参数，以及 Llama 和 Falcon 家族。如果你不想使用 C++ 并自行移植预训练模型，可以使用一个纯 Python 替代方案：带 CUDA 支持的 PyTorch port。源码可在 GitHub 上获得，也发布在 PyPI 上，可以从那里安装：
pip install bitnet
这个包提供了 BitLinear、BitNet Transformer、BitNet Attention 和 FeedForward 的 PyTorch 实现。
使用这个 API，你可以按如下方式创建一个简单的 BitNet Transformer：
from bitnet import BitNetTransformer
bitnet = BitNetTransformer(
num_tokens=20000,
dim=1024,
depth=6,
heads=8,
ff_mult=4,
)
我们使用 BitNetTransformer 类，并指定唯一输入 token 数、embedding 维度、层数、attention head 数，以及 feed-forward hidden dimension multiplier。生成的模型架构如下：
BitNetTransformer(
(emb): Embedding(20000, 1024)
(transformer): Transformer(
(layers): ModuleList(
(0-5): 6 x BitMGQA(
(q_proj): BitLinear(in_features=1024, out_features=1024, bias=True)
(k_proj): BitLinear(in_features=1024, out_features=512, bias=True)
(v_proj): BitLinear(in_features=1024, out_features=512, bias=True)
(norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)
(out_proj): BitLinear(in_features=512, out_features=1024, bias=True)
)
)
(ffn_layers): ModuleList(
(0-5): 6 x BitFeedForward(
(ff): Sequential(
(0): Sequential(
(0): BitLinear(in_features=1024, out_features=4096, bias=True)
(1): SiLU()
)
(1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)
(2): Dropout(p=0.1, inplace=False)
(3): BitLinear(in_features=4096, out_features=1024, bias=True)
)
)
)
)
(to_logits): Sequential(
(0): RMSNorm()
(1): Linear(in_features=1024, out_features=20000, bias=True)
)
)
我们可以很容易识别出这个包提供的所有 PyTorch 实现。该包集成了 Hugging Face Transformers 库，并支持托管在 Hugging Face Hub 上的预训练 LLM，而不只是原始 BitNet 框架支持的少数模型家族。我们用第 8.2 节讨论过的 ProtGPT2 试一下。你可以使用 Transformers API 从 Hugging Face Hub 下载预训练模型和 tokenizer：
from transformers import AutoTokenizer, AutoModelForCausalLM
model_id = 'nferruz/ProtGPT2'
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id)
接下来，使用 BitNet API 替换那些有 BitNet 实现的线性层：
from bitnet import replace_linears_in_hf
replace_linears_in_hf(model)
无论预训练 LLM 架构是什么，只要模型通过 Transformers API 使用，我们总是调用 replace_linears_in_hf。下面可以看到，head linear layers 已经被替换：
GPT2LMHeadModel(
(transformer): GPT2Model(
(wte): Embedding(50257, 1280)
(wpe): Embedding(1024, 1280)
(drop): Dropout(p=0.1, inplace=False)
(h): ModuleList(
(0-35): 36 x GPT2Block(
(ln_1): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)
(attn): GPT2SdpaAttention(
(c_attn): Conv1D()
(c_proj): Conv1D()
(attn_dropout): Dropout(p=0.1, inplace=False)
(resid_dropout): Dropout(p=0.1, inplace=False)
)
(ln_2): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)
(mlp): GPT2MLP(
(c_fc): Conv1D()
(c_proj): Conv1D()
(act): NewGELUActivation()
(dropout): Dropout(p=0.1, inplace=False)
)
)
)
(ln_f): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)
)
(lm_head): BitLinear(in_features=1280, out_features=50257, bias=False)
)
现在可以像往常一样使用 Transformers API 运行推理，包括使用 pipeline：
from transformers import pipeline
pipe = pipeline(
"text-generation",
model=model,
tokenizer=tokenizer,
)
sequences = pipe("<|endoftext|>",
max_length=100,
do_sample=True,
top_k=950,
repetition_penalty=1.2,
num_return_sequences=10,
eos_token_id=0)
Python BitNet 包还为 PyTorch 模型提供了 drop-in replacement：
import torch
from torch import nn
from bitnet import replace_linears_in_pytorch_model
model = nn.Sequential(
nn.Linear(10, 20),
nn.ReLU(),
nn.Linear(20, 30),
)
replace_linears_in_pytorch_model(model)
在前面的示例中，所有 PyTorch 线性层都已经被 BitNet 线性层替换。
也可以使用已有 PyTorch checkpoint 运行推理：
from bitnet import BitNetInference
bitnet = BitNetInference()
bitnet.load_model("../model_checkpoint.pth")
output_str = bitnet.generate("The dog jumped over the ", 512)
当存在计算或可持续性约束时，BitNet 可能是其他量化策略的可行替代方案。
到这里，我们对高级模型量化技术的介绍就结束了。对于许多使用本地小型语言模型的生产用例，传统量化方法已经很合适。但对于涉及更复杂、非结构化内容的小众任务，你可能会想使用这些高级策略中的一种。
总结
- FlexGen 将模型权重、activations 和 attention cache 分布在 GPU、CPU 和磁盘之间，使大模型可以在受限硬件上推理。
- FlexGen 使用线性规划，在不同存储类型之间找到高效存储和访问 tensor 的模式。
- FlexGen 将权重和 attention cache 压缩到 4 bit，以最小准确率损失提升吞吐。
- 在 FlexGen 中，
Policy类设置权重、cache 和 activations 的 offloading 百分比与压缩方式。 - FlexGen 优先考虑吞吐而不是延迟，因此最适合包含数百万 token 的批处理。
- SmoothQuant 通过将量化难点从 activations 转移到 weights，解决大语言模型中的 activation outlier 问题。
- SmoothQuant 使用数学变换，在量化前平滑跨通道的 activation 幅度。
- 超过 60 亿到 70 亿参数的模型会出现系统性 activation outlier，从而降低量化性能。
- SmoothQuant 在大型模型中进行 8-bit 权重和 activation 量化时可以保持准确率。
- SmoothQuant 会替换标准线性层，以实现 8-bit 权重和 activation 量化。
- BitNet 引入 BitLinear 组件，用 1-bit 权重替代传统线性层。
- BitNet 在 attention 层使用 1-bit 权重，同时让 optimizer states 和 gradients 保持高精度。
- BitLinear 组件可以降低计算成本，同时保留 input/output embeddings 的精度。
- 在高学习率下，BitNet 比传统 FP16 Transformer 训练更稳定，收敛更可靠。
- BitNet 在使用远少于 8-bit 量化的内存和能耗时，仍能取得有竞争力的性能。
- Python BitNet 实现为现有模型中的线性层提供了 PyTorch 兼容替代。
