NVIDIA炸裂！一行代码AI模型提速2.4倍，LoRA重训无需等待！

NVIDIA TensorRT是一个专为优化机器学习模型在NVIDIA GPU上部署的人工智能推理库。TensorRT针对现代架构中的专用硬件，例如NVIDIA Blackwell Tensor Cores，加速高级机器学习模型中的常见操作。它还可以通过使用层融合和自动内核策略选择等优化技术，修改人工智能模型以在特定硬件上更有效地运行。

PyTorch等流行的框架为人工智能模型提供了一个直观且一致的界面，但它们并不总能实现最佳性能。Torch-TensorRT弥合了这一差距。它是一个强大的PyTorch模型编译器，可在NVIDIA GPU上提供TensorRT级别的性能，同时保持PyTorch类的可用性。它使您能够在原生PyTorch的基础上提高一倍的性能，而无需更改PyTorch API。

本文将展示Torch-TensorRT如何使优化变得非常简单，只需极少的代码更改即可实现显著的加速。仅使用一行代码，FLUX.1-dev（一个120亿参数的修正流转换器）的性能与原生PyTorch FP16相比提高了1.5倍。此外，通过在此基础上应用简单的FP8量化程序，性能提高到2.4倍。

我们还将展示如何使用Torch-TensorRT来支持高级扩散器工作流程，例如通过即时模型重拟合实现低秩自适应（LoRA）。

模型加速

HuggingFace Diffusers是一个SDK，为开发人员提供了对各种高级模型的便捷访问。它还支持许多高级用例，例如微调和LoRA，用于自定义现成的模型。

通常，模型优化是在易用性和性能之间进行权衡。许多优化工作流程都需要将模型从PyTorch导出为第三方格式。这种劳动密集型过程使得围绕人工智能模型实施复杂的工作流程变得困难，例如支持多个GPU模型或在运行时修改权重（例如使用LoRA时）。

Torch-TensorRT优化了扩散器管道的关键组件，而无需中间步骤，只需极少的代码。如果您更改管道的某个部分、添加controlnet或加载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提供了强大的加速功能，并支持各种部署方案。但是，将其与涉及动态权重、图或第三方API（如扩散器）的应用程序集成可能需要额外的开发工作。Torch-TensorRT通过可变Torch-TensorRT模块（MTTM）简化了这些动态用例。

MTTM被设计为PyTorch模块的透明包装器，并增加了使用TensorRT即时优化前向函数的行为。它保持了源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
# 编译从第一次调用模块时开始

首次执行时，MTTM会自动捕获输入模式并复制原始模块的所有行为。因此，用户无需在PyTorch和TensorRT模块之间提供虚假输入或管理其他必要的属性（例如设备或配置）。编译后，可以使用以下命令轻松地将优化的模块序列化并保存到磁盘：

torch_trt.MutableTorchTensorRTModule.save(trt_gm, "mutable_module.pt2")
# 将模块保存到磁盘

透明地支持LoRA

图像生成应用程序最流行的先进工作流程之一是使用LoRA自定义FLUX等模型的输出。当用户想要对FLUX.1-dev模型权重进行有针对性的调整时（例如，以不同的艺术风格生成图像），他们通常会加载不同的LoRA模块。但是，切换LoRA对于模型优化器来说可能具有挑战性，因为权重更新通常需要重新编译，这是一个耗时的过程，无法在同一运行时会话中执行。

Torch-TensorRT通过权重重新拟合克服了此限制，从而可以在同一运行时内切换LoRA，而无需重新编译。通过使用权重重新拟合，权重更改时的周转时间大大缩短，从而提高了GenAI应用程序的活跃度。

使用MTTM，重新拟合过程在后台无缝处理。用户只需应用Hugging Face的load_lora_weights API将所需的LoRA加载到管道中，MTTM会自动检测权重更改。在后续执行中，MTTM会在内部执行必要的重新拟合，而无需额外的用户操作。

# 标准的HuggingFace LoRA加载程序
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
# 重新拟合发生在这里

量化

为了进一步优化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倍的加速。值得注意的是，在利用FP8精度时，在B200 GPU上获得了性能提升。具体而言，生成一批2个图像的平均时间进一步降至2.72秒，每步68毫秒，表示加速了240％。
A bar chart titled "FLUX.1-dev Image Generation Speedup" comparing performance speedups across three GPUs (H100, H200, B200) using different configurations: PyTorch FP16 (baseline), Torch-TensorRT FP16, and Torch-TensorRT FP8.

图2. FLUX.1-dev基准可视化

此外，FP8使得可以在消费级硬件（例如GeForce RTX 5090）上运行FLUX.1-dev，GeForce RTX 5090是一款具有32 GB内存的消费级GPU。在批量大小为1时，RTX 5090可以成功编译并以平均每步260毫秒的延迟运行。

结论

通过使用像MutableTorchTensorRT模块这样的Torch-TensorRT API，开发人员可以创建具有高吞吐量、低延迟性能的生成式人工智能应用程序。通过结合使用JIT和AOT编译以及简单的量化技术，用户可以显著减少推理时间和GPU内存占用，同时支持动态工作流程，例如通过最少更改其现有PyTorch代码来支持LoRA。

新媒网跨境了解到，未来，FP4精度也将添加到受支持的精度列表中，以进一步优化模型的内存占用和推理速度。虽然本文以FLUX为特色，但此工作流程通常适用于HuggingFace Diffusers支持的许多扩散模型，例如Stable Diffusion和Kandinsky。

有关更多信息并查看GitHub上的FLUX.1-dev演示。

新媒网（公号: 新媒网跨境发布），是一个专业的跨境电商、游戏、支付、贸易和广告社区平台，为百万跨境人传递最新的海外淘金精准资讯情报。

本文来源：新媒网 https://nmedialink.com/posts/6758.html