使用 QAT 微调 Llama3

量化感知训练 （QAT） 是用户量化 模型，而不会导致准确性或困惑度的显著下降。在这个 教程中，我们将介绍如何在微调过程中应用 QAT，量化 生成的模型，然后使用 Torchtune 评估您的量化模型。

您将学到什么

• 什么是 QAT 以及它如何帮助减少量化退化

• 如何在 torchtune 中微调期间运行 QAT

• 连接 QAT、量化和评估配方的端到端示例

先决条件

• 熟悉 torchtune

• 确保安装 torchtune

• 确保您已下载 Llama3-8B 模型权重

什么是 QAT？

量化感知训练 （QAT） 是指在 训练或微调，最终目标是最终产生更高的质量 量化模型与简单的训练后量化 （PTQ） 的比较。在 QAT 期间， 权重和/或激活是 “假量化的”，这意味着它们被转换 就像它们被量化一样，但保持原始数据类型（例如 bfloat16） 而无需实际转换为较低的位宽。因此，假量化允许 model 来调整量化噪声，因此进行训练 过程“意识到”模型在训练后最终会被量化。

# PTQ: x_q is quantized and cast to int8
# scale and zero point (zp) refer to parameters used to quantize x_float
# qmin and qmax refer to the range of quantized values
x_q = (x_float / scale + zp).round().clamp(qmin, qmax).cast(int8)

# QAT: x_fq is still in float
# Fake quantize simulates the numerics of quantize + dequantize
x_fq = (x_float / scale + zp).round().clamp(qmin, qmax)
x_fq = (x_fq - zp) * scale

QAT 通常涉及在训练之前和之后对模型进行转换。 例如，在 torchao QAT 实现中， 这些表示为 AND 步骤：（1） 插入假量化 操作转换为线性层，以及 （2） 将假量化操作 在训练后对操作进行实际量化和去量化操作，从而产生一个量化的 模型（Dequantize 操作通常在降低后与 Linear 融合）。 在这两个步骤之间，训练可以完全像以前一样进行。prepare()convert()prepare()convert()

将 QAT 应用于 Llama3 模型

我们可以轻松地将上述 QAT 转换应用于 torchtune 中的 Llama3 以进行微调：

from torchtune.training.quantization import Int8DynActInt4WeightQATQuantizer
from torchtune.models.llama3 import llama3_8b

model = llama3_8b()

# Quantizer for int8 dynamic per token activations +
# int4 grouped per channel weights, only for linear layers
quantizer = Int8DynActInt4WeightQATQuantizer()

# Insert "fake quantize" operations into linear layers.
# These operations simulate quantization numerics during
# fine-tuning without performing any dtype casting
prepared_model = quantizer.prepare(model)

如果我们打印模型，我们将看到所有线性层都已与 交换，它模拟了 int8 的数值 动态每个令牌激活 + 按通道权重分组的 int4。现在模型 已准备好进行微调。Int8DynActInt4WeightQATLinear

>>> print(model.layers[0].attn)
MultiHeadAttention(
  (q_proj): Linear(in_features=4096, out_features=4096, bias=False)
  (k_proj): Linear(in_features=4096, out_features=1024, bias=False)
  (v_proj): Linear(in_features=4096, out_features=1024, bias=False)
  (output_proj): Linear(in_features=4096, out_features=4096, bias=False)
  (pos_embeddings): RotaryPositionalEmbeddings()
)

>>> print(prepared_model.layers[0].attn)
MultiHeadAttention(
  (q_proj): Int8DynActInt4WeightQATLinear(in_features=4096, out_features=4096, bias=False)
  (k_proj): Int8DynActInt4WeightQATLinear(in_features=4096, out_features=1024, bias=False)
  (v_proj): Int8DynActInt4WeightQATLinear(in_features=4096, out_features=1024, bias=False)
  (output_proj): Int8DynActInt4WeightQATLinear(in_features=4096, out_features=4096, bias=False)
  (pos_embeddings): RotaryPositionalEmbeddings()
)

微调后，我们可以将模型转换为实际的量化模型。 如果我们打印转换后的模型，我们将看到 QAT 线性已经 与 Int8DynActInt4WeightLinear 交换，这是量化版本 线性层。然后可以将这个量化模型保存到 checkpoint 和 用于推理或生成。

# Fine-tune as before
train_loop(prepared_model)

# Convert fake quantize to actual quantize operations
converted_model = quantizer.convert(prepared_model)

>>> print(converted_model.layers[0].attn)
MultiHeadAttention(
  (q_proj): Int8DynActInt4WeightLinear()
  (k_proj): Int8DynActInt4WeightLinear()
  (v_proj): Int8DynActInt4WeightLinear()
  (output_proj): Int8DynActInt4WeightLinear()
  (pos_embeddings): RotaryPositionalEmbeddings()
)

