How nvidia/tensorrt-model-optimizer works
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: quantization Quantization · torch/export Unified HF export · onnx/export ONNX exporters · distill Distillation
Components 20
PyTorch Optimization
- Quantization
modelopt/torch/quantization - Pruning
modelopt/torch/prune - Architecture search
modelopt/torch/nas - Sparsity
modelopt/torch/sparsity - Speculative decoding
modelopt/torch/fastgen - Puzzletron
modelopt/torch/puzzletron - Optimization recipes
modelopt/recipe
ONNX Optimization
- ONNX quantization
modelopt/onnx/quantization - Precision conversion
modelopt/onnx/autocast - Graph surgery
modelopt/onnx/graph_surgery
Export and Deployment
- Unified HF export
modelopt/torch/export - ONNX exporters
modelopt/onnx/export - Optimized checkpoint
- Inference frameworks
- Model deployment runtime
modelopt/torch/_deploy
Training and Adaptation
- Distillation
modelopt/torch/distill - DMS distillation training
experimental/dms/dms/training/engine.py - Training integrations
modelopt/torch/quantization/plugins
Other components
- Model practitioner
- HF / PyTorch / ONNX models
Connections 25
- 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