使用 NVIDIA TensorRT-LLM 优化 OpenAI GPT-OSS 模型

本 Notebook 提供了使用 NVIDIA TensorRT-LLM 优化 gpt-oss 模型以实现高性能推理的分步指南。TensorRT-LLM 为用户提供了一个易于使用的 Python API 来定义大型语言模型 (LLM)，并支持最先进的优化，以便在 NVIDIA GPU 上高效地执行推理。TensorRT-LLM 还包含用于创建 Python 和 C++ 运行时以高性能方式协调推理执行的组件。

TensorRT-LLM 支持以下模型：

  • gpt-oss-20b
  • gpt-oss-120b

在本指南中，我们将运行 gpt-oss-20b，如果您想尝试更大的模型或需要更多自定义设置，请参阅 此 部署指南。

注意：您的输入提示应使用 harmony 响应 格式，模型才能正常工作，尽管本指南不强制要求。

在 NVIDIA Brev 上启动

您可以使用 NVIDIA Brev 来简化环境设置。单击下面的按钮，在具有预配置必要依赖项的 Brev 实例上启动此项目。

部署后，单击“打开 Notebook”按钮开始本指南

在 Brev 上启动

先决条件

硬件

要运行 gpt-oss-20b 模型，您需要一个至少具有 20 GB VRAM 的 NVIDIA GPU。

推荐 GPU：NVIDIA Hopper（例如，H100、H200）、NVIDIA Blackwell（例如，B100、B200）、NVIDIA RTX PRO、NVIDIA RTX 50 系列（例如，RTX 5090）。

软件

  • CUDA Toolkit 12.8 或更高版本
  • Python 3.12 或更高版本

安装 TensorRT-LLM

有多种安装 TensorRT-LLM 的方法。在本指南中，我们将介绍使用 NVIDIA NGC 提供的预构建 Docker 容器以及从源代码构建。

如果您正在使用 NVIDIA Brev，则可以跳过此部分。

使用 NVIDIA NGC

从 NVIDIA NGC 拉取用于 GPT-OSS 的预构建 TensorRT-LLM 容器。 这是最简单的入门方法，并确保包含所有依赖项。

docker pull nvcr.io/nvidia/tensorrt-llm/release:gpt-oss-dev
docker run --gpus all -it --rm -v $(pwd):/workspace nvcr.io/nvidia/tensorrt-llm/release:gpt-oss-dev

使用 Docker（从源代码构建）

或者，您可以从源代码构建 TensorRT-LLM 容器。 如果您想修改源代码或使用自定义分支，此方法很有用。 有关详细说明，请参阅 官方文档。

TensorRT-LLM 将很快通过 pip 提供

关于 GPU 架构的注意事项：首次运行模型时，TensorRT-LLM 将为您的特定 GPU 架构（例如 Hopper、Ada 或 Blackwell）构建优化的引擎。如果您看到有关 GPU 的 CUDA 功能（例如 sm_90、sm_120）与 PyTorch 安装不兼容的警告，请确保您拥有最新的 NVIDIA 驱动程序和与 PyTorch 版本匹配的 CUDA Toolkit 版本。

验证 TensorRT-LLM 安装

from tensorrt_llm import LLM, SamplingParams

利用 TensorRT-LLM Python API

在下一个代码单元中，我们将演示如何使用 TensorRT-LLM Python API 来：

  1. 从 Hugging Face 下载指定的模型权重（使用您的 HF_TOKEN 进行身份验证）。
  2. 如果尚不存在，则自动为您的 GPU 架构构建 TensorRT 引擎。
  3. 加载模型并为其准备推理。
  4. 运行一个简单的文本生成示例以验证一切是否正常。

注意：首次运行可能需要几分钟时间，因为它会下载模型并构建引擎。 后续运行将快得多，因为引擎将被缓存。

llm = LLM(model="openai/gpt-oss-20b")

prompts = ["Hello, my name is", "The capital of France is"]
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
for output in llm.generate(prompts, sampling_params):
    print(f"Prompt: {output.prompt!r}, Generated text: {output.outputs[0].text!r}")

结论和后续步骤

恭喜！您已成功使用 TensorRT-LLM Python API 优化并运行了大型语言模型。

在本 Notebook 中，您已学会：

  • 使用必要的依赖项设置您的环境。
  • 使用 tensorrt_llm.LLM API 从 Hugging Face Hub 下载模型。
  • 自动构建针对 GPU 优化的、高性能的 TensorRT 引擎。
  • 使用优化后的模型运行推理。

您可以探索更高级的功能以进一步提高性能和效率：

  • 基准测试：尝试运行 基准测试，将 TensorRT-LLM 引擎的延迟和吞吐量与原始 Hugging Face 模型进行比较。您可以通过迭代更多提示并测量执行时间来完成此操作。

  • 量化：TensorRT-LLM 支持 各种量化技术（如 INT8 或 FP8），以减小模型大小并加速推理，同时对准确性影响最小。这是在资源受限的硬件上部署模型的强大功能。

  • 使用 NVIDIA Dynamo 部署：对于生产环境，您可以使用 NVIDIA Dynamo 部署您的 TensorRT-LLM 引擎，以实现健壮、可扩展且支持多模型服务。