NVIDIA推出了最新版的NVIDIA TensorRT Model Optimizer v0.15,这是一款尖端的量化工具包,旨在增强如量化、稀疏和修剪等模型优化技术。据NVIDIA技术博客,此次更新旨在减少模型复杂度并优化生成式AI模型的推理速度。
缓存扩散
新版本包括对缓存扩散的支持,基于之前建立的8-bit后训练量化(PTQ)技术。此功能通过重用之前去噪步骤的缓存输出来加速推理时间的扩散模型。DeepCache和块缓存等方法无需额外训练即可优化推理速度。此机制利用连续去噪步骤之间的高级特征的时间一致性,使其兼容如DiT和UNet等模型。
开发者可以通过在Model Optimizer中使用单个‘cachify’实例来启用缓存扩散功能。例如,在NVIDIA H100 Tensor Core GPU上启用缓存扩散的Stable Diffusion XL(SDXL)模型,可以在每秒生成图像上获得1.67倍的加速。如果同时启用了FP8,这种加速效果会进一步增加。
支持NVIDIA NeMo的量化感知训练
量化感知训练(QAT)在神经网络训练期间模拟量化效果,以恢复量化后的模型精度。此过程计算缩放因子并在微调过程中纳入模拟量化损失。Model Optimizer使用自定义CUDA内核进行模拟量化,实现了低精度模型权重和激活以便于高效的硬件部署。
Model Optimizer v0.15扩展了对NVIDIA NeMo的QAT集成支持,NeMo是一个开发定制生成型AI模型的企业级平台。对NeMo模型的一流支持允许用户直接使用原始训练管道微调模型。更多细节,请参见NeMo GitHub仓库中的QAT示例。
QLoRA 工作流程
量化低秩适应(QLoRA)是一种在模型训练期间减少内存使用和计算复杂度的微调技术。它结合了量化和低秩适应(LoRA),使大型语言模型(LLM)的微调更加可行。Model Optimizer现在使用NF4数据类型支持使用NVIDIA NeMo的QLoRA工作流程。在Alpaca数据集上的Llama 13B模型中,QLoRA可以减少峰值内存使用量29-51%,同时保持模型精度。
扩展AI模型支持
最新版本还扩展了对更多AI模型的支持,包括Stability.ai的Stable Diffusion 3、Google的RecurrentGemma、Microsoft的Phi-3、Snowflake的Arctic 2和Databricks的DBRX。有关更多详细信息,请参见Model Optimizer GitHub仓库中的示例脚本和支持矩阵。
开始使用
NVIDIA TensorRT Model Optimizer提供与NVIDIA TensorRT-LLM和TensorRT的无缝集成,可在PyPI上以nvidia-modelopt进行安装。访问NVIDIA TensorRT Model Optimizer GitHub页面获取推理优化的示例脚本和方案。也提供了全面的文档。
Image source: Shutterstock