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

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

本文将逐步介绍如何使用 NVIDIA Model Optimizer，通过训练后量化（PTQ）方法将 CLIP 模型量化为 FP8 格式。关于模型量化的一般介绍，请参见《Model Quantization: Concepts, Methods, and Why It Matters》。

什么是 NVIDIA Model Optimizer？

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

ModelOpt 支持高性能量化格式，如 FP4、FP8、INT8 和 INT4，并支持包括 SmoothQuant、AWQ、SVDQuant 和 Double Quantization 在内的高级算法。它同时支持 PTQ 和量化感知训练（QAT）。

什么是 CLIP？

CLIP（Contrastive Language-Image Pretraining）由 OpenAI 于 2021 年提出，是一种基础视觉语言模型（VLM），通过在大规模图文对上进行对比学习，为图像和文本学习共享嵌入空间。它能够生成语义对齐的表示，因此已成为现代多模态系统中的核心构建模块。

CLIP 文本编码器被广泛复用为文本到图像（例如 Stable Diffusion）和文本到视频（例如 AnimateDiff）生成的条件模块。其视觉编码器则作为多模态 LLM（如 LLaVA）以及开放词汇感知模型（如 OWL-ViT）中的视觉骨干。OpenCLIP 和 SigLIP 等后继模型扩展了数据规模并改进了目标函数，但保留了双编码器对比学习范式。

量化方案

本文将以下量化方案作为分步指南，用于通过 ModelOpt 运行 CLIP 模型量化并理解其流程。

首先，按如下方式准备相应的模型和数据集：

• 基础 CLIP 模型：CLIP-ViT-L-14-laion2B-s32B-b82K
• 量化校准数据集：来自 MS-COCO 的 10K 子集
• 模型准确性评估任务聚焦于 CLIP_benchmark 中的三项：cifar100（零样本分类）、imagenet1k（零样本分类）、mscoco_captions（零样本检索）

如何使用 ModelOpt 运行 PTQ

以下代码示例展示了如何使用 ModelOpt 以 FP8 格式对 CLIP 模型运行 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) per-tensor static quantization
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)

# Calibration: 8k MS-COCO image-text pairs
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)

# Save quantized modelopt checkpoint
q_model.save_pretrained(ckpt_path)
mtq.print_quant_summary(q_model)

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

有关详细配置架构，请参见 ModelOpt 量化指南。量化配置中的超参数始终可以按需微调，而找到最优值通常需要若干轮迭代。

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

要将注意力纳入量化范围，请为 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 基准评估

保存的 ModelOpt 检查点可恢复到任何下游评估脚本中。详情请参考《Restoring ModelOpt Models》。量化后的 CLIP 检查点在三个基准上进行了评估：零样本分类（CIFAR-100、ImageNet-1k）和零样本检索（MS-COCO Captions）。FP16 CLIP 模型作为基线。

根据评估结果，CLIP-FP8 量化模型展现出与 CLIP-FP16 模型相当的质量。值得注意的是，当 patch embedding 层中的量化器被禁用时，量化对模型质量的影响变得可以忽略。

ModelOpt PTQ 流程内部机制

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

伪量化过程主要通过两种方式工作：

• 统计信息收集：量化器在数据通过时收集张量统计信息（例如最小值和最大值）。这些统计信息用于计算缩放因子等最优量化参数。
• 量化模拟：量化器对流经网络的张量执行量化再反量化（QDQ）操作。它仅模拟低精度计算；真正的加速和内存节省应通过将模型导出到 NVIDIA TensorRT 等部署框架来实现。

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

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

总体而言，ModelOpt PTQ 流程包括六个阶段：

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

QAT 通过在冻结量化器状态的情况下微调模型权重，恢复量化导致的质量损失。它比 PTQ 计算量更大，但可以更好地提升量化模型质量。更多详情请参见 ModelOpt 示例。

开始使用 NVIDIA Model Optimizer

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

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

要探索其他工作流并根据自身用例调整 ModelOpt，请参见 ModelOpt 文档。