借助 NVIDIA TensorRT 轻量级引擎，为 AI 应用提供卓越性能和最小占用空间

NVIDIA TensorRT一个成熟的数据中心推理库，已迅速成为 NVIDIA GeForce RTX 和 NVIDIA RTX GPU 的理想推理后端。现在，部署 TensorRT 借助预构建的 TensorRT 引擎，应用变得更加简单。

新发布的重量去除 TensorRT 10.0 引擎提供了一种独特的解决方案，通过将其缩小到仅执行代码来最大限度地减少引擎出货量，从而实现 > 95%的引擎尺寸压缩。

在本文中，我们将讨论如何构建轻量级引擎，以及如何使用 TensorRT 40MB 精简运行时直接在最终用户设备上对其进行改装。

什么是轻量级引擎？

TensorRT 10.0 中引入的权重剥离引擎包含无权重的执行代码 (CUDA 内核).在构建阶段启用权重剥离后，引擎的体积比传统引擎小 95%以上，并且仅保留必要的权重以优化性能。

这些引擎支持 ONNX 模型和其他网络定义，提供与 refittable 引擎类似的优势。在反序列化期间，可以直接在最终用户设备上使用模型中的原始权重对其进行 refit，从而将延迟降至最低，同时保持不影响 inference 性能。

为什么要进行重量去除？

传统上，TensorRT 引擎包含网络的所有权重，从而在特定于硬件的引擎之间产生冗余权重。除非直接在最终用户设备上构建引擎，否则这就需要为整个安装基础提供重磅的预构建引擎。

假设一个场景是，工作站 AI 应用通过 M 个独立的 GPU SKU，创建 P 个引擎用于各个优化配置文件的每个 N 个深度学习模型。这会导致最终应用二进制文件中的大量权重重复（M × P），可以通过使用权重 stripped 的引擎来避免。

对于卷积神经网络（CNN）和大型语言模型（LLM）用例，Weight-stripped 引擎可实现超过 95% 的引擎大小压缩，从而在不增加应用程序大小的情况下打包更多的人工智能功能。这些引擎在 TensorRT 小更新之间版本兼容，并且在使用 kVERSION_COMPATIBLE 标志构建时可以使用约 40 MB 的 lean runtime。

NVIDIA TensorRT Cloud，目前精选合作伙伴抢先体验还提供在各种 NVIDIA GPU 上构建轻量级引擎的选项。即将支持构建和改装轻量级 NVIDIA TensorRT-LLM 引擎。

注册以随时了解所有 TensorRT 相关内容 。

构建轻量级引擎

在本地构建权重受限的引擎时，TensorRT 构建器仍然需要模型权重来做出优化决策，从而确保与普通或有权重的引擎相比，稍后改装时的性能保持一致。

在构建过程中使用真实权重使 TensorRT 能够通过不断折叠静态节点和引入融合优化来优化计算，例如，在引擎构建过程中 GELU 文字未标记为可改装时，将其直接嵌入到单个 CUDA 内核中。

TensorRT Cloud 还有助于从 ONNX 模型中创建权重受限的引擎。

有关更多信息，请参阅开始使用稍后的部分。

部署权重受限的引擎

如果您的应用附带权重剥离的引擎，您可以在最终用户设备上使用 ONNX 文件中的权重，在数秒内轻松地对其进行重装。在序列化后，重装的引擎可以消除重复性的重装成本，同时保持 TensorRT 所知的快速反序列化效率。

您可以通过提供精简的运行时（~40 MB in TensorRT 10.0）来使用 refit 功能，而无需整个 TensorRT 构建器。这样改装的权重剥离引擎保留了版本前向兼容性（VFC）和硬件前向兼容性（HFC）的优势，使其能够在下一代 GPUs 上运行，而无需应用程序更新。

案例研究

我们使用 SDXL 实现了 99%以上的压缩率，demoDiffusion在 NVIDIA GeForce RTX 4090 GPU 上运行 (表 1)。

虽然我们即将支持权重受限的 TensorRT-LLM 引擎，但以下是基于 NVIDIA GeForce RTX 4090 GPU 内部构建的一些测量结果，在多个 LLM 上实现了 99%以上的压缩 (表 2)。

即使在NVIDIA H100等数据中心GPU上，您也可以在 TensorRT LLM Llama 模型上看到类似的压缩（表3）。

开始使用

在 TensorRT 10.0 中，一个新的标志 kREFIT_IDENTICAL 优化了 builder，以便在引擎使用 identical 权重进行 refitted 时发挥作用。当与 kSTRIP_PLAN 标志一起使用，因为在 build-time 和 refit-time 期间使用所有 identical 权重，可以确保实现最大限度的性能，从而最大限度地减小引擎尺寸。

新的序列化标志SerializationFlag::kEXCLUDE_WEIGHTS此外，在 refit 过程中进行设置时，它还会将 refitted weight-stripped 引擎永久保存到磁盘中，从而避免在未来应用启动期间产生 recurring refit 成本。如果不设置，它支持使用 refit API 使用新的权重进行 continuous refitting。

以下是完整重量减轻和改装工作流程的说明：

• 构建轻量级引擎。
• 从 INetworkDefinition 中改装权重受限的引擎。
• 通过 ONNX 解析器改装权重受限引擎。
• 序列化为全权重引擎。

构建轻量化引擎

设置相应的构建器标志以启用权重限制构建。

config->setFlag(BuilderFlag::kSTRIP_PLAN);
config->setFlag(BuilderFlag::kREFIT_IDENTICAL);
builder->buildSerializedNetwork(*network, *config);

