TorchAO 中的量化感知训练 (II)

在我们之前的量化感知训练 (QAT) 博客中，我们介绍了 TorchAO 中针对使用 ExecuTorch 的边缘设备的大型语言模型的初步 QAT 流程。自那时起，我们将此流程扩展到了针对 MSLK 等快速 CUDA 内核，以实现在 vLLM 中的快速推理，并将此流程整合进了 Unsloth 和 Axolotl 等流行的微调框架中。我们还探索了更先进的 QAT 技术，例如用于低位量化的 PARQ（原型阶段）：

• Unsloth 集成：使用 INT4 QAT 恢复高达 66.9% 的精度损失，与 BF16 相比实现了 1.73 倍的推理加速。请查看我们的笔记本和 HuggingFace 检查点以获取端到端指南。
• Axolotl 集成：使用 NVFP4 QAT（原型阶段）恢复高达 71.6% 的精度损失，在 B200 GPU 上与 BF16 相比，以 1/4 的 HBM 使用量实现了 1.35 倍的推理加速。
• PARQ：（原型阶段）一种基于优化器的替代性 QAT 技术，用于低位量化。在仅使用约 58% 内存占用并实现约 1.57 倍更快的解码吞吐量的同时，使每行 3 位模型达到与每组 4 位基准模型相当的精度。

只需几行代码，即可尝试我们最新的 QAT 流程！

from torchao.quantization import quantize_, Int4WeightOnlyConfig
from torchao.quantization.qat import QATConfig

# The same config to use for Post-Training Quantization (PTQ)
base_config = Int4WeightOnlyConfig(group_size=32)


# Prepare step: model is now "fake quantized" and ready for training
quantize_(model, QATConfig(base_config, step="prepare")

train(model)

# Convert step: model is now quantized and ready for inference
quantize_(model, QATConfig(base_config, step="convert")

量化感知训练 (QAT)

减轻训练后量化 (PTQ) 造成的精度损失的一种知名技术是 QAT，它是一个可选的微调步骤，旨在使模型权重适应一种更“意识到”最终将被量化的表示形式。QAT 通过在训练过程中对权重（以及可选的激活值）进行“伪量化”来实现这一点，这意味着尽可能地模拟 PTQ 数值，然后在前向传播过程中立即反量化回高精度值，同时保持反向传播不变。

QAT 还可以与低秩自适应 (LoRA)相结合，从而获得两者的优势：在显著减少训练过程中的存储和计算需求的同时，减轻量化带来的精度损失。LoRA 是一种流行的微调技术，它通过冻结原始模型权重并仅训练一小部分新的 LoRA“适配器”（在模型规模中占比很小）来大幅减少可训练参数的数量。在训练过程中，伪量化被应用于 LoRA 适配器（图 1 中的 A 和 B）和冻结权重（以及可选的输入激活），因为在量化模型中，适配器最终会合并回权重中。该技术已被证明可将 QAT 速度提升 1.89 倍，同时将所需内存减少 36.1%。

图 1：将 QAT 与 LoRA 结合使用，用户可以在减轻量化损失的同时加快训练速度并减少训练过程中的内存占用。伪量化被动态应用于冻结权重、LoRA 适配器以及（可选的）输入激活。

请注意，QAT + LoRA 与 QLoRA 的区别在于，权重（以及可选的激活值）是在训练过程中进行伪量化的，而不是在训练前实际量化为较低位的 dtype（例如 NF4）并以该表示形式进行存储。通常，QAT 期间模拟的量化方案应尽可能与实际的训练后量化方案相匹配。

截至 TorchAO 0.16.0，我们支持以下 dtype 组合：

权重 dtype	输入激活 dtype	TorchAO 配置
INT4	FP32, BF16	Int4WeightOnlyConfig
INT4	FP8	Float8DynamicActivationInt4WeightConfig
INT4	INT8	Int8DynamicActivationIntxWeightConfig (用于边缘设备)
NVFP4	NVFP4	NVFP4DynamicActivationNVFP4WeightConfig (原型阶段)

与 Unsloth 的集成

TorchAO 的 QAT 支持已集成到 Unsloth 的完整微调工作流（训练所有模型参数）和 LoRA 微调工作流（仅训练适配器）中。在我们的初步实验中，与没有 QAT 的量化基准相比，INT4 权重仅 QAT 在 GPQA 上为 Gemma3-4B恢复了 66.9% 的精度损失，在 BBH 上为 Gemma3-12B恢复了 45.5% 的精度损失（图 2）。仅通过在微调期间应用伪量化，这便转化为 2.1% 和 1.0% 的原始精度提升（图 3）。生成的 QAT 量化模型可以用作非 QAT 量化模型的直接替代品，因为模型结构未更改，但具有更高的精度。

图 2：Unsloth 利用 TorchAO 中的 QAT，通过 INT4 QAT + LoRA 恢复了高达 66.9% 的精度损失 (来源)。QAT 恢复了 Gemma3-12B BBH 上 4.5% 损失精度的 45.5%，Gemma3-4B GPTQ 上 1.5% 损失精度的 66.9%，Qwen3-4B WikiText 上 1.11 困惑度增量的 36.3%，以及 Llama-3.2-1B MMLU Pro 上 4.5% 损失精度的 36.0%。

图 3：Unsloth 利用 TorchAO 中的 QAT，通过 INT4 QAT + LoRA 将量化和微调模型的原始精度提高了高达 2.1% (来源)。QAT 恢复了 Gemma3-12B BBH 上 4.5% 损失精度中的 2.1%，Gemma3-4B GPQA 上 1.5% 损失精度中的 1.0%，Qwen3-4B MMLU Pro 上 5.8% 损失精度中的 2.0%，以及 Llama-3.2-3B MMLU Pro 上 2.0% 损失精度中的 0.7%。

Unsloth 用户可以通过简单地指定额外的 qat_scheme 标志，在他们的微调工作流中启用 QAT。有关端到端的示例，请查看 Unsloth 的免费 QAT笔记本，或者查看我们 HuggingFace 检查点中的模型卡，它们也是使用 Unsloth 进行微调的。

from unsloth import FastLanguageModel
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = "unsloth/gemma3-12b-it",
    max_seq_length = 2048,
    load_in_16bit = True,
)
model = FastLanguageModel.get_peft_model(
    model,
    r = 16,
    target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",
                      "gate_proj", "up_proj", "down_proj",],
    lora_alpha = 32,   
    # We support fp8-int4, fp8-fp8, int4, int8-int4, phone-deployment
    qat_scheme = "int4",
)

