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) |
構建引擎后,保存引擎計劃文件,并通過應用安裝程序分發。
改裝輕量級引擎
根據指定網絡的方式,TensorRTINetworkDefinition
API 或 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_enable
TensorRT 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。
?