config.flags |= 1 << int(trt.BuilderFlag.STRIP_PLAN)
config.flags |= 1 << int(trt.BuilderFlag.REFIT_IDENTICAL)
builder.build_serialized_network(network, config)

构建引擎后，保存引擎计划文件，并通过应用安装程序分发。

改装轻量级引擎

根据指定网络的方式，TensorRTINetworkDefinitionAPI 或 ONNX 运行时 (ORT)，请参阅相关的改装说明。

改装 INetworkDefinition 中的权重受限引擎

从引擎创建 refitter。

auto refitter = std::unique_ptr<nvinfer1::IRefitter>(
            	nvinfer1::createInferRefitter(*engine, logger));

refitter = trt.Refitter(engine, TRT_LOGGER)

在客户端上，首次启动网络时，更新引擎中的所有权重。在此处，使用getAllWeights因为引擎计划中的所有权重都已删除。

int32_t const nbWts = refitter->getAllWeights(0, nullptr);
std::vector<char const*> allWtsNames(nbWts);
refitter->getAllWeights(nbWts, allWtsNames.data());

all_weights_names = refitter.get_all_weights()

逐一更新权重。

for (int32_t i = 0; i < nbWts; ++i)
    refitter->setNamedWeights(allWtsNames[i], Weights{...});

for name in all_weights_names:
    refitter.set_named_weights(name, trt.Weights(...))

通过 ONNX 解析器改装权重受限引擎

从引擎创建 refitter。

auto refitter = std::unique_ptr<nvinfer1::IRefitter>(
     nvinfer1::createInferRefitter(*engine, logger));

refitter = trt.Refitter(engine, TRT_LOGGER)

创建 Refit 解析器。

auto parser_refitter = std::unique_ptr<nvonnxparser::IParserRefitter>(
            	nvonnxparser::createParserRefitter(*refitter, logger));

parser_refitter = trt.OnnxParserRefitter(refitter, TRT_LOGGER)

根据原始 ONNX 文件改装引擎。

if (parser_refitter->refitFromFile(onnx_model_path.c_str())) {
    assert(refitter->refitCudaEngine());
}

parser_refitter.refit_from_file(onnx_model_path)
refitter.refit_cuda_engine()

序列化为全权重引擎

保存完整的引擎平面图文件。

auto serializationConfig = std::unique_ptr<ISerializationConfig>(engine->createSerializationConfig());
auto serializationFlag = serializationConfig->getFlags();
serializationFlag &= ~(1<< static_cast<uint32_t>(SerializationFlag::kEXCLUDE_WEIGHTS));
serializationConfig->setFlags(serializationFlag);
auto hostMemory = std::unique_ptr<nvinfer1::IHostMemory>(engine->serializeWithConfig(*serializationConfig));

serialization_config = engine.create_serialization_config()
serialization_config.flags &= ~(1 << int(trt.SerializationFlag.EXCLUDE_WEIGHTS))
binary = engine.serialize_with_config(serialization_config)

应用程序现在可以使用新的完整引擎计划文件进行未来推理。

限制

在 TensorRT 10.0 中，weight-stripped 功能仅限于使用相同的 build-time 权重进行 refitting，以确保正确的功能和最大性能。refit-time 和 build-time 权重之间的差异可能会导致未定义的行为。builder 控制要剥离的权重，用户无法做出 layer-level 决策。在将来的版本中，可能会取消此限制，以实现更加精细的控制。

在即将发布的版本中，我们将支持 TensorRT-LLM 中的权重限制引擎，并且同样需要使用相同的构建时权重来确保功能正确。

与 ONNX 运行时集成

TensorRT 10.0 权重剥离功能已经集成到 开源神经网络推理优化工具包（ORT）中，并将从 ORT 1.18.1 版本开始提供。ORT 在 Windows 生态系统中的多个硬件提供商中提供了一个标准推理平台。这种集成使 TensorRT 能够通过 ORT API 提供相同的功能，在满足不同客户硬件需求时减少了 shipment sizes。

在 ORT 中，weight-stripped 功能使用相同的基于 EP 上下文节点的逻辑，嵌入 serialized TensorRT 引擎在 ONNX 模型中绕过 builder 实例化。这种方法避免了 builder 资源的分配，并显著缩短了 TensorRT EP 会话设置时间。有关更多信息，请参阅ONNX Runtime: Accelerated AI Deployment for PC Apps(GTC 会议)。

现在，您可以提供权重剥离引擎，而不是完整引擎。为此，请使用 Boolean 创建具有权重剥离的 TensorRT 引擎 (稍后将进行调整)trt_weightless_engine_enableTensorRT EP 会话选项以及现有的trt_dump_ep_context_model标志。

生成的 ONNX EP 上下文元数据模型包含重构权重所需的原始 ONNX 模型的文件名。这意味着不再需要 builder 库进行运输。只需将原始 ONNX 模型的权重和嵌入式权重剥离引擎与主 nvinfer 库一起运输即可。

在针对该特定模型进行下一次 TensorRT EP 调用时，仍启用 trt_weightless_engine_enable 和 trt_dump_ep_context_model 后，引擎将使用原始权重进行改装。

结束语

借助 TensorRT 轻量级引擎，您可以在无需担心增加应用大小的情况下将大量 AI 功能融入应用，同时利用 NVIDIA GPU 上的 TensorRT 峰值性能。

设备上改装可实现持续更新，同时改进权重，而无需重建引擎。TensorRT-LLM 即将支持权重剥离和改装，为生成、部署和维护未来的生成式 AI 模型提供了一种新颖的方法。

有关 ONNX 模型的端到端权重剥离示例的更多信息，请参阅/NVIDIA/TensorRTGitHub 库或/NVIDIA/TensorRT notebook。