模型优化与训练后量化

模型量化是一种有效的方法，可减少显存使用并提升消费级设备上的推理性能。通过降低计算和内存需求同时保持模型质量，量化有助于AI模型在资源受限的环境中更高效地运行。

本文介绍如何使用 NVIDIA Model Optimizer 通过训练后量化（PTQ）方法将 CLIP 模型量化为 FP8 格式，并包含导出 PyTorch 检查点的示例工作流。

阅读完整博客文章：

模型量化：使用 NVIDIA Model Optimizer 进行训练后量化

来源：https://developer.nvidia.com/blog/model-quantization-post-training-quantization-using-nvidia-model-optimizer/ 模型量化是一种有效的方法，可减少显存使用并提升消费级设备（如 NVIDIA GeForce RTX（https://www.nvidia.com/en-us/geforce/graphics-cards/）GPU）上的推理性能。通过降低计算和内存需求同时保持模型质量，量化有助于 AI 模型在资源受限的环境中更高效地运行。

本文介绍如何使用 NVIDIA Model Optimizer（https://github.com/NVIDIA/Model-Optimizer）通过训练后量化（PTQ）（https://developer.nvidia.com/blog/optimizing-llms-for-performance-and-accuracy-with-post-training-quantization/）方法将 CLIP 模型量化为 FP8 格式。有关模型量化的通用介绍，请参阅模型量化：概念、方法及其重要性（https://developer.nvidia.com/blog/model-quantization-concepts-methods-and-why-it-matters/）。

什么是 NVIDIA Model Optimizer？https://developer.nvidia.com/blog/model-quantization-post-training-quantization-using-nvidia-model-optimizer/#what_is_nvidia_model_optimizer

NVIDIA Model Optimizer（https://github.com/NVIDIA/Model-Optimizer）（ModelOpt）库集成了最先进的模型优化技术，以压缩和加速 AI 模型。这些技术包括量化、蒸馏、剪枝、推测性解码和稀疏性。ModelOpt 接受 Hugging Face、PyTorch 或 ONNX 格式的模型作为输入，并提供 Python API，使用户能够轻松组合不同的优化技术，生成优化的检查点。

ModelOpt 支持高性能量化格式，如 FP4、FP8、INT8 和 INT4，以及高级算法，包括 SmoothQuant、AWQ、SVDQuant 和 Double Quantization。它同时支持 PTQ（https://developer.nvidia.com/blog/optimizing-llms-for-performance-and-accuracy-with-post-training-quantization/）和量化感知训练（QAT）（https://developer.nvidia.com/blog/how-quantization-aware-training-enables-low-precision-accuracy-recovery/）。

什么是 CLIP？https://developer.nvidia.com/blog/model-quantization-post-training-quantization-using-nvidia-model-optimizer/#what_is_clip

CLIP（https://github.com/openai/CLIP）（对比语言-图像预训练）由 OpenAI 于 2021 年提出，是一个基础视觉语言模型（VLM）（https://www.nvidia.com/en-us/glossary/vision-language-models/），通过在大规模图像-文本对上进行对比学习，学习图像和文本的共享嵌入空间。其生成语义对齐表示的能力使其成为现代多模态系统的核心构建块。

CLIP 文本编码器被广泛复用为文本到图像（例如 Stable Diffusion）和文本到视频（例如 AnimateDiff）生成的条件模块。其视觉编码器在多模态大语言模型（如 LLaVA）和开放词汇感知模型（如 OWL-ViT）中充当视觉骨干。后续模型如 OpenCLIP 和 SigLIP 扩展了数据规模并优化了目标函数，但保留了双编码器对比范式。

量化配方https://developer.nvidia.com/blog/model-quantization-post-training-quantization-using-nvidia-model-optimizer/#quantization_recipe

以下量化配方作为分步指南，用于在 ModelOpt 中运行 CLIP 模型量化，以了解该过程的工作原理。

首先，准备相应的模型和数据集，如下所示：

• 基础 CLIP 模型：CLIP-ViT-L-14-laion2B-s32B-b82K（https://huggingface.co/laion/CLIP-ViT-L-14-laion2B-s32B-b82K）
• 用于量化的校准数据集：来自 MS-COCO（https://github.com/revantteotia/clip-training?tab=readme-ov-file#preparing-training-dataset）的 10K 子集
• 模型精度评估任务集中在 CLIP_benchmark（https://github.com/LAION-AI/CLIP_benchmark）中的三个任务上：
  • cifar100（零样本分类）
  • imagenet1k（零样本分类）
  • mscoco_captions（零样本检索）

