Linux 安装指南

系统要求

最新模型优化器 (nvidia-modelopt) 目前有以下系统要求：

操作系统	Linux
架构	x86_64, aarch64 (SBSA)
Python	>=3.8,<3.13
CUDA	>=11.8 (推荐 12.x)
PyTorch（可选）	>=2.0
ONNX 运行时（可选）	1.18
TensorRT-LLM（可选）	0.15

环境设置

Docker image (Recommended)

使用ModelOpt的docker镜像

开始使用模型优化器和额外依赖项（例如 TensorRT-LLM 部署）的最简单方法是从我们的 Docker 镜像开始。

安装NVIDIA Container Toolkit后，请运行以下命令来构建Model Optimizer docker容器，该容器已预装所有必要的依赖项以运行示例。

# 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.
bash 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__)"

使用替代的NVIDIA Docker镜像

对于PyTorch，你也可以使用NVIDIA NGC PyTorch容器，而对于NVIDIA NeMo框架，你可以使用NeMo容器。这两个容器都预装了Model Optimizer。NeMo容器还包含了HuggingFace和TensorRT-LLM的依赖。如果尚未更新，请确保将Model Optimizer更新到最新版本。

Local environment (PIP / Conda)

设置虚拟环境

我们建议如果您还没有虚拟环境，请设置一个。运行以下命令来设置并激活一个名为modelopt的conda虚拟环境，使用Python 3.12：

conda create -n modelopt python=3.12 pip
conda activate modelopt

（可选）安装所需的PyTorch版本

默认情况下，将安装 pip 上可用的最新 PyTorch 版本 (torch>=2.0)。如果您想为特定的 CUDA 版本安装特定的 PyTorch 版本，请首先按照说明安装您所需的 PyTorch 版本。

（可选）安装ONNX Runtime

如果您希望使用ModelOpt的ONNX功能，请按照ONNX Runtime文档中的步骤为您的CUDA版本正确安装推荐的ONNX Runtime版本。如果ONNX Runtime未预安装，ModelOpt将使用默认的ONNX Runtime，这可能与您的CUDA版本不兼容。

（可选）安装其他NVIDIA依赖项

如果您希望将ModelOpt与其他NVIDIA库（例如TensorRT-LLM、NeMo、Triton等）结合使用，请确保检查这些库在本地环境中的安装便捷性。如果您遇到任何问题，我们建议使用docker镜像以获得无缝体验。例如，TensorRT-LLM文档要求在docker镜像中安装。您仍然可以选择在本地使用ModelOpt的其他功能，例如量化HuggingFace模型，然后使用docker镜像进行部署。

安装模型优化器

ModelOpt 及其依赖项可以通过 pip 安装。在使用前，请查看 ModelOpt 及其任何依赖项的许可条款。

如果您构建并使用ModelOpt的docker镜像，您可以跳过此步骤，因为该镜像已经包含了ModelOpt和所有预安装的可选依赖项。如果您使用其他推荐的docker镜像，ModelOpt已经预安装了一些以下的可选依赖项。请确保使用pip升级到最新版本的ModelOpt（包含您需要的适当可选依赖项），如下所示。

pip install "nvidia-modelopt[all]" -U --extra-index-url https://pypi.nvidia.com

如果您只想安装部分依赖项，请将[all]替换为下面描述的所需可选依赖项。

识别正确的部分依赖

请注意，当安装nvidia-modelopt时，如果没有任何可选依赖项，只会安装最基本的需求，并且没有任何模块可以在没有适当的可选依赖项或[all]可选依赖项的情况下工作。以下是需要安装的可选依赖项列表，以正确使用相应的模块：

模块	可选依赖项
`modelopt.deploy`	`[deploy]`
`modelopt.onnx`	`[onnx]`
`modelopt.torch`	`[torch]`
`modelopt.torch._deploy`	`[torch, deploy]`

此外，我们支持为以下第三方包安装依赖项：

第三方包	可选依赖项
Huggingface (`transformers`, `diffusers`, 等)	`[hf]`

检查安装

提示

当你第一次使用ModelOpt的PyTorch量化API时，如果安装了torch和CUDA，它将编译快速量化内核。这可能需要几分钟，但后续的量化调用将会快得多。要调用编译并检查是否成功或为docker构建预编译，请运行以下命令：

python -c "import modelopt.torch.quantization.extensions as ext; ext.precompile()"