NVIDIA 宣布推出新版 v0.15 NVIDIA TensorRT 模型优化器，这是一种先进的模型优化技术量化工具包，包含量化、稀疏和剪枝。这些技术可降低模型复杂性，使 NVIDIA TensorRT-LLM 和 NVIDIA TensorRT 等下游推理框架能够更有效地优化生成式 AI 模型的推理速度。
本文将概述近期 TensorRT Model Optimizer 版本的一些主要特性和升级，包括缓存扩散、使用 NVIDIA NeMo 的全新量化感知训练工作流程以及 QLoRA 支持。
缓存扩散
以前，TensorRT Model Optimizer（简称‘Model Optimizer’）利用其 8 位训练后量化（PTQ）技术强效助力 NVIDIA TensorRT，为 Stable Diffusion XL 的性能设立了标杆。为进一步普及扩散模型的快速推理，Model Optimizer v0.15 增加了对缓存扩散的支持，可与 FP8 或 INT8 PTQ 一起使用，以在推理时进一步加速扩散模型。
DeepCache 和块缓存等缓存扩散方法通过重复使用先前降噪步骤中的缓存输出，无需额外训练即可优化推理速度。缓存机制利用了扩散模型反向降噪过程的内在特性，其中连续步骤之间的高级特征具有显著的时间一致性，并且可以缓存和重复使用。缓存扩散与 DiT 和 UNet 等各种主干模型兼容，在不影响质量或训练成本的情况下实现显著的推理加速。
要启用缓存扩散，开发者只需在 Model Optimizer 中使用单个’cachify’实例和扩散管道即可。有关详细示例，请参阅缓存扩散教程notebook。对于基于 NVIDIA H100 Tensor Core GPU 的 FP16 Stable Diffusion XL（SDXL），在 Model Optimizer 中启用缓存扩散可将每秒图像的速度提高 1.67 倍（图 1）。此外，当还启用 FP8 时，此加速会提高。此外，Model Optimizer 使用户能够自定义缓存配置，以加快推理速度。在不久的将来，使用 TensorRT 的缓存扩散管道将支持更多的扩散模型。
为提供更公平的比较，我们使用禁用缓存的Model Optimizer缓存扩散管线对无缓存的FP16基准进行了基准测试，而不是使用TensorRT中的具有批量大小限制的demoDiffusion管线。NVIDIA H100 80 GB HBM3 GPU；步长大小30；批量大小16；TensorRT v10.2.0；TensorRT Model Optimizer v0.15
使用 NVIDIA NeMo 进行量化感知训练
量化感知训练 (QAT) 是一种在模拟量化效果的同时训练神经网络的技术，旨在在量化后恢复模型精度。该过程涉及在训练期间计算扩展系数，并将模拟量化损失纳入微调过程，从而提高神经网络对量化的弹性。在 Model Optimizer 中，QAT 使用自定义 CUDA 内核进行模拟量化，实现较低精度的模型权重和激活，以实现高效的硬件部署。
使用 Model Optimizer mtq.quantize() API 量化的模型可以直接使用原始训练管线进行微调。在 QAT 期间，量化器中的缩放系数会被冻结，模型权重则会进行微调。由于 QAT 过程需要更短的微调时间，因此建议使用较小的学习率。
Model Optimizer v0.15 将 QAT 集成支持从 Hugging Face Trainer 和 Megatron 扩展到 NVIDIA NeMo，这是一个用于开发自定义生成式 AI 模型的企业级平台。现在，Model Optimizer 对 NeMo 模型提供一流的支持。要了解如何使用现有的 NeMo 训练管线执行 QAT，请参阅 NeMo GitHub 库中的新 QAT 示例。了解更多关于 QAT 的信息。
QLoRA 工作流程
量化低阶适应（QLoRA）是一种高效的微调技术，可以在模型训练期间减少内存占用和计算复杂性。通过将量化与低阶适应（LoRA）相结合，QLoRA 使硬件资源有限的开发者更容易使用 LLM 微调。
Model Optimizer 已经使用 NVIDIA NeMo 和 NF4 数据类型添加了对 QLoRA 工作流程的支持。有关工作流程的详细信息，请参阅 NeMo 文档。对于 Alpaka 数据集上的 Llama 13B 模型，QLoRA 可以根据批量大小将峰值显存占用率降低 29%-51%，同时保持相同的模型精度（图 2）。请注意，与 LoRA 相比，QLoRA 具有更长的训练步长时间的权衡（表 1）。
NVIDIA H100 GPU；序列长度为 512；全局批量大小为 256；NeMo 24.07；TensorRT 模型优化器 v0.13
批量大小 |
每个全局批量的时间 | ||
| LoRA | QLoRA | 增加百分比 | |
| 2 | 2.7 | 6.7 | 148% |
| 4 | 2.3 | 4.4 | 91% |
| 8 | 2.2 | 3.2 | 46% |
NVIDIA H100 GPU；序列长度为 512；全局批量大小为 256；NeMo 24.07；TensorRT 模型优化器 v0.13
对 AI 模型的扩展支持
TensorRT 模型优化器已经扩展对更广泛的热门 AI 模型的支持，包括 Stability.ai 的 Stable Diffusion 3、Google 的 RecurrentGemma、Microsoft 的 Phi-3、Snowflake 的 Arctic 2 和 Databricks 的 DBRX。有关更多详细信息，请参阅示例脚本以获取教程和支持矩阵。
开始使用
NVIDIA TensorRT 模型优化器提供与 NVIDIA TensorRT-LLM 和 TensorRT 的无缝集成以进行部署。它可作为 ‘nvidia-modelopt‘ 在 PyPI 上安装。访问 GitHub 上的 NVIDIA/TensorRT-Model-Optimizer，了解推理优化的示例脚本和方法。有关更多详细信息，请参阅 Model Optimizer 文档。
我们重视您对 TensorRT Model Optimizer 的反馈。如果您有任何建议、问题或功能请求，请在 GitHub 上打开新的 NVIDIA/TensorRT-Model-Optimizer 问题。您的输入将有助于我们迭代量化工具包，以更好地满足您的需求。
