NVIDIA TensorRT 是一个 AI 推理库，用于优化机器学习模型，以便在 NVIDIA GPU 上部署。TensorRT 针对现代架构中的专用硬件（例如 NVIDIA Blackwell Tensor Core）进行优化，可加速高级机器学习模型中的常见操作。它还可以通过使用层融合和自动内核策略选择等优化技术，修改 AI 模型，使其在特定硬件上更高效地运行。
PyTorch 等热门框架为 AI 模型提供了直观且一致的界面；但是，它们并不总是能达到最佳性能。Torch-TensorRT 弥合了这一差距。它是一款功能强大的 PyTorch 模型编译器，可在 NVIDIA GPU 上提供 TensorRT 级别的性能，同时保持 PyTorch 的类可用性。它使您能够在不更改 PyTorch API 的情况下将性能提高一倍。
在这篇博客文章中，我们将展示 Torch-TensorRT 如何使优化变得非常简单，只需对代码进行少量更改即可实现显著加速。在 FLUX.1-dev（一个 120 亿参数的整流流 Transformer）上，只需一行代码，性能即可提升至原生 PyTorch FP16 的 1.5 倍。此外，通过在此基础上应用简单的 FP8 量化过程，性能可提高至原来的 2.4 倍。
我们还将展示如何使用 Torch-TensorRT 来支持高级扩散器工作流，例如通过即时模型重构进行低秩自适应 (LoRA) 。
模型加速
HuggingFace Diffusers 是一个 SDK，可让开发者轻松访问各种高级模型。它还支持许多高级用例，例如用于定制现成模型的微调和 LoRA。
通常，模型优化是易用性和性能之间的权衡。许多优化工作流需要将模型从 PyTorch 导出为第三方格式。这一劳动密集型流程使得围绕 AI 模型的复杂工作流难以实现，例如支持多个 GPU 模型或在运行时修改权重（例如使用 LoRA 时）。
Torch-TensorRT 优化了扩散器管线的关键组件，无需中间步骤，代码量非常少。如果要更改管线的一部分、添加控制网络或加载 LoRA，则无需额外工作。新权重可实时重新安装，而现有工作流可能需要您在工作流之外手动重新导出和重新优化模型。
我们使用 FLUX.1-dev 作为示例，展示了 Torch-TensorRT 加速的有效性以及易于访问的集成工作流。FLUX.1-dev 可以从 HuggingFace 下载，运行方式如下：
代码：Hugging Face 加载流程图
import torch
import torch_tensorrt
from diffusers import FLUXPipeline
DEVICE = "cuda:0"
pipe = FLUXPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev",
torch_dtype=torch.float16,
)
pipe.to(DEVICE).to(torch.float16)
NVIDIA HGX B200 GPU 可以开箱即用，为许多低延迟应用提供足够的性能；但是，通过一个简单的额外步骤，延迟可以显著降低，从而进一步增强用户体验。
使用可变的 Torch-TensorRT 模块进行单行优化
Torch-TensorRT 通过优化模型并为将部署应用程序的 GPU 生成特定的 TensorRT 引擎，从而提高了标准 PyTorch 的模型性能。TensorRT 通过层融合和内核自动调优等技术实现这一目标，从而最大限度地提高吞吐量并最大限度地降低延迟。
对于静态计算图，TensorRT 可提供强大的加速功能，并支持各种部署场景。但是，将其与涉及动态权重、图形或 diffusers 等第三方 API 的应用程序集成可能需要额外的开发工作。Torch-TensorRT 通过可变 Torch-TensorRT 模块 (MTTM) 简化了这些动态用例。
MTTM 旨在作为 PyTorch 模块的透明包装器，具有使用 TensorRT 即时优化 forward 函数的附加功能。它保留了源 PyTorch 模型中的所有功能，可将 Torch-TensorRT 加速无缝集成到 Hugging Face 管线等复杂系统中。因此，在运行时修改管线的流程（例如插入 LoRA 适配器）无需任何额外的代码更改即可正常运行。
当检测到图形或权重发生变化时，该模块会通过重新安装或重新编译前向函数来自动调整这些变化。此外，与现有的 JIT 工作流不同，MTTM 是可串行的，因此是 AOT 和 JIT 的混合方法。开发者可以发布预编译的 MTTM，如果运行时条件发生变化，可以随时进行改装或重新编译。
pipe.transformer = torch_tensorrt.MutableTorchTensorRTModule(
pipe.transformer,
strict=False,
allow_complex_guards_as_runtime_asserts=True,
enabled_precisions={torch.float16},
truncate_double=True,
immutable_weights=False,
offload_module_to_cpu=True,
)
images = pipe(
prompt,
output_type="pil",
num_inference_steps=20,
num_images_per_prompt=batch_size,
).images # Compilation starts at the first time the module is called
首次执行时，MTTM 会自动捕获输入模式并复制原始模块的所有行为。因此，用户无需在 PyTorch 和 TensorRT 模块之间提供虚假输入或管理其他必要属性（例如设备或配置）。编译完成后，可以执行以下操作轻松地将优化后的模块序列化并保存到磁盘：
torch_trt.MutableTorchTensorRTModule.save(trt_gm, "mutable_module.pt2") # Save the module to the disk
透明地支持 LoRA
图像生成应用中最流行的高级工作流之一是使用 LoRA 来定制 FLUX 等模型的输出。当用户想要对 FLUX.1-dev 模型权重进行有针对性的调整时（例如，生成不同艺术风格的图像），他们通常会加载不同的 LoRA 模块。然而，对于模型优化器来说，切换 LoRAs 可能会很困难，因为权重更新通常需要重新编译，而重新编译是一个耗时的过程，无法在同一个运行时会话中执行。
Torch-TensorRT 通过权重重装克服了这一限制，可以在不重新编译的情况下在同一运行时内切换 LoRA。通过使用权重重装，在权重发生变化时，周转时间会显著缩短，从而提高 GenAI 应用的实时性。
使用 MTTM，改装过程在后台无缝处理。用户只需应用 Hugging Face 的 load_lora_weights API 将所需的 LoRA 加载到工作流中，MTTM 就会自动检测权重变化。在后续执行中，MTTM 会在内部执行必要的改装，无需用户采取任何额外操作。
# Standard HuggingFace LoRA Loading Procedure
pipe.load_lora_weights(path, adapter_name="lora1")
pipe.set_adapters(["lora1"], adapter_weights=[1])
pipe.fuse_lora()
pipe.unload_lora_weights()
images = pipe(
prompt,
output_type="pil",
num_inference_steps=20,
num_images_per_prompt=batch_size,
).images # Refitting happens here
量化（Quantization）
为了进一步优化 FLUX.1-dev 模型以在较小的 GPU 上运行，我们应用量化技术来提高推理性能，同时通过将权重和激活从标准 16 位浮点表示转换为 8 位浮点 (fp8) 等低精度格式来减少模型大小和 GPU 显存占用。
模型量化使用 NVIDIA TensorRT 模型优化器 (nvidia-modelopt) 执行，这是一个提供高级模型优化技术的综合库，包括量化、修剪、蒸馏和推测性解码。该工具可有效压缩深度学习模型，以便在 TensorRT-LLM 或 TensorRT 等框架上进行下游部署，从而最大限度地提高推理速度。
将模型量化为目标精度后，我们可以使用相同的 Torch-TensorRT 编译过程，并将启用精度设置为目标量化精度：
quantized_transformer = mtq.diffusers.quantize_diffusers_module(pipe, "transformer", mtq.FP8_DEFAULT_CONFIG),
pipe.transformer = torch_tensorrt.MutableTorchTensorRTModule(
quantized_transformer,
strict=False,
allow_complex_guards_as_runtime_asserts=True,
enabled_precisions={torch.float16},
truncate_double=True,
immutable_weights=False,
offload_module_to_cpu=True,
)
images = pipe(
prompt,
output_type="pil",
num_inference_steps=20,
num_images_per_prompt=batch_size,
).images
FLUX 性能基准测试
在 B200 GPU 上使用 PyTorch 运行 FLUX.1-dev，20 个降噪步骤，批量大小为 2，平均运行时间约为 6.56 秒，每个步骤的延迟为 164 毫秒。
MTTM 提供加速功能，只需一行代码，生成两张图像的平均时间就减少到 4.28 秒，FP16 精度下的每一步延迟降低到 107 毫秒。与在 FP16 下使用 PyTorch 运行 FLUX.1-dev 相比，速度可提升 1.5 倍。
通过进一步量化为 FP8，优化的 TensorRT 引擎与在 FP16 下运行的原始 Hugging Face 实现相比，速度可提升至原来的 2.4 倍。值得注意的是，在 B200 GPU 上，利用 FP8 精度可以获得性能提升。具体而言，生成一批 2 张图像的平均时间进一步缩短至 2.72 秒，每一步 68 毫秒，速度提升了 240%。
此外，FP8 使得 FLUX.1-dev 能够在 GeForce RTX 5090 等消费级硬件上运行，而 GeForce RTX 5090 是一款具有 32 GB 显存的消费级 GPU。在批量大小为 1 时，RTX 5090 可以成功编译并以每步平均 260 毫秒的延迟运行。
结论
通过使用 MutableTorchTensorRT 模块等 Torch-TensorRT API，开发者可以创建具有高吞吐量、低延迟性能的生成式 AI 应用。通过使用 JIT 和 AOT 编译的混合体以及简单的量化技术，用户可以显著减少推理时间和 GPU 显存占用，同时通过对现有 PyTorch 代码进行最小更改，支持 LoRA 等动态工作流。
未来，FP4 精度也将被添加到支持的精度列表中，以进一步优化模型的内存占用和推理速度。虽然这篇博客文章介绍了 FLUX，但这种工作流适用于 HuggingFace Diffusers 支持的许多扩散模型，例如 Stable Diffusion 和 Kandinsky。详细了解 FLUX.1-dev 并查看 GitHub 上的 FLUX.1-dev 演示。
