Awesome

NVIDIA TensorRT Model Optimizer

A Library to Quantize and Compress Deep Learning Models for Optimized Inference on GPUs

Examples | Documentation | Benchmark Results | Roadmap | ModelOpt-Windows

Latest News

• [2024/10/23] Quantized FP8 Llama-3.1 Instruct models available on Hugging Face for download: 8B, 70B, 405B
• [2024/9/10] Post-Training Quantization of LLMs with NVIDIA NeMo and TensorRT Model Optimizer
• [2024/8/28] Boosting Llama 3.1 405B Performance up to 44% with TensorRT Model Optimizer on NVIDIA H200 GPUs
• [2024/8/28] Up to 1.9X Higher Llama 3.1 Performance with Medusa
• [2024/08/15] New features in recent releases: Cache Diffusion, QLoRA workflow with NVIDIA NeMo, and more. Check out our blog for details.
• [2024/06/03] Model Optimizer now has an experimental feature to deploy to vLLM as part of our effort to support popular deployment frameworks. Check out the workflow here
• [2024/05/08] Announcement: Model Optimizer Now Formally Available to Further Accelerate GenAI Inference Performance
• [2024/03/27] Model Optimizer supercharges TensorRT-LLM to set MLPerf LLM inference records
• [2024/03/18] GTC Session: Optimize Generative AI Inference with Quantization in TensorRT-LLM and TensorRT
• [2024/03/07] Model Optimizer's 8-bit Post-Training Quantization enables TensorRT to accelerate Stable Diffusion to nearly 2x faster
• [2024/02/01] Speed up inference with Model Optimizer quantization techniques in TRT-LLM

Table of Contents

• Model Optimizer Overview
• Installation
• Techniques
  • Quantization
  • Distillation
  • Pruning
  • Sparsity
• Examples
• Support Matrix
• Benchmark
• Quantized Checkpoints
• Roadmap
• Release Notes
• Contributing

Model Optimizer Overview

Minimizing inference costs presents a significant challenge as generative AI models continue to grow in complexity and size. The NVIDIA TensorRT Model Optimizer (referred to as Model Optimizer, or ModelOpt) is a library comprising state-of-the-art model optimization techniques including quantization, distillation, pruning, and sparsity to compress models. It accepts a torch or ONNX model as inputs and provides Python APIs for users to easily stack different model optimization techniques to produce an optimized quantized checkpoint. Seamlessly integrated within the NVIDIA AI software ecosystem, the quantized checkpoint generated from Model Optimizer is ready for deployment in downstream inference frameworks like TensorRT-LLM or TensorRT. ModelOpt is integrated with NVIDIA NeMo and Megatron-LM for training-in-the-loop optimization techniques. For enterprise users, the 8-bit quantization with Stable Diffusion is also available on NVIDIA NIM.

Model Optimizer for both Linux and Windows are available for free for all developers on NVIDIA PyPI. This repository is for sharing examples and GPU-optimized recipes as well as collecting feedback from the community.

Installation / Docker

Easiest way to get started with using Model Optimizer and additional dependencies (e.g. TensorRT-LLM deployment) is to start from our docker image.

After installing the NVIDIA Container Toolkit, please run the following commands to build the Model Optimizer docker container which has all the necessary dependencies pre-installed for running the examples.

# Clone the ModelOpt repository
git clone https://github.com/NVIDIA/TensorRT-Model-Optimizer.git
cd TensorRT-Model-Optimizer

# Build the docker (will be tagged `docker.io/library/modelopt_examples:latest`)
# You may customize `docker/Dockerfile` to include or exclude certain dependencies you may or may not need.
./docker/build.sh

# Run the docker image
docker run --gpus all -it --shm-size 20g --rm docker.io/library/modelopt_examples:latest bash

# Check installation (inside the docker container)
python -c "import modelopt; print(modelopt.__version__)"

See the installation guide for more details on alternate pre-built docker images or installation in a local environment.

NOTE: Unless specified otherwise, all example READMEs assume they are using the above ModelOpt docker image for running the examples. The example specific dependencies are required to be install separately from their respective requirements.txt files if not using the ModelOpt's docker image.

