NVIDIA Model Optimizer (ModelOpt)

repository·main·Indexed 25 days ago

https://github.com/nvidia/model-optimizer

A unified library of SOTA model optimization techniques including quantization, pruning, Neural Architecture Search (NAS), distillation, and speculative decoding. It supports Hugging Face, PyTorch, and ONNX models, compressing deep learning models for deployment in frameworks such as TensorRT-LLM, TensorRT, and vLLM to optimize inference speed.

Tokens
178.1K
Snippets
453
Records
777
Agent score
86%

What's inside nvidia-model-optimizer

  1. Overview of ModelOpt Quantization

    main

    ModelOpt provides a quantization toolkit designed for NVIDIA's hardware and software stack. It supports quantization for both PyTorch and ONNX frameworks.

    ModelOpt uses simulated quantization in the original precision to help developers test and optimize the trade-off between model accuracy and low-precision formats. To realize actual performance speedups and memory savings, models quantized via simulation must be exported to deployment frameworks such as TensorRT or TensorRT-LLM.

  2. Overview of LLM Pruning and Distillation with ModelOpt

    main

    NVIDIA ModelOpt (Model Optimizer) provides two primary methods for structural compression of Large Language Models (LLMs) to reduce size and improve efficiency:

    1. Minitron (Homogeneous Pruning): Best for quick, reliable, and moderate compression (10–20% parameter reduction). It applies uniform pruning decisions across all layers, resulting in a smaller model with the same architecture as the original. It targets a specific parameter count.
    2. Puzzletron (Heterogeneous NAS-based Pruning): Best for aggressive compression (>30% memory reduction) under strict hardware memory budgets. It uses Neural Architecture Search (NAS) and Mixed-Integer Programming (MIP) to find optimal, variable layer widths for each layer independently. This results in a heterogeneous architecture tailored to specific constraints.

    Both methods are typically followed by Knowledge Distillation to recover accuracy lost during the pruning process.

    Note: This guide focuses on dense Transformer models. For MoE (Mixture of Experts) or hybrid architectures, refer to the Nemotron-3-Nano-30B-A3B-BF16 tutorial.

  3. Overview of the EAGLE3 Pipeline Stages

    main

    The EAGLE3 pipeline consists of 4 stages used to add support for new models. These stages are mapped to tasks in the pipeline YAML configuration:

    TaskStageContainerScriptDescription
    task_0Data synthesisvllm/vllm-openaicommon/vllm/query.shServes model with vLLM and generates synthetic conversations.
    task_1Hidden state dumpvllm/vllm-openaicommon/eagle3/dump_offline_data*.shDumps hidden states from generated conversations.
    task_2Training + Exporttensorrt-llm/releasecommon/eagle3/train_eagle.shTrains the EAGLE3 draft model and exports the HF checkpoint.
    task_3Benchmarkvllm/vllm-openaicommon/specdec_bench/quick_check.shRuns the speculative decoding benchmark.
  4. Overview of NVIDIA Model Optimizer

    main

    NVIDIA Model Optimizer (ModelOpt) is a library designed to accelerate AI models using state-of-the-art optimization techniques.

    Supported Inputs:

    • Hugging Face models
    • PyTorch models
    • ONNX models

    Optimization Techniques:

    • Quantization
    • Pruning
    • Neural Architecture Search (NAS)
    • Distillation
    • Speculative decoding
    • Sparsity

    Deployment Integration: Optimized checkpoints are designed for seamless deployment in inference frameworks including:

    • SGLang
    • TensorRT-LLM
    • TensorRT
    • vLLM

    The unified Hugging Face export API supports both transformers and diffusers models.

  5. Overview of NVIDIA Model Optimizer (ModelOpt)

    main

    NVIDIA Model Optimizer (ModelOpt) is a library designed to minimize inference costs for generative AI models through compression techniques like quantization and sparsity.

    Key Capabilities:

    • Input Support: Accepts torch or ONNX models.
    • Output/Deployment: Generates optimized checkpoints ready for TensorRT-LLM, TensorRT (Linux), Microsoft DirectML, and TensorRT-RTX (Windows).
    • Ecosystem Integration:
      • Supports Hugging Face export for transformers and diffusers models.
      • Integrates with NVIDIA NeMo and Megatron-LM for training-in-the-loop optimization.
      • Supports Microsoft Olive and ONNX Runtime via the DirectML path on Windows.
    • Platform Support:
      • Linux: Standard ModelOpt library.
      • Windows: ModelOpt-Windows optimized for RTX PC systems, featuring local GPU calibration and reduced memory consumption.

    Installation: Available via NVIDIA PyPI: pip install nvidia-modelopt (or search for nvidia-modelopt on PyPI).

  6. Future directions for the Model Optimizer recipe system

    main

    The recipe system is planned to expand with the following capabilities:

    • QAT recipes: Support for recipe_type: qat including training hyperparameters, distillation settings, and dataset configuration.
    • Sparsity recipes: Configurations for structured and unstructured pruning.
    • Speculative decoding recipes: Settings for draft models and vocabulary calibration.
    • Composite recipes: Ability to chain multiple optimization stages (e.g., quantize then prune) in a single recipe.
    • Dataset configuration: A standardized dataset section for specifying calibration data.
    • Recipe merging and override utilities: Programmatic tools for composing and customizing recipes.
    • Unified entry point: A nv-modelopt CLI that uses --recipe as the primary configuration mechanism to replace per-example scripts.
  7. Explore the modelopt API submodules

    main

    The modelopt API is organized into several specialized submodules. Depending on your deployment target or framework, you should use the following modules:

    • modelopt.torch: For PyTorch-based model optimization tasks.
    • modelopt.onnx: For optimization tasks involving the ONNX format.
    • modelopt.deploy: For deployment-related optimization and model preparation.

    Refer to the specific submodule documentation for detailed function signatures and usage patterns.

  8. Understand quantization format components

    main

    A quantization format in NVIDIA Model Optimizer is defined by three primary components:

    1. Precision format: Defines the bit-width of the quantized values. This includes integer formats (sign bit + mantissa bits) and floating-point formats (sign bit + exponent bits + mantissa bits).
    2. Block format: Defines how a tensor is divided into blocks to share scaling factors. Common formats include:
      • Per-tensor quantization: The entire tensor uses a single global scaling factor.
      • Per-channel quantization: Each channel is quantized separately.
      • Fine-grained per-block quantization: The tensor is divided into fixed-size blocks along the channel dimension (typically used for low-bit quantization like 4-bit to preserve accuracy).
    3. Calibration algorithm: Calculates scaling factors and may adjust weights to maximize accuracy. Examples include max calibration, Entropy Calibration, SmoothQuant, AWQ, and SVDQuant.
  9. Pruning Overview and Techniques

    main

    Model pruning reduces neural network complexity and improves efficiency by removing redundant parameters or connections. Supported techniques in NVIDIA Model Optimizer include:

    • Minitron: Designed for GPT, Mamba, MoE, and Hybrid Transformer Mamba models in NVIDIA Megatron-LM or Megatron-Bridge frameworks. It prunes embedding hidden size, MLP FFN hidden size, attention heads, Mamba heads/dimensions, MoE experts/sizes, and model depth.
    • Puzzletron: Uses a Mixed Integer Programming (MIP) based NAS search algorithm.
    • FastNAS: Recommended for Computer Vision models; finds a subnet that maximizes a score function under given constraints.
  10. Understand PTQ quantization recipes and schemes

    main

    Post-Training Quantization (PTQ) in NVIDIA Model Optimizer is organized into two main categories:

    1. General Recipes: Located under general/ptq/. These are model-agnostic and serve as the recommended starting point for any model.
    2. Model-Specific Recipes: Located under huggingface/. These are tailored to specific model_type folders or specific checkpoints (e.g., huggingface/models/<org>/<checkpoint>/). These recipes deviate from the general baseline to optimize for specific architectures.

    When starting quantization, it is recommended to begin with a general recipe and move to model-specific recipes if specialized optimization is required.

  11. Sparsity Concepts: Structured, Unstructured, and N:M Sparsity

    main

    Unstructured Sparsity

    Zero weights are randomly distributed across the weight matrix. It is flexible but often results in poor utilization on parallel hardware like GPUs.

    Structured Sparsity

    Enforces a specific sparsity pattern on weights. This is more efficient for memory access and can achieve higher mathematical throughput.

    N:M Sparsity

    A fine-grained structured pattern where in each block of $M$ contiguous elements, at most $N$ are non-zeros.

    • Benefits: Reduced memory bandwidth requirements and higher math throughput on Sparse Tensor Cores.
    • NVIDIA Support: On Ampere or later architectures, 2:4 sparsity (2 non-zeros in every block of 4) is supported for accelerated inference.