分布式量化感知训练 (QAT)

QAT 可以在推理时利用量化带来的内存节省优化，同时不会显著降低模型性能。在 torchtune 中，我们使用 torchao 来实现 QAT。其工作原理是 在微调期间模拟量化数值。虽然这可能会在训练期间引入内存和计算开销，但我们的测试发现 QAT 在评估量化模型时显著降低了性能下降，同时不影响模型尺寸缩减收益。请参阅 关于 QAT 的 PyTorch 博客文章，以深入了解该技术的工作原理。

我们提供经过预测试的即用型配置，您只需两步即可使用最新的 Llama 模型运行它们

tune download meta-llama/Meta-Llama-3-8B-Instruct  \
--output-dir /tmp/Meta-Llama-3-8B-Instruct \
--ignore-patterns "original/consolidated.00.pth" \
--HF_TOKEN <HF_TOKEN>

tune run --nproc_per_node 6 qat_distributed \
--config llama3/8B_qat_full

注意

您可能需要获得您感兴趣的 Llama 模型的访问权限。有关访问门控仓库的详细信息，请参阅此处。此外，此工作负载至少需要 6 个 GPU，每个 GPU 的显存至少为 80GB，例如 A100 或 H100。

目前，QAT 的主要可调杠杆是使用延迟伪量化（delayed fake quantization）。延迟伪量化允许控制伪量化发生的步骤。经验上，最初允许模型在没有伪量化的情况下进行微调，可以使权重和激活值在进行伪量化之前稳定下来，从而可能提高量化精度。这可以通过 fake_quant_after_n_steps 指定。为了让您大致了解如何配置此参数，我们在 fake_quant_after_n_steps ~= total_steps // 2 时取得了最佳结果。

未来我们计划支持不同的量化策略。目前，请注意您至少需要 torch>=2.4.0 才能使用 Int8DynActInt4WeightQATQuantizer 策略。通常，使用 QAT 训练、量化和评估模型的流程如下：

1. 使用上述命令或遵循教程运行 qat_distributed recipe。默认情况下，这将使用 Int8DynActInt4WeightQATQuantizer。

2. 这将生成一个原始数据类型的未量化模型。要获得实际的量化模型，请在此之后运行 tune run quantize，并在配置中指定相同的量化器，例如：

   # QAT specific args
   quantizer:
     _component_: torchtune.training.quantization.Int8DynActInt4WeightQATQuantizer
     groupsize: 256
   

3. 通过指定相应的训练后量化器，评估或使用您的量化模型进行推理

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

注意

在这些示例中，我们使用配置文件来展示如何自定义 recipe。请查看configs 教程了解更多信息。

我们许多其他的内存优化功能也可以在此 recipe 中使用

• 调整模型精度。

• 使用激活检查点。

• 启用梯度累积。

• 使用低精度优化器。

您可以在我们的内存优化概述中了解更多关于我们所有内存优化功能的信息。

对看到此 recipe 的实际应用感兴趣吗？请查看我们的一些教程，了解如何使用它

• 使用 QAT 微调 Llama3