torchtune 中的 QAT 微调配方

综上所述，我们现在可以使用 torchtune 的 QAT 配方微调模型。 确保您首先通过以下方式下载了 Llama3 权重和分词器 按照这些说明进行操作。在本教程中， 我们使用以下设置来演示 QAT 在恢复方面的有效性 量化降级与直接量化微调模型相比 无 QAT。您可以复制默认 QAT 配置并进行以下操作 相应地修改：

tune cp llama3/8B_qat_full custom_8B_qat_full.yaml

# Dataset
dataset:
  _component_: torchtune.datasets.text_completion_dataset
  source: allenai/c4
  max_seq_len: 8192
  column: text
  name: en
  split: train
seed: null
shuffle: True

...

epochs: 1
max_steps_per_epoch: 2000
fake_quant_after_n_steps: 1000
memory_efficient_fsdp_wrap: False

根据经验，我们观察到禁用前 N 步的假量化 导致了更好的结果，大概是因为这样做可以让权重稳定下来 在我们开始将量化噪声引入微调过程之前。 因此，我们在这里禁用前 1000 步的假量化。

然后，您可以使用以下命令通过上述命令通过 QAT 运行微调 config。此工作负载至少需要 6 个 GPU，每个 GPU 的 VRAM 至少为 80GB。 默认情况下，这将使用每个令牌激活的 int8 动态 + 每个令牌激活的 int4 分组 Channel Weights 量化配置，如上所示：

tune run --nnodes 1 --nproc_per_node 6 qat_distributed --config custom_8B_qat_full.yaml

注意

确保指向您的 Llama3 权重和分词器的位置。这是可以做到的 通过添加或直接修改文件。有关如何轻松克隆和修改 torchtune 配置的更多详细信息，请参阅我们的 All About Configs。checkpointer.checkpoint_files=[my_model_checkpoint_path] tokenizer_checkpoint=my_tokenizer_checkpoint_path8B_qat_full.yaml

注意

与常规微调相比，QAT 引入了内存和计算开销， 由于假量化从根本上涉及额外的操作并且需要克隆 权重，以避免在计算假量化值时改变它们。 一般来说，我们预计 30% 左右的模型的微调速度会降低 美洲驼 3-8B。使用激活检查点时，每 GPU 最小（每个 GPU < 5GB）。

量化 QAT 模型

请注意，上面的 QAT 配方会生成一个未量化的 bfloat16 模型。模型 结构与通过定期完全微调生成的模型完全相同 没有 QAT，只是具有不同的权重。要实际获得量化模型， 复制 quantization 配置并对其进行以下修改：

tune cp quantization custom_quantization.yaml

# Model arguments
model:
  _component_: torchtune.models.llama3.llama3_8b

checkpointer:
  _component_: torchtune.training.FullModelMetaCheckpointer
  checkpoint_dir: <your QAT checkpoint dir>
  checkpoint_files: [meta_model_0.pt]
  recipe_checkpoint: null
  output_dir: <your QAT checkpoint dir>
  model_type: LLAMA3

...

quantizer:
  _component_: torchtune.training.quantization.Int8DynActInt4WeightQATQuantizer
  groupsize: 256

以下命令在 QAT 流中执行 convert 步骤，该步骤实际上 将 float 模型量化为具有量化权重的模型：

tune run quantize --config custom_quantization.yaml

注意

确保使用您用于微调模型的相同 QAT 量化器。 否则，数值将关闭，量化模型的性能将很差。

评估量化模型

现在我们有一个量化模型，我们可以对它运行一些评估并比较 与没有 QAT 的常规微调（即 训练后量化）相比的结果。 为此，我们使用了 torchtune 中集成的 EleutherAI 评估工具。首先，复制评估配置并进行以下更改：

tune cp eleuther_evaluation custom_eleuther_evaluation.yaml

# Model arguments
model:
  _component_: torchtune.models.llama3.llama3_8b

checkpointer:
  _component_: torchtune.training.FullModelTorchTuneCheckpointer
  checkpoint_dir: <your quantized model checkpoint dir>
  checkpoint_files: [meta_model_0-8da4w.pt]
  recipe_checkpoint: null
  output_dir: <your quantized model checkpoint dir>
  model_type: LLAMA3

...

# EleutherAI specific eval args
tasks: ["hellaswag", "wikitext"]
limit: null
max_seq_length: 8192
batch_size: 8

quantizer:
  _component_: torchtune.training.quantization.Int8DynActInt4WeightQuantizer
  groupsize: 256

注意

