# nvidia/tensorrt-model-optimizer architecture

> Architecture diagram of the GitHub repository nvidia/tensorrt-model-optimizer, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-09-29). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/nvidia/tensorrt-model-optimizer
- Repository: https://github.com/nvidia/tensorrt-model-optimizer

## Overview

ModelOpt is a model-optimization library: users provide Hugging Face, PyTorch, or ONNX models, apply optimization techniques, and export checkpoints for inference deployment. The graph separates PyTorch and ONNX workflows, includes the README’s pruning, NAS, distillation, sparsity, and speculative-decoding capabilities, and shows deployment as a downstream stage. Paths identify relevant source areas from the tree; beyond the sampled files, capability placement follows the README and module structure, so uncertain internal call links are omitted.

Main entry points: [`modelopt/torch/quantization`](https://github.com/nvidia/tensorrt-model-optimizer/tree/main/modelopt/torch/quantization) (Quantization), [`modelopt/torch/export`](https://github.com/nvidia/tensorrt-model-optimizer/tree/main/modelopt/torch/export) (Unified HF export), [`modelopt/onnx/export`](https://github.com/nvidia/tensorrt-model-optimizer/tree/main/modelopt/onnx/export) (ONNX exporters), [`modelopt/torch/distill`](https://github.com/nvidia/tensorrt-model-optimizer/tree/main/modelopt/torch/distill) (Distillation)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_torch_opt["PyTorch Optimization"]
  node_torch_quant["Quantization"]
  node_pruning["Pruning"]
  node_nas["Architecture search"]
  node_sparsity["Sparsity"]
  node_fastgen["Speculative decoding"]
  node_puzzletron["Puzzletron"]
  node_recipes["Optimization recipes"]
end

subgraph group_onnx_opt["ONNX Optimization"]
  node_onnx_quant["ONNX quantization"]
  node_autocast["Precision conversion"]
  node_graph_surgery["Graph surgery"]
end

subgraph group_export["Export and Deployment"]
  node_torch_export["Unified HF export"]
  node_onnx_export["ONNX exporters"]
  node_checkpoint["Optimized checkpoint"]
  node_deploy["Inference frameworks"]
  node_deploy_runtime["Model deployment runtime"]
end

subgraph group_training["Training and Adaptation"]
  node_distill["Distillation"]
  node_dms["DMS distillation training<br/>[engine.py]"]
  node_training_frameworks["Training integrations"]
end

node_user(("Model practitioner"))
node_model_inputs["HF / PyTorch / ONNX models"]

node_user -->|"provides"| node_model_inputs
node_model_inputs -->|"optimize"| node_torch_quant
node_model_inputs -->|"optimize"| node_pruning
node_model_inputs -->|"search"| node_nas
node_model_inputs -->|"optimize"| node_sparsity
node_model_inputs -->|"distill"| node_distill
node_model_inputs -->|"optimize"| node_fastgen
node_model_inputs -->|"optimize"| node_puzzletron
node_recipes -->|"configure"| node_torch_quant
node_model_inputs -->|"quantize"| node_onnx_quant
node_model_inputs -->|"convert precision"| node_autocast
node_model_inputs -->|"transform graph"| node_graph_surgery
node_torch_quant -->|"export model"| node_torch_export
node_pruning -->|"export model"| node_torch_export
node_distill -->|"export model"| node_torch_export
node_onnx_quant -->|"export model"| node_onnx_export
node_autocast -->|"export model"| node_onnx_export
node_graph_surgery -->|"export model"| node_onnx_export
node_torch_export -->|"produces"| node_checkpoint
node_onnx_export -->|"produces"| node_checkpoint
node_checkpoint -->|"deploy to"| node_deploy
node_checkpoint -->|"load model"| node_deploy_runtime
node_training_frameworks -.->|"supports training"| node_distill
node_training_frameworks -.->|"supports training"| node_torch_quant
node_model_inputs -.->|"distill"| node_dms

click node_torch_quant "https://github.com/nvidia/tensorrt-model-optimizer/tree/main/modelopt/torch/quantization"
click node_pruning "https://github.com/nvidia/tensorrt-model-optimizer/tree/main/modelopt/torch/prune"
click node_nas "https://github.com/nvidia/tensorrt-model-optimizer/tree/main/modelopt/torch/nas"
click node_sparsity "https://github.com/nvidia/tensorrt-model-optimizer/tree/main/modelopt/torch/sparsity"
click node_distill "https://github.com/nvidia/tensorrt-model-optimizer/tree/main/modelopt/torch/distill"
click node_fastgen "https://github.com/nvidia/tensorrt-model-optimizer/tree/main/modelopt/torch/fastgen"
click node_puzzletron "https://github.com/nvidia/tensorrt-model-optimizer/tree/main/modelopt/torch/puzzletron"
click node_recipes "https://github.com/nvidia/tensorrt-model-optimizer/tree/main/modelopt/recipe"
click node_onnx_quant "https://github.com/nvidia/tensorrt-model-optimizer/tree/main/modelopt/onnx/quantization"
click node_autocast "https://github.com/nvidia/tensorrt-model-optimizer/tree/main/modelopt/onnx/autocast"
click node_graph_surgery "https://github.com/nvidia/tensorrt-model-optimizer/tree/main/modelopt/onnx/graph_surgery"
click node_torch_export "https://github.com/nvidia/tensorrt-model-optimizer/tree/main/modelopt/torch/export"
click node_onnx_export "https://github.com/nvidia/tensorrt-model-optimizer/tree/main/modelopt/onnx/export"
click node_deploy_runtime "https://github.com/nvidia/tensorrt-model-optimizer/tree/main/modelopt/torch/_deploy"
click node_dms "https://github.com/nvidia/tensorrt-model-optimizer/blob/main/experimental/dms/dms/training/engine.py"
click node_training_frameworks "https://github.com/nvidia/tensorrt-model-optimizer/tree/main/modelopt/torch/quantization/plugins"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_torch_quant,node_pruning,node_nas,node_sparsity,node_fastgen,node_puzzletron,node_recipes toneBlue
class node_onnx_quant,node_autocast,node_graph_surgery toneAmber
class node_torch_export,node_onnx_export,node_checkpoint,node_deploy,node_deploy_runtime toneMint
class node_distill,node_dms,node_training_frameworks toneRose
class node_user,node_model_inputs toneIndigo
```

## Components

### PyTorch Optimization

- [Quantization](https://github.com/nvidia/tensorrt-model-optimizer/tree/main/modelopt/torch/quantization): `modelopt/torch/quantization`
- [Pruning](https://github.com/nvidia/tensorrt-model-optimizer/tree/main/modelopt/torch/prune): `modelopt/torch/prune`
- [Architecture search](https://github.com/nvidia/tensorrt-model-optimizer/tree/main/modelopt/torch/nas): `modelopt/torch/nas`
- [Sparsity](https://github.com/nvidia/tensorrt-model-optimizer/tree/main/modelopt/torch/sparsity): `modelopt/torch/sparsity`
- [Speculative decoding](https://github.com/nvidia/tensorrt-model-optimizer/tree/main/modelopt/torch/fastgen): `modelopt/torch/fastgen`
- [Puzzletron](https://github.com/nvidia/tensorrt-model-optimizer/tree/main/modelopt/torch/puzzletron): `modelopt/torch/puzzletron`
- [Optimization recipes](https://github.com/nvidia/tensorrt-model-optimizer/tree/main/modelopt/recipe): `modelopt/recipe`

### ONNX Optimization

- [ONNX quantization](https://github.com/nvidia/tensorrt-model-optimizer/tree/main/modelopt/onnx/quantization): `modelopt/onnx/quantization`
- [Precision conversion](https://github.com/nvidia/tensorrt-model-optimizer/tree/main/modelopt/onnx/autocast): `modelopt/onnx/autocast`
- [Graph surgery](https://github.com/nvidia/tensorrt-model-optimizer/tree/main/modelopt/onnx/graph_surgery): `modelopt/onnx/graph_surgery`

### Export and Deployment

- [Unified HF export](https://github.com/nvidia/tensorrt-model-optimizer/tree/main/modelopt/torch/export): `modelopt/torch/export`
- [ONNX exporters](https://github.com/nvidia/tensorrt-model-optimizer/tree/main/modelopt/onnx/export): `modelopt/onnx/export`
- **Optimized checkpoint**
- **Inference frameworks**
- [Model deployment runtime](https://github.com/nvidia/tensorrt-model-optimizer/tree/main/modelopt/torch/_deploy): `modelopt/torch/_deploy`

### Training and Adaptation

- [Distillation](https://github.com/nvidia/tensorrt-model-optimizer/tree/main/modelopt/torch/distill): `modelopt/torch/distill`
- [DMS distillation training](https://github.com/nvidia/tensorrt-model-optimizer/blob/main/experimental/dms/dms/training/engine.py): `experimental/dms/dms/training/engine.py`
- [Training integrations](https://github.com/nvidia/tensorrt-model-optimizer/tree/main/modelopt/torch/quantization/plugins): `modelopt/torch/quantization/plugins`

### Other components

- **Model practitioner**
- **HF / PyTorch / ONNX models**

## Connections

- Model practitioner → HF / PyTorch / ONNX models: provides
- HF / PyTorch / ONNX models → Quantization: optimize
- HF / PyTorch / ONNX models → Pruning: optimize
- HF / PyTorch / ONNX models → Architecture search: search
- HF / PyTorch / ONNX models → Sparsity: optimize
- HF / PyTorch / ONNX models → Distillation: distill
- HF / PyTorch / ONNX models → Speculative decoding: optimize
- HF / PyTorch / ONNX models → Puzzletron: optimize
- Optimization recipes → Quantization: configure
- HF / PyTorch / ONNX models → ONNX quantization: quantize
- HF / PyTorch / ONNX models → Precision conversion: convert precision
- HF / PyTorch / ONNX models → Graph surgery: transform graph
- Quantization → Unified HF export: export model
- Pruning → Unified HF export: export model
- Distillation → Unified HF export: export model
- ONNX quantization → ONNX exporters: export model
- Precision conversion → ONNX exporters: export model
- Graph surgery → ONNX exporters: export model
- Unified HF export → Optimized checkpoint: produces
- ONNX exporters → Optimized checkpoint: produces
- Optimized checkpoint → Inference frameworks: deploy to
- Optimized checkpoint → Model deployment runtime: load model
- Training integrations → Distillation: supports training
- Training integrations → Quantization: supports training
- HF / PyTorch / ONNX models → DMS distillation training: distill