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 能夠通過不斷折疊靜態節點和引入融合優化來優化計算，例如，在引擎構建過程中 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)。
| SDXL fp16 | Full engine size (MB) | Weight-stripped engine size (MB) |
| clip | 237.51 | 4.37 |
| clip2 | 1329.40 | 8.28 |
| unet | 6493.25 | 58.19 |
雖然我們即將支持權重受限的 TensorRT-LLM 引擎，但以下是基于 NVIDIA GeForce RTX 4090 GPU 內部構建的一些測量結果，在多個 LLM 上實現了 99%以上的壓縮 (表 2)。
| Model | Full engine size (MB) | Weight-stripped engine size (MB) |
| chatglm-6b int4 | 3831.09 | 4.50 |
| chatglm-6b int8 | 6470.10 | 34.12 |
| gemma2.5b int4 | 2987.81 | 6.62 |
| gemma2.5b int8 | 3905.64 | 26.03 |
| gemma7b int4 | 6826.14 | 7.86 |
| gemma7b int8 | 10409.58 | 54.07 |
| llama2-7b fp16 | 12856.44 | 4.04 |
| llama2-7b int4 | 3691.12 | 5.41 |
| llama2-7b int8 | 6683.72 | 41.54 |
| mistral fp16 | 13831.44 | 4.06 |
| mistral int4 | 3954.45 | 8.12 |
| mistral int8 | 7178.73 | 41.91 |
| phi2 fp16 | 5304.53 | 4.04 |
即使在NVIDIA H100等數據中心GPU上，您也可以在 TensorRT LLM Llama 模型上看到類似的壓縮（表3）。
| Model | Full engine size (MB) | Weight-stripped engine size (MB) |
| llama-7b fp16 + WoQ int8 | 6704.55 | 28.69 |
| llama2-70b fp8 + TP=2 | 66341.72 | 60.61 |
開始使用
在 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。
?