Unsloth 还利用 TorchAO QAT 通过 ExecuTorch 将模型部署到 Pixel 8 和 iPhone 15 Pro 等智能手机上，在此过程中恢复了高达 70% 的精度损失。用户可以为此用例指定 qat_scheme = "phone-deployment" 或 "int8-int4"。有关更多详细信息，请参阅此博客。

与 Axolotl 的集成

我们还将 TorchAO 的 QAT 支持集成到了 Axolotl 的多 GPU 完整微调工作流中。Axolotl 支持多种 QAT 方案，包括 INT4 仅权重、FP8 动态激活 + INT4 权重以及 NVFP4 动态激活 + NVFP4 权重（原型阶段）。 

感谢 Lambda Labs 的支持，我们能够在由两个 B200 节点组成的 Lambda Labs 即时集群（通过节点内 NVLINK 和节点间 Infiniband 连接）上，通过训练多达 72B 参数的模型，证明了 QAT 的有效性。在我们的初步实验中，NVFP4 QAT 在指令遵循和数学推理任务及基准测试的微调与评估过程中，恢复了显著的精度损失。我们恢复了 Gemma3-12B 上高达 63.2% 的精度损失，以及 Gemma3-27B 上 71.6% 的精度损失，分别转化为 +3.2% 和 +2.3% 的原始精度提升（图 4 和 5）。

图 4：Axolotl 利用 TorchAO 中的 NVFP4 QAT 为 Gemma3-12B 恢复了高达 63.2% 的精度损失（绝对精度提升 +3.2%）。