Techniques

Quantization

Quantization is an effective model optimization technique for large models. Quantization with Model Optimizer can compress model size by 2x-4x, speeding up inference while preserving model quality. Model Optimizer enables highly performant quantization formats including FP8, INT8, INT4, etc and supports advanced algorithms such as SmoothQuant, AWQ, and Double Quantization with easy-to-use Python APIs. Both Post-training quantization (PTQ) and Quantization-aware training (QAT) are supported.

Distillation

Knowledge Distillation allows for increasing the accuracy and/or convergence speed of a desired model architecture by using a more powerful model's learned features to guide a student model's objective function into imitating it.

Pruning

Pruning is a technique to reduce the model size and accelerate the inference by removing unnecessary weights. Model Optimizer provides Python APIs to prune Linear and Conv layers, and Transformer attention heads, MLP, embedding hidden size and number of layers (depth).

Sparsity

Sparsity is a technique to further reduce the memory footprint of deep learning models and accelerate the inference. Model Optimizer Python APIs to apply weight sparsity to a given model. It also supports NVIDIA 2:4 sparsity pattern and various sparsification methods, such as NVIDIA ASP and SparseGPT.

Examples

• PTQ for LLMs covers how to use Post-training quantization (PTQ) and export to TensorRT-LLM for deployment for popular pre-trained models from frameworks like
  • Hugging Face
  • NVIDIA NeMo
  • NVIDIA Megatron-LM
  • Medusa
• PTQ for Diffusers walks through how to quantize a diffusion model with FP8 or INT8, export to ONNX, and deploy with TensorRT. The Diffusers example in this repo is complementary to the demoDiffusion example in TensorRT repo and includes FP8 plugins as well as the latest updates on INT8 quantization.
• QAT for LLMs demonstrates the recipe and workflow for Quantization-aware Training (QAT), which can further preserve model accuracy at low precisions (e.g., INT4, or 4-bit in NVIDIA Blackwell platform).
• Sparsity for LLMs shows how to perform Post-training Sparsification and Sparsity-aware fine-tuning on a pre-trained Hugging Face model.
• Pruning demonstrates how to optimally prune Linear and Conv layers, and Transformer attention heads, MLP, and depth using the Model Optimizer for following frameworks:
  • NVIDIA NeMo / NVIDIA Megatron-LM GPT-style models (e.g. Llama 3, Mistral NeMo, etc.)
  • Hugging Face language models like BERT and GPT-J
  • Computer Vision models like NVIDIA Tao framework detection models.
• ONNX PTQ shows how to quantize the ONNX models in INT4 or INT8 quantization mode. The examples also include the deployment of quantized ONNX models using TensorRT.
• Distillation for LLMs demonstrates how to use Knowledge Distillation, which can increasing the accuracy and/or convergence speed for finetuning / QAT.
• Chained Optimizations shows how to chain multiple optimizations together (e.g. Pruning + Distillation + Quantization).
• Model Hub provides an example to deploy and run quantized Llama 3.1 8B instruct model from Nvidia's Hugging Face model hub on both TensorRT-LLM and vLLM.

Model Support Matrix

• For LLM quantization, please refer to this support matrix.
• For VLM quantization, please refer to this support matrix.
• For Diffusion, Model Optimizer supports FLUX, Stable Diffusion 3, Stable Diffusion XL, SDXL-Turbo, and Stable Diffusion 2.1.
• For speculative decoding, please refer to this support matrix.

Benchmark

Please find the benchmarks at here.

Quantized Checkpoints

Quantized checkpoints in Hugging Face model hub are ready for TensorRT-LLM and vLLM deployments. More models coming soon.

Roadmap

Please see our product roadmap.

Release Notes

Please see Model Optimizer Changelog here.

Contributing

At the moment, we are not accepting external contributions. However, this will soon change after we open source our library in early 2025 with a focus on extensibility. We welcome any feedback and feature requests. Please open an issue if you have any suggestions or questions.