
torchtune 量化实战指南QAT 量化感知训练、QLoRA 与 torchao 集成全解析【免费下载链接】torchtunePyTorch native post-training library项目地址: https://gitcode.com/GitHub_Trending/to/torchtunetorchtune 通过集成 torchao 支持量化感知训练QAT与 QLoRA覆盖从低比特微调、量化权重导出到量化模型评估与推理的完整闭环。读完本篇你将掌握如何用tune run一条命令启动 4 卡 QAT 微调、如何用tune run quantize导出真正的量化权重、以及如何在 EleutherAI 评测与文本生成中加载并运行量化模型同时理解 quantizer 组件在源码中的真实调用链路与命名规则。支持的量化模式8da4w 与 8da4w-qattorchtune 当前只集成了部分量化技术具体支持情况以配方文件的 docstring 为准。量化配方的QuantizationRecipe类文档字符串明确列出了两种模式模式量化器PyTorch 版本要求说明8da4wtorchtune.training.quantization.Int8DynActInt4WeightQuantizer2.3int8 动态逐 token 激活量化 int4 逐组逐轴权重量化8da4w-qattorchtune.training.quantization.Int8DynActInt4WeightQATQuantizer2.4与8da4w相同但用于量化 QAT 微调产出的 checkpoint两种模式共享一个关键超参groupsize来自quantize.pydocstring 的原始说明groupsizeintint4 权重量化的分组大小指获得独立量化参数的分组维度例如 32、64、128、256。数值越小量化粒度越细、精度越高但内存开销也越大。从 torchtune/training/quantization.py 的源码可以确认这些模式的底层实现Int8DynActInt4WeightQuantizer.quantize()内部调用 torchao 的Int8DynamicActivationIntxWeightConfig(weight_dtypetorch.int4, weight_granularityPerGroup(self.groupsize))再通过quantize_(model, quantize_fn)完成权重转换。模块中还有一个_quantizer_to_mode注册表将量化器类映射到模式字符串8da4w、8da4w-qat等get_quantizer_mode()函数即基于该表返回模式名——这个模式名既用于quantize.py中判断是否走 QAT 分支也参与最终 checkpoint 文件命名见下文。此外该模块还导出了Int4WeightOnlyQuantizer模式4w走 torchao tinygemm kernel 的 int4 纯权重量化等组件但quantize.py的 docstring 表明当前量化配方主推的是上表两种模式。对于训练后量化PTQ仓库给出的建议是直接使用 torchao 完成量化并在 torchao 中做评测与基准测试而不是通过 torchtune 的量化配方——torchtune 的定位在于 QAT 与 QLoRA 微调流程。量化感知训练QAT用假量化微调出高精度模型PyTorch 2.4QAT 指在微调过程中对权重和/或激活应用假量化fake quantization只模拟量化的数学计算而不真正将原始 dtype 转换为低精度。这样做能让模型权重在训练过程中适应未来的低比特表示从而在最终量化后保持更好的精度。启动 QAT 微调的命令来自 llama3/8B_qat_full.yaml 配置头注释tune run --nproc_per_node 4 qat_distributed --config llama3/8B_qat_full该配置的完整要点如下摘自 recipes/configs/llama3/8B_qat_full.yaml# QAT arguments quantizer: _component_: torchtune.training.quantization.Int8DynActInt4WeightQATQuantizer groupsize: 256 # Fine-tuning arguments batch_size: 2 epochs: 1 dtype: bf16 optimizer: _component_: torch.optim.AdamW lr: 2e-5 fused: True loss: _component_: torchtune.modules.loss.LinearCrossEntropyLoss gradient_accumulation_steps: 1 compile: False # torch.compile the model loss, True increases speed decreases memory # Memory management enable_activation_checkpointing: True # True reduces memory enable_activation_offloading: False # True reduces memory注意compile: False不是可选项qat_distributed.py 的构造函数中会显式检查并在compileTrue时抛出ValueError(Compile is not yet supported for QAT. ...)。该配置以 Llama 3 8B Instruct 为例依赖预先下载好的权重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也支持命令行覆盖参数例如在启动时替换 checkpointer 目录tune run --nproc_per_node 4 qat_distributed --config llama3/8B_qat_full checkpointer.checkpoint_dirYOUR_CHECKPOINT_DIRQAT 配方中的关键机制从 qat_distributed.py 的 docstring 可以提炼出 QAT 配方的几个核心特性产出物是未量化模型QAT 配方最终保存的仍是原始 dtype 的未量化模型需要再用tune run quantize单独量化见下一节。延迟假量化可通过fake_quant_after_n_steps指定从第 N 步之后才开始假量化。配方注释中说明让模型先不带假量化地微调若干步使权重与激活值趋于稳定后再量化可能带来更好的量化精度。源码中该值来自cfg.get(fake_quant_after_n_steps, None)qat_distributed.py默认None即立即启用。并行与省显存基于 PyTorch FSDP API支持fsdp_cpu_offload、激活检查点enable_activation_checkpointing与激活卸载enable_activation_offloading需 PyTorch 2.5不支持 DDP不支持 CPU 训练仅支持 fp32 与 bf16 精度dtypefp16会直接抛ValueError。QAT 检查点要求由于 QAT 训练会周期性落盘评测与生成量化模型时必须使用 torchtune 格式的 checkpoint这一点直接影响后两节的配置。在模型构建环节qat_distributed.py中会执行qat_distributed.pyquantizer config.instantiate(quantizer_cfg) quantizer.precision self._dtype quantizer_mode training.quantization.get_quantizer_mode(quantizer) if qat not in quantizer_mode: raise ValueError(Quantizer mode %s is not supported for finetuning % quantizer_mode) model quantizer.prepare(model)可以看到配方会校验 quantizer 模式必须包含qat子串随后调用 torchao 量化器的prepare(model)将线性层替换为带假量化逻辑的模块。用 tune run quantize 导出真正的量化模型QAT 训练结束后运行tune run quantize并指定与训练相同的 quantizer才能得到真实的低比特权重。命令与量化配置如下来自 quantization.md 与 configs/quantization.yamltune run quantize --config quantization# QAT specific args quantizer: _component_: torchtune.training.quantization.Int8DynActInt4WeightQATQuantizer groupsize: 256完整的 quantization.yaml 默认以 Llama 2 7B 为例核心字段output_dir: /tmp/torchtune/llama2_7B/quantized # /tmp 可能被系统清理按需修改 model: _component_: torchtune.models.llama2.llama2_7b checkpointer: _component_: torchtune.training.FullModelHFCheckpointer checkpoint_dir: /tmp/Llama-2-7b-hf checkpoint_files: [ pytorch_model-00001-of-00002.bin, pytorch_model-00002-of-00002.bin, ] recipe_checkpoint: null output_dir: ${output_dir} model_type: LLAMA2 device: cuda dtype: bf16 seed: 1234 quantizer: _component_: torchtune.training.quantization.Int8DynActInt4WeightQuantizer groupsize: 256量化配方的内部执行流程quantize.py 中main的调用链非常清晰config.parse def main(cfg: DictConfig) - None: config.log_config(recipe_nameQuantizationRecipe, cfgcfg) recipe QuantizationRecipe(cfgcfg) recipe.setup(cfgcfg) # 加载 checkpoint 构建模型 recipe.quantize(cfgcfg) # 真正的量化 recipe.save_checkpoint(cfgcfg)三个环节的实现细节setup→_setup_modelquantize.py先以cfg.dtypebf16为默认 dtype 实例化模型若量化模式包含qat先调用self._quantizer.prepare(model)将模型准备为 QAT 结构然后再load_state_dict——这就是为什么 QAT checkpoint 必须用 QATQuantizer 才能正确加载最后用training.validate_expected_param_dtype校验参数 dtype 与配置一致。quantizequantize.py在torch.no_grad()下执行QAT 模式走self._quantizer.convert(self._model)把假量化结构转换为真实低比特表示PTQ 模式走self._quantizer.quantize(self._model)torchao 的quantize_路径。函数还记录量化耗时与 GPU 峰值显存。save_checkpointquantize.py输出文件名为f{file_name}-{self._quantization_mode}.rstrip(-qat)再补上.ckpt后缀。也就是说QAT 模式8da4w-qat产出的文件会以8da4w命名-qat后缀被截掉PTQ 模式则直接叫8da4w——这与后文用Int8DynActInt4WeightQuantizer加载 QAT 量化模型的说法相呼应两者最终是同一类量化结构。评测量化模型修改 eleuther_evaluation 配置量化模型的质量需要用下游任务验证。基于默认的 EleutherAI 评测配置只需做两处改动# Currently we only support torchtune checkpoints when # evaluating quantized models. For more details on checkpointing see # https://pytorch.org/torchtune/main/deep_dives/checkpointer.html # Make sure to change the default checkpointer component checkpointer: _component_: torchtune.training.FullModelTorchTuneCheckpointer .. checkpoint_files: [quantized_model_checkpoint] # Quantization specific args quantizer: _component_: torchtune.training.quantization.Int8DynActInt4WeightQuantizer groupsize: 256说明默认 eleuther_evaluation.yaml 的checkpointer是FullModelHFCheckpointer用于加载 HF 权重而量化产物是 torchtune 保存的.ckpt因此必须换成FullModelTorchTuneCheckpointer并把checkpoint_files指向量化后的 checkpointquantizer字段默认是null该文件末尾即为quantizer: null加载量化模型时必须填写文档特别指出可以用Int8DynActInt4WeightQuantizer加载 QAT 量化模型因为它与 QAT 导出结果属于同一类型量化8da4w。然后运行评测tune run eleuther_eval --config eleuther_evaluation配置中还可按需调整tasks默认[truthfulqa_mc2]配置文件头注释示例tasks[truthfulqa_mc2,hellaswag]、limit、max_seq_length默认 4096、batch_size默认 8与enable_kv_cache默认 True。生成式推理量化模型修改 generation 配置推理流程与评测一致基于默认的 generation 配置 做同样的两处改动checkpointer: _component_: torchtune.training.FullModelTorchTuneCheckpointer .. checkpoint_files: [quantized_model_checkpoint] # Quantization Arguments quantizer: _component_: torchtune.training.quantization.Int8DynActInt4WeightQuantizer groupsize: 256运行生成tune run generate --config generation该配置的其他常用字段默认以 Llama 2 7B 为例prompt.user默认 Tell me a joke.、max_new_tokens: 300、temperature: 0.6、top_k: 300、enable_kv_cache: True同样存在默认的quantizer: null需要覆盖。QLoRAQAT LoRA 的融合实现recipes/quantization.md标题中提到的 QLoRA在源码侧的落点是 torchtune/training/quantization.py 中的swap_lora_linear_with_qat函数它递归遍历模型将每个LoRALinear替换为QATLoRALinear替换后线性层的计算变为x - fake_quantize(W_frozen) fake_quantize(x) BAx即冻结的低比特主干走假量化路径LoRA 低秩分支BAx保持原精度训练。文档字符串明确其目的假量化模拟量化数值而不做真正的 dtype 转换从而让微调后的模型在最终量化时精度损失更小。仓库中对应的实战配方可参考 qat_lora_finetune_distributed.py 及各模型目录下的qat_lora/qlora_single_device配置如 llama3_2/3B_qat_lora.yaml、llama2/7B_qlora_single_device.yaml其测试覆盖见 tests/recipes/test_qat_lora_finetune_distributed.py。小结与适用前提流程主线QAT 微调qat_distributed产出未量化模型→tune run quantize导出低比特权重QAT 模式走quantizer.convert→ 用FullModelTorchTuneCheckpointerInt8DynActInt4WeightQuantizer评测eleuther_eval与生成generate。版本前提8da4w需要 PyTorch 2.38da4w-qat需要 PyTorch 2.4quantization.py 中还要求 torchao 0.7.0否则导入即抛错。参数取舍groupsize越小精度越细但内存开销越大仓库各配置默认使用 256QAT 不支持torch.compileQAT 评测/生成只接受 torchtune 格式 checkpoint。PTQ 场景训练后量化与基准测试建议直接使用 torchaotorchtune 的量化配方聚焦于 QAT 与 QLoRA 微调链路。【免费下载链接】torchtunePyTorch native post-training library项目地址: https://gitcode.com/GitHub_Trending/to/torchtune创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考