图 5：Axolotl 利用 TorchAO 中的 NVFP4 QAT 在各种 MMLU 任务中为 Gemma3-27B 恢复了高达 71.6% 的精度损失（绝对精度提升 +2.4%）。

将 QAT 应用于 Axolotl 微调工作流非常简单。只需在配置文件中添加以下部分（参见这些示例配置）：

# my_qat_workflow.yaml
qat:
  activation_dtype: nvfp4
  weight_dtype: nvfp4
  group_size: 16 # NVFP4 only supports group size of 16 per specification

然后运行以下命令以使用相同的配置进行微调和量化。请注意，使用 NVFP4 QAT 进行微调目前仅在 Blackwell GPU 或更高版本上受支持。

axolotl train my_qat_workflow.yaml
axolotl quantize my_qat_workflow.yaml

有关更详细的说明，请参阅以下文档：

• https://docs.axolotl.org.cn/docs/qat.html
• https://docs.axolotl.org.cn/docs/quantize.html

分段仿射正则化量化 (PARQ)

PARQ 是 TorchAO 中的一项新技术，它使执行自定义低位方案的 QAT 变得简单直接（参见 2025 ICML论文）。特别是，它支持使用一种“拉伸弹性量化”函数进行 QAT，该函数使输出值在量化网格上的分布比通常的仿射网格更均匀。近期研究表明，这对于低于 4 位的量化质量至关重要。在我们的实验中，使用 PARQ 训练的每行 3 位模型达到了与每组 4 位基准模型相当的精度，同时仅使用了约 58% 的内存占用，解码吞吐量提升了约 1.57 倍。

除了简化超低位 QAT 的实验之外，PARQ 还与 ExecuTorch 无缝集成，类似于现有的 TorchAO QAT 流程，实现了从训练到使用 PyTorch 生态系统进行高效端侧部署的端到端路径。

TorchAO 集成

现有的 TorchAO QAT 通过在向前传播中交换模块来工作。PARQ 采用了一种不同的方法：它直接在优化器的 step 函数内部执行量化。用户在定义优化器的 param_groups 时，指定要量化的参数以及量化函数和粒度。因此，无需更改模型代码。

PARQ 通过应用 TorchAO 的 quantize_ API 中的配置类和量化原语来量化权重。根据配置，它可以选择基于仿射网格或适合低位 QAT 的新型拉伸网格来确定量化参数。PARQ 与现有 QAT 方法之间的这种紧密集成确保了熟悉的设置过程和数值一致性。此外，它还可以轻松地与来自同一 API 的动态激活量化相组合。

在 TorchAO 中，我们提供了简单的 API，用于配置 PARQ 优化器以匹配多种量化策略。例如，您可以将所有线性层量化为 2 位，同时将 token 嵌入层保持为 4 位。这种混合量化器和位宽的能力使得探索高级配置变得容易，并显著加快了实验速度。

import torch
from torchao.prototype.parq.api import QuantConfig, create_optimizer

# Define how to quantize linear weights from `model.parameters()`
def linear_filter_fn(module, fqn):
    return isinstance(module, torch.nn.Linear) and fqn.endswith("weight")

linear_config = QuantConfig(bitwidth=2, group_size=None)
quant_configs_and_filter_fns = [(linear_config, linear_filter_fn)]

# Apply `parq.optim.QuantOptimizer` to quantize in `optimizer.step()`
optimizer = create_optimizer(
    model,
    quant_configs_and_filter_fns,
    base_optimizer_cls=torch.optim.AdamW,
    base_optimizer_kwargs={"weight_decay": 1e-2},
    quant_per_channel=True,
)

Phi-4-mini-instruct 上的低位 QAT 及 ExecuTorch 部署

为了在现实环境中展示 PARQ，我们对微软的 Phi-4-mini-instruct 模型进行了低位微调，使用了多种量化方案，并在精度和端侧性能方面将其与 4 位 PTQ基准进行了比较。

我们发布了 HuggingFace 模型卡和脚本，用于在公共数据上重现这些结果。