由于我们传入的是量化模型，因此请务必使用相应的 后训练量化器，而不是 QAT 量化器。例如，如果你 在微调过程中，您 应在此步骤中指定。有关支持的量化器的完整列表，请参阅 quantization recipe。Int8DynActInt4WeightQATQuantizerInt8DynActInt4WeightQuantizer

现在运行评估配方：

tune run eleuther_eval --config my_eleuther_evaluation.yaml

结果应如下所示：

# QAT quantized model evaluation results (int8 activations + int4 weights)

|  Tasks  |Version|Filter|n-shot|    Metric     |Value |   |Stderr|
|---------|------:|------|-----:|---------------|-----:|---|------|
|wikitext |      2|none  |     0|word_perplexity|9.9148|±  |N/A   |
|         |       |none  |     0|byte_perplexity|1.5357|±  |N/A   |
|         |       |none  |     0|bits_per_byte  |0.6189|±  |N/A   |
|hellaswag|      1|none  |     0|acc            |0.5687|±  |0.0049|
|         |       |none  |     0|acc_norm       |0.7536|±  |0.0043|

将这些结果与没有 QAT 微调的模型进行比较，我们可以看到 QAT 能够恢复很大一部分量化退化 与 PTQ 相比。例如，归一化 使用 PTQ 时，hellaswag 任务的准确率下降了 2.20%，但使用 PTQ 时仅下降了 0.74% QAT 与原始未量化模型相比。同样，单词困惑 在 wikitext 任务中，PTQ 增加了 2.048，但 QAT 仅增加了 1.190（较低 更好）。

# PTQ quantized model evaluation results (int8 activations + int4 weights)

|  Tasks  |Version|Filter|n-shot|    Metric     | Value |   |Stderr|
|---------|------:|------|-----:|---------------|------:|---|------|
|wikitext |      2|none  |     0|word_perplexity|10.7735|±  |N/A   |
|         |       |none  |     0|byte_perplexity| 1.5598|±  |N/A   |
|         |       |none  |     0|bits_per_byte  | 0.6413|±  |N/A   |
|hellaswag|      1|none  |     0|acc            | 0.5481|±  |0.0050|
|         |       |none  |     0|acc_norm       | 0.7390|±  |0.0044|

# Float model evaluation results (bfloat16)

|  Tasks  |Version|Filter|n-shot|    Metric     |Value |   |Stderr|
|---------|------:|------|-----:|---------------|-----:|---|------|
|wikitext |      2|none  |     0|word_perplexity|8.7251|±  |N/A   |
|         |       |none  |     0|byte_perplexity|1.4994|±  |N/A   |
|         |       |none  |     0|bits_per_byte  |0.5844|±  |N/A   |
|hellaswag|      1|none  |     0|acc            |0.5740|±  |0.0049|
|         |       |none  |     0|acc_norm       |0.7610|±  |0.0043|

因此，QAT 流程产生了一个量化模型，其性能优于后训练 量化模型。重要的是，量化模型结构在两者中是相同的 flows，以及模型大小、内存使用情况和所有其他性能 特性也是一样的。

请注意，尽管权重被量化为 int4，但量化的模型大小 对于 QAT 和 PTQ 流，流量均为 8.187 GB，而原始浮点模型 为 14.958 GB。这是因为这个量化器使用 int8 来表示权重 因为 PyTorch 没有原生 int4 dtype 支持。更高效的表示 是打包 int4 权重，这将使量化模型大小减半。这是 Int4WeightOnlyQuantizer 的作用，以及相应的 QAT 量化器将 在将来添加。

将 QAT 模型降低到设备（可选）

量化模型的一个重要动机是能够在 resource 中运行它 受约束的环境。您可以进一步将 QAT Llama3 模型降低到边缘设备 例如使用 ExecuTorch 的智能手机 按照这些说明进行操作。 例如，以下命令将模型降低到 XNNPACK 后端：

python -m examples.models.llama2.export_llama --checkpoint <your QAT checkpoint> -p <params.json> -kv --use_sdpa_with_kv_cache -X -qmode 8da4w --group_size 256 -d fp32 --metadata '{"get_bos_id":128000, "get_eos_id":128001}' --embedding-quantize 4,32 --output_name="llama3_8da4w.pte"

这会产生一个大小为 3.881 GB 的更小的量化模型。在 OnePlus 12 智能手机上进行基准测试时，该模型还实现了与训练后量化模型相同的推理和生成速度。这是因为模型结构在两个流中是相同的：

	质量检查	PTQ
量化模型大小	3.881 吉字节	3.881 吉字节
推理速度	9.709 托克/秒	9.815 托克/秒
生成速度	11.316 托克/秒	11.364 托克/秒