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 問題。您的輸入將有助于我們迭代量化工具包,以更好地滿足您的需求。
?