如何使用 ModelOpt 运行 PTQhttps://developer.nvidia.com/blog/model-quantization-post-training-quantization-using-nvidia-model-optimizer/#how_to_run_ptq_with_modelopt

以下代码示例展示了如何使用 ModelOpt 对 CLIP 模型进行 FP8 的 PTQ 量化：

import torch
from torch.utils.data import DataLoader, Subset
from transformers import CLIPModel, CLIPTokenizer, CLIPImageProcessor
from transformers.models.clip.modeling_clip import CLIPAttention

import modelopt.torch.opt as mto
import modelopt.torch.quantization as mtq
from modelopt.torch.quantization.plugins.diffusion.diffusers import _QuantAttention

# FP8 (E4M3) 逐张量静态量化
FP8_CFG = {
    "quant_cfg": {
        "*weight_quantizer":      {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "Half"},
        "*input_quantizer":       {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "Half"},
        "*[qkv]_bmm_quantizer":   {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "Half"},
        "*bmm2_output_quantizer": {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "Half"},
        "default": {"enable": False},
    },
    "algorithm": "max",
}

mto.enable_huggingface_checkpointing()
mtq.QuantModuleRegistry.register({CLIPAttention: "CLIPAttention"})(_QuantAttention)

model = CLIPModel.from_pretrained(args.model_ckpt, attn_implementation="sdpa").half().eval().cuda()

tokenizer = CLIPTokenizer.from_pretrained(args.model_ckpt)
processor = CLIPImageProcessor.from_pretrained(args.model_ckpt)
calib_set = Subset(CLIP_COCO_dataset(ANN, IMG_DIR, tokenizer, processor), range(8192))
loader    = DataLoader(calib_set, batch_size=512, num_workers=4)

# 校准：8k MS-COCO 图像-文本对
def calibrate(m):
    for img, txt in loader:
        m.get_text_features(input_ids=txt.cuda())
        m.get_image_features(pixel_values=img.cuda())

q_model = mtq.quantize(model, FP8_CFG, forward_loop=calibrate)

# 保存量化的 ModelOpt 检查点
q_model.save_pretrained(ckpt_path)
mtq.print_quant_summary(q_model)

FP8_CFG 只是其中一种配方：W8A8（权重和激活均为 FP8）、逐张量、静态量化，使用简单的 AbsMax 算法进行校准。ModelOpt 支持更多维度的选择（逐通道/逐块粒度、动态激活量化、高级校准算法如 AWQ/GPTQ 等）。

有关详细的配置模式，请参阅 ModelOpt 量化指南（https://nvidia.github.io/Model-Optimizer/guides/_quant_cfg.html）。量化配置中的超参数始终可以根据需要进行微调，找到最优值通常需要一些迭代。

在 mtq.quantize 返回后，CLIP 的 Linear 层都带有权重和激活量化器——但注意力块仍未受影响。这是因为多头注意力分发到 torch.nn.functional.scaled_dot_product_attention，这是一个 ModelOpt 模块遍历器无法自行拦截的函数式 API。

要将注意力纳入量化范围，为 CLIPAttention 注册一个量化替换版本：

mtq.QuantModuleRegistry.register({CLIPAttention: 
"CLIPAttention"})(_QuantAttention)

现在，每个 CLIPAttention 实例都被升级为来自 ModelOpt diffusers 插件的 _QuantAttention。在其前向传播中，_QuantAttention 透明地拦截 SDPA 调用，并在融合内核周围插入四个量化器：

• q_bmm_quantizer、k_bmm_quantizer、v_bmm_quantizer 包装投影后的 Q/K/V 张量，然后它们进入内核
• bmm2_output_quantizer 包装内核输出（softmax @ V），然后它流入 out_proj

这确保了整个注意力机制的适当量化。

为了恢复一些精度，通常建议使用 mtq.disable_quantizer 禁用部分量化器。该函数接受一个函数作为输入，该函数本身接受模块名称作为输入。通过正则表达式或字符串匹配，可以选择要禁用的层。在以下示例中，CLIP 模型的 patch_embedding 层中的量化器被禁用。

import re
def filter_func(name):
    pattern = re.compile( 
        r".*(patch_embedding).*"
    )
    return pattern.match(name) is not None

mtq.disable_quantizer(q_model, filter)

CLIP 基准评估https://developer.nvidia.com/blog/model-quantization-post-training-quantization-using-nvidia-model-optimizer/#clip_benchmark_evaluation

