分布式量化感知训练 (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 的 VRAM 至少为 80GB，例如 A100 或 H100。

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

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

1. 使用上述命令或按照教程运行 qat_distributed 配方。默认情况下，这将使用 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
   

注意

我们使用配置文件来展示如何自定义这些示例中的配方。查看 配置教程 以了解更多信息。

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

• 调整模型精度。

• 使用激活检查点。

• 启用梯度累积。

• 使用较低精度优化器。

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

有兴趣了解此配方的实际应用吗？查看我们的一些教程，展示如何使用它

• 使用 QAT 微调 Llama3