• 2位：Phi-4-mini-instruct-parq-2w-4e-shared
• 3位：Phi-4-mini-instruct-parq-3w-4e-shared
• 4位：Phi-4-mini-instruct-parq-4w-4e-shared-gsm

	线性量化	嵌入层量化	微调任务	ExecuTorch *.pte 模型大小
2位 QAT	2位每行	4位每行，与 lm_head 绑定	常规对话	1.13GB
3位 QAT	3位每行	4位每行，与 lm_head 绑定	常规对话	1.53GB
4位 QAT	4位每行	4位每行，与 lm_head 绑定	小学数学问题	1.93GB
4位 PTQ	4位每组 (32)，包括 lm_head	8位每行	不适用	2.78GB

尽管使用了更粗糙的每行量化粒度，但我们的 3 位 QAT 模型在一系列推理基准测试中表现与 4 位 PTQ 模型相当（见下文）。此外，我们针对小学数学优化的 4 位 QAT 模型在此类问题上比 4 位 PTQ 基准准确率高出近 17%。

图 7a：使用 PARQ 训练的低于 4 位模型与 4 位 PTQ 和基础 bf16 模型的精度对比。

图 7b：图 7a 的续集。

图 8：PARQ 在数学应用题基准测试上的精度。

借助 ExecuTorch，我们可以在移动设备上运行上述 2、3 和 4 位模型。下面我们展示了在 iPhone 15 Pro 上运行模型时的性能数据和屏幕截图：我们看到较低的位宽带来了解码速度和内存使用方面的显著提升。实际上，导出的 3 位模型仅占用 4 位模型约 58% 的内存空间，同时解码速度提升了约 1.57 倍，尽管它在上述基准测试中达到了相似的精度。

模型	ExecuTorch 后端	内存 (GB)	解码速度 (tok/sec)
2位 QAT	TorchAO 低位内核	1.4	27
3位 QAT	TorchAO 低位内核	1.8	22
4位 QAT	TorchAO 低位内核	2.2	18
4位 PTQ	XNNPACK	3.1	14

图 9：使用和不使用 PARQ QAT 进行 2、3 和 4 位宽量化，并使用 ExecuTorch 部署在 iPhone 15 Pro 上的 Phi-4-mini-instruct 的输出示例。

展望未来

在这篇博客中，我们重点介绍了 TorchAO QAT 与 Unsloth 和 Axolotl 的两项新集成，并介绍了 PARQ，这是一种针对低位设置的、基于优化器的创新 QAT 技术。在不久的将来，我们计划探索以下方向以继续这项工作：

• 强化学习：流行的 RL 算法（如 PPO 和 GRPO）可以从即时量化输出中受益。然而，为了保持真正的策略内 (on-policy) RL 训练，需要 QAT 来匹配推理数值。探索如何将 TorchAO QAT 应用于此设置将是一个有趣的探索方向。

• 在 QAT 期间利用 GPU 内核：通过使用专为 QAT 设计的自定义内核（如 FP-Quant 中的内核）来加快训练速度存在机会，例如在该自动微分函数中，前向传播执行 MXFP4 GEMM，反向传播执行 MXFP8 GEMM。对于针对整数量化 dtypes 的 QAT 工作流，在训练期间复制整数数值还可以减少端到端工作流中潜在的数值差异。

• 进一步集成：我们还计划将 QAT 合并到更新的微调框架（如 TorchForge）中，并通过例如在 Axolotl QAT 中增加 LoRA 支持来改进我们现有的集成。

• PARQ：我们计划扩展对实验性新优化算法的支持。这些算法类似于 PARQ 的现有算法，但在非凸问题上具有更精确的收敛保证。

致谢

我们深切感谢我们的外部合作者 Daniel Han、Michael Han、Datta Nimmaturi (Unsloth)、Salman Mohammadi 和 Wing Lian (Axolotl) 为实现富有成效的 TorchAO QAT 集成所做的贡献。我们还要感谢 Lambda Labs 的 Nick Harvey 为运行我们的 NVFP4 QAT 实验提供了基础设施。最后，我们感谢所有为该项目和本博客提供宝贵反馈的人，包括 Driss Guessous、Vasiliy Kuznetsov、Supriya Rao、Mark Saroufim 和 Lin Xiao。