保存的 ModelOpt 检查点可以恢复到任何下游评估脚本中。有关详细信息，请参阅恢复 ModelOpt 模型（https://nvidia.github.io/Model-Optimizer/guides/2_save_load.html#restoring-modelopt-models）。量化后的 CLIP 检查点在三项基准上进行了评估：零样本分类（CIFAR-100、ImageNet-1k）和零样本检索（MS-COCO Captions）。FP16 CLIP 模型作为基线。

柱状图比较了量化模型 CLIP-FP8-PTQ 与基线 CLIP-FP16 模型在 CIFAR100、ImageNet-1k 和 MS-COCO 指标上的质量。FP8 PTQ 在这些评估基准上与 FP16 非常接近。 图 1. CLIP 模型质量比较：FP16 基线 vs FP8-PTQ 量化模型 根据评估结果，CLIP-FP8 量化模型表现出与 CLIP-FP16 模型相当的质量。值得注意的是，当 patch embedding 层中的量化器被禁用时，量化对模型质量的影响变得微不足道。

ModelOpt PTQ 流程内部机制https://developer.nvidia.com/blog/model-quantization-post-training-quantization-using-nvidia-model-optimizer/#inside_the_modelopt_ptq_flow

重要的是要理解，这个阶段涉及的是“伪量化”操作，因为模型的实际数据类型并未改变。相反，这些插入的量化器作为观察者，模拟量化的效果，同时保持模型为原始浮点格式。

伪量化过程通过两种关键方式工作：

• 统计信息收集：当数据通过量化器时，它们收集张量统计信息（例如最小值和最大值）。这些统计信息用于计算最优的量化参数，例如缩放因子。
• 量化模拟：量化器对网络中流动的张量执行量化-反量化（QDQ）操作。它仅模拟低精度计算，而真正的加速和内存节省应通过将模型导出到部署框架（如 NVIDIA TensorRT（https://developer.nvidia.com/tensorrt））来实现。

这种模拟至关重要，因为它使您能够在实际量化之前评估模型的精度。量化器应用与部署后量化模型在下游推理框架中相同的舍入和精度限制，因此您可以：

• 在部署前衡量精度影响
• 尝试不同的量化配置
• 识别可能需要特殊处理的问题层

通常，ModelOpt PTQ 流程包括六个阶段：

1. 准备：设置量化配置，在模型权重和/或激活周围插入量化器模块。
2. 校准：通过模型前向传播一小批代表性数据，以便每个量化器收集统计信息（例如激活的 amax）并推导出其缩放因子。
3. 伪量化：量化器现在以浮点形式执行 Q→DQ 往返，忠实地模拟目标格式的精度损失，同时模型仍在 FP16/BF16 中运行。
4. 评估：在保留的评估集上测量精度，并与未量化的基线进行比较。
5. 迭代：如果差距不可接受，调整量化配置（粒度、算法、量化层），禁用敏感层的量化，并重新校准。
6. 导出与部署：一旦精度可接受，伪量化权重被压缩为真正的低精度形式，并导出为检查点以供下游引擎使用。在我们的案例中，我们将 PyTorch 检查点导出为 ONNX，并使用 TensorRT 运行推理。速度和内存节省将在那里实现。

ModelOpt PTQ 工作流程图，展示了准备、校准、伪量化、评估步骤，然后进入 OK 决策，如果是，则导出并部署，否则迭代回准备步骤。 图 2. ModelOpt PTQ 工作流程 QAT 通过微调模型权重（冻结量化器状态）来恢复量化引起的质量损失。它比 PTQ 计算密集更高，但可以更好地提升量化模型质量。有关更多详细信息，请参阅 ModelOpt 示例（https://nvidia.github.io/Model-Optimizer/guides/_pytorch_quantization.html#quantization-aware-training-qat）。

开始使用 NVIDIA Model Optimizerhttps://developer.nvidia.com/blog/model-quantization-post-training-quantization-using-nvidia-model-optimizer/#get_started_with_nvidia_model_optimizer

本文介绍了 NVIDIA Model Optimizer，并通过一个实用的代码示例，将 CLIP 模型量化为 FP8，展示了典型的训练后量化工作流程。在三个评估数据集上的结果表明，FP8 量化可以在保持模型质量的同时，实现更高效的部署路径。

准备好开始将 ModelOpt 用于您自己的模型了吗？请遵循此工作流程：准备模型和校准数据，设置量化配置，校准，根据特定任务的质量指标验证量化模型，保存和恢复 ModelOpt 检查点。

要探索其他工作流程并使 ModelOpt 适应您自己的用例，请参阅 ModelOpt 文档（https://nvidia.github.io/Model-Optimizer/）。