From c2d4493705f0a215c2eb8f7ce344523dcdea084b Mon Sep 17 00:00:00 2001 From: Kevin Chen Date: Tue, 15 Jul 2025 13:04:50 -0700 Subject: [PATCH 1/2] Add loadModelProto APIs --- .../tensorrt/tensorrt_provider_options.h | 11 +- .../core/graph/graph_proto_serializer.cc | 4 +- .../tensorrt/onnx_ctx_model_helper.cc | 4 + .../tensorrt/onnx_ctx_model_helper.h | 6 + .../tensorrt/tensorrt_execution_provider.cc | 221 ++++++++++++++++-- .../tensorrt/tensorrt_execution_provider.h | 14 ++ .../tensorrt_execution_provider_info.cc | 24 ++ .../tensorrt_execution_provider_info.h | 3 + .../tensorrt/tensorrt_provider_factory.cc | 3 + 9 files changed, 271 insertions(+), 19 deletions(-) diff --git a/include/onnxruntime/core/providers/tensorrt/tensorrt_provider_options.h b/include/onnxruntime/core/providers/tensorrt/tensorrt_provider_options.h index 9fb1eb9107774..0aad80c4ddab9 100644 --- a/include/onnxruntime/core/providers/tensorrt/tensorrt_provider_options.h +++ b/include/onnxruntime/core/providers/tensorrt/tensorrt_provider_options.h @@ -89,7 +89,12 @@ struct OrtTensorRTProviderOptionsV2 { size_t trt_onnx_bytestream_size{0}; // size of the byte stream provided as "trt_onnx_bytestream" // can be updated using: UpdateTensorRTProviderOptionsWithValue - const char* trt_engine_cache_prefix{nullptr}; // specify engine cache prefix - int trt_engine_hw_compatible{0}; // Enable hardware compatibility. Default 0 = false, nonzero = true - const char* trt_op_types_to_exclude{}; // Exclude specific ops from running on TRT. + const void* trt_external_data_bytestream{nullptr}; // The byte stream containing the weights to override the ones provided in the ONNX model. + // can be updated using: UpdateTensorRTProviderOptionsWithValue + size_t trt_external_data_bytestream_size{0}; // size of the byte stream provided as "trt_external_data_bytestream" + // can be updated using: UpdateTensorRTProviderOptionsWithValue + const char* trt_engine_cache_prefix{nullptr}; // specify engine cache prefix + int trt_engine_hw_compatible{0}; // Enable hardware compatibility. Default 0 = false, nonzero = true + const char* trt_op_types_to_exclude{}; // Exclude specific ops from running on TRT. + int trt_load_user_initializer{0}; // Save initializers locally instead of to disk. Default 0 = false, nonzero = true }; diff --git a/onnxruntime/core/graph/graph_proto_serializer.cc b/onnxruntime/core/graph/graph_proto_serializer.cc index 993020278eb03..0fbcea2719ce8 100644 --- a/onnxruntime/core/graph/graph_proto_serializer.cc +++ b/onnxruntime/core/graph/graph_proto_serializer.cc @@ -94,8 +94,8 @@ void GraphViewerToProto(const GraphViewer& graph_view, current_scope_initializer_set.insert(name); auto* p_initializer = graph_proto.add_initializer(); - // Do not save raw or external data into the graph, only the metadata - if (!include_initializer_data && (init->has_raw_data() || init->has_data_location())) { + // Do not save raw into the graph, only the metadata + if (!include_initializer_data && init->has_raw_data()) { // Set datatype if (init->has_data_type()) { p_initializer->set_data_type(init->data_type()); diff --git a/onnxruntime/core/providers/tensorrt/onnx_ctx_model_helper.cc b/onnxruntime/core/providers/tensorrt/onnx_ctx_model_helper.cc index fbccd7d4a286b..69fcda98dc75f 100644 --- a/onnxruntime/core/providers/tensorrt/onnx_ctx_model_helper.cc +++ b/onnxruntime/core/providers/tensorrt/onnx_ctx_model_helper.cc @@ -298,6 +298,8 @@ Status TensorRTCacheModelHandler::GetEpContextFromGraph(const GraphViewer& graph make_secure_path_checks, onnx_model_bytestream_, onnx_model_bytestream_size_, + onnx_external_data_bytestream_, + onnx_external_data_bytestream_size_, (*trt_engine_).get(), false /* serialize refitted engine to disk */, detailed_build_log_); @@ -367,6 +369,8 @@ Status TensorRTCacheModelHandler::GetEpContextFromGraph(const GraphViewer& graph make_secure_path_checks, onnx_model_bytestream_, onnx_model_bytestream_size_, + onnx_external_data_bytestream_, + onnx_external_data_bytestream_size_, (*trt_engine_).get(), true /* serialize refitted engine to disk */, detailed_build_log_); diff --git a/onnxruntime/core/providers/tensorrt/onnx_ctx_model_helper.h b/onnxruntime/core/providers/tensorrt/onnx_ctx_model_helper.h index 3af0143cbf14e..e89b047919cd6 100644 --- a/onnxruntime/core/providers/tensorrt/onnx_ctx_model_helper.h +++ b/onnxruntime/core/providers/tensorrt/onnx_ctx_model_helper.h @@ -54,6 +54,8 @@ class TensorRTCacheModelHandler { std::string onnx_model_folder_path, const void* onnx_model_bytestream, size_t onnx_model_bytestream_size, + const void* onnx_external_data_bytestream, + size_t onnx_external_data_bytestream_size, bool detailed_build_log) : trt_engine_(trt_engine), trt_runtime_(trt_runtime), @@ -63,6 +65,8 @@ class TensorRTCacheModelHandler { onnx_model_folder_path_(onnx_model_folder_path), onnx_model_bytestream_(onnx_model_bytestream), onnx_model_bytestream_size_(onnx_model_bytestream_size), + onnx_external_data_bytestream_(onnx_external_data_bytestream), + onnx_external_data_bytestream_size_(onnx_external_data_bytestream_size), detailed_build_log_(detailed_build_log) { } ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(TensorRTCacheModelHandler); @@ -80,6 +84,8 @@ class TensorRTCacheModelHandler { std::string onnx_model_folder_path_; const void* onnx_model_bytestream_; size_t onnx_model_bytestream_size_; + const void* onnx_external_data_bytestream_; + size_t onnx_external_data_bytestream_size_; bool detailed_build_log_; }; // TRTCacheModelHandler } // namespace onnxruntime diff --git a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc index 1121775bf5ef7..31354ef543cb9 100644 --- a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc +++ b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc @@ -1395,6 +1395,14 @@ TensorrtExecutionProvider::TensorrtExecutionProvider(const TensorrtExecutionProv "When providing either 'trt_onnx_bytestream_size' or " "'trt_onnx_bytestream' both have to be provided")); } + onnx_external_data_bytestream_ = info.external_data_bytestream; + onnx_external_data_bytestream_size_ = info.external_data_bytestream_size; + if ((onnx_external_data_bytestream_ != nullptr && onnx_external_data_bytestream_size_ == 0) || + (onnx_external_data_bytestream_ == nullptr && onnx_external_data_bytestream_size_ != 0)) { + ORT_THROW_IF_ERROR(ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, + "When providing either 'trt_external_data_bytestream_size' or " + "'trt_external_data_bytestream' both have to be provided")); + } timing_cache_enable_ = info.timing_cache_enable; force_timing_cache_match_ = info.force_timing_cache; detailed_build_log_ = info.detailed_build_log; @@ -1435,6 +1443,7 @@ TensorrtExecutionProvider::TensorrtExecutionProvider(const TensorrtExecutionProv engine_hw_compatible_ = info.engine_hw_compatible; op_types_to_exclude_ = info.op_types_to_exclude; preview_features_ = ParseTrtPreviewFeatures(info.preview_features); + load_user_initializer_ = info.load_user_initializer; } else { try { const std::string max_partition_iterations_env = onnxruntime::GetEnvironmentVar(tensorrt_env_vars::kMaxPartitionIterations); @@ -1836,7 +1845,9 @@ TensorrtExecutionProvider::TensorrtExecutionProvider(const TensorrtExecutionProv << ", trt_cache_prefix: " << cache_prefix_ << ", trt_engine_hw_compatible: " << engine_hw_compatible_ << ", trt_onnx_model_bytestream_size_: " << onnx_model_bytestream_size_ - << ", trt_op_types_to_exclude: " << op_types_to_exclude_; + << ", trt_onnx_external_data_bytestream_size: " << onnx_external_data_bytestream_size_ + << ", trt_op_types_to_exclude: " << op_types_to_exclude_ + << ", trt_load_user_initializer: " << load_user_initializer_; } TensorrtExecutionProvider::~TensorrtExecutionProvider() { @@ -2318,7 +2329,24 @@ SubGraphCollection_t TensorrtExecutionProvider::GetSupportedList(SubGraphCollect // When creating model proto from graph viewer, let ORT use priority-based topological sort based on node index. // The reason is, in some cases, for example ResNet50, using default topological sort will end up with generating // the model proto that has different node ordering compared to original onnx model. - graph_viewer->ToProto(*model_proto->mutable_graph(), true, true, 1 /*priority-based topological sort*/); + // Save Initializer Data. + + std::vector userWeights; + + // Keep inits in memory instead of writing to ModelProto. + if (load_user_initializer_) { + auto allInitializers = graph_viewer->GetAllInitializedTensors(); + + for (auto entry : allInitializers) { + auto* tp = entry.second; + if (tp->has_raw_data()) { + userWeights.push_back( + TensorrtUserWeights{tp->name(), tp->raw_data(), (int64_t)tp->raw_data().size()}); + } + } + } + + graph_viewer->ToProto(*model_proto->mutable_graph(), true, true, 1 /*priority-based topological sort*/, !load_user_initializer_ /*include_initializer_data*/); model_proto->set_ir_version(ONNX_NAMESPACE::Version::IR_VERSION); std::string string_buf; @@ -2343,9 +2371,22 @@ SubGraphCollection_t TensorrtExecutionProvider::GetSupportedList(SubGraphCollect auto trt_network = std::unique_ptr(trt_builder->createNetworkV2(network_flags)); auto trt_parser = tensorrt_ptr::unique_pointer(nvonnxparser::createParser(*trt_network, trt_logger)); + bool is_model_supported = false; #if (NV_TENSORRT_MAJOR == 10 && NV_TENSORRT_MINOR > 1) || NV_TENSORRT_MAJOR > 10 - auto is_model_supported = trt_parser->supportsModelV2(string_buf.data(), string_buf.size(), model_path_); +#if (NV_TENSORRT_MAJOR == 10 && NV_TENSORRT_MINOR > 12) || NV_TENSORRT_MAJOR > 10 + if (load_user_initializer_) { + trt_parser->loadModelProto(string_buf.data(), string_buf.size(), model_path_); + for (auto const& userWeight : userWeights) { + trt_parser->loadInitializer(userWeight.name.c_str(), static_cast(userWeight.data.c_str()), userWeight.size); + } + is_model_supported = trt_parser->parseModelProto(); + } else { + is_model_supported = trt_parser->supportsModelV2(string_buf.data(), string_buf.size(), model_path_); + } +#else + is_model_supported = trt_parser->supportsModelV2(string_buf.data(), string_buf.size(), model_path_); +#endif // (NV_TENSORRT_MAJOR == 10 && NV_TENSORRT_MINOR > 12) || NV_TENSORRT_MAJOR > 10 // Note: Calling getNbSubgraphs or getSubgraphNodes before calling supportsModelV2 results in undefined behavior. auto num_subgraphs = trt_parser->getNbSubgraphs(); @@ -2363,7 +2404,7 @@ SubGraphCollection_t TensorrtExecutionProvider::GetSupportedList(SubGraphCollect } #else trt_parser->supportsModel(string_buf.data(), string_buf.size(), parser_nodes_list, model_path_); -#endif +#endif // (NV_TENSORRT_MAJOR == 10 && NV_TENSORRT_MINOR > 1) || NV_TENSORRT_MAJOR > 10 SubGraphCollection_t next_nodes_list; const std::vector& subgraph_node_index = graph_viewer->GetNodesInTopologicalOrder(1 /*priority-based topological sort*/); @@ -2804,11 +2845,14 @@ common::Status TensorrtExecutionProvider::RefitEngine(std::string onnx_model_fil bool path_check, const void* onnx_model_bytestream, size_t onnx_model_bytestream_size, + const void* onnx_external_data_bytestream, + size_t onnx_external_data_bytestream_size, nvinfer1::ICudaEngine* trt_engine, bool serialize_refitted_engine, bool detailed_build_log) { #if NV_TENSORRT_MAJOR >= 10 bool refit_from_file = onnx_model_bytestream == nullptr && onnx_model_bytestream_size == 0; + bool refit_complete = false; std::filesystem::path onnx_model_path{onnx_model_folder_path}; if (refit_from_file) { if (!onnx_model_filename.empty()) { @@ -2845,17 +2889,131 @@ common::Status TensorrtExecutionProvider::RefitEngine(std::string onnx_model_fil auto refitter = std::unique_ptr(nvinfer1::createInferRefitter(*trt_engine, trt_logger)); auto parser_refitter = std::unique_ptr( nvonnxparser::createParserRefitter(*refitter, trt_logger)); - if (refit_from_file) { - LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Refitting from file on disk: " << onnx_model_path.string(); - if (!parser_refitter->refitFromFile(onnx_model_path.string().c_str())) { + +#if (NV_TENSORRT_MAJOR == 10 && NV_TENSORRT_MINOR > 12) || NV_TENSORRT_MAJOR > 10 + bool refit_with_external_data = onnx_external_data_bytestream != nullptr && onnx_external_data_bytestream_size != 0; + + // New refit APIs + if (refit_with_external_data) { + // A valid model bytestream must be passed. + if (refit_from_file) { return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "TensorRT EP's IParserRefitter could not refit deserialized weight-stripped engine with weights contained in: " + onnx_model_path.string()); + "TensorRT EP's refit with external data must be called with a valid ONNX model bytestream"); } - } else { - LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Refitting from byte array"; - if (!parser_refitter->refitFromBytes(onnx_model_bytestream, onnx_model_bytestream_size)) { + + if (!parser_refitter->loadModelProto(onnx_model_bytestream, onnx_model_bytestream_size, nullptr)) { + return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, + "TensorRT EP's IParserRefitter could not load model from provided onnx_model_bytestream"); + } + + // Extract weight information from the Refitter. + int required_weights = refitter->getAllWeights(0, nullptr); + std::vector refit_names(required_weights); + refitter->getAllWeights(required_weights, refit_names.data()); + + // Vectors to keep track of data pointers. + std::vector names; + names.reserve(required_weights); + std::vector bytes; + bytes.reserve(required_weights); + std::vector sizes; + sizes.reserve(required_weights); + + if (refit_with_external_data) { + auto onnx_model = ModelProto::Create(); + TensorProtos* allInitializers_byte_stream; + + // Reconstruct onnx model view. + const auto onnx_model_view = std::string((const char*)onnx_model_bytestream, + onnx_model_bytestream_size); + if (!onnx_model->ParseFromString(onnx_model_view)) { + return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, + "The provided ONNX bytestream to refit could not be parsed."); + } + + // Extract graph and initializer information. + auto const& graph = onnx_model->mutable_graph(); + allInitializers_byte_stream = graph->mutable_initializer(); + LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Initializers that were found " << allInitializers_byte_stream->size(); + + // Loop through all initializers + for (int initializer_idx = 0; initializer_idx < allInitializers_byte_stream->size(); ++initializer_idx) { + auto& proto = allInitializers_byte_stream->at(initializer_idx); + auto& proto_name = proto.name(); + bool weight_is_refittable = std::find(refit_names.begin(), refit_names.end(), proto_name) != refit_names.end(); + if (weight_is_refittable) { + if (proto.has_data_location()) { + if (proto.data_location() == TensorProto_DataLocation_EXTERNAL) { + // Default values for reading into external_data blob. + int64_t offset = 0; + size_t length = 0; + auto external_data = proto.mutable_external_data(); + const std::string kOffset = "offset", kLength = "length"; + for (int entry_idx = 0; entry_idx < external_data->size(); ++entry_idx) { + auto current_key = external_data->at(entry_idx).mutable_key(); + auto current_value = external_data->at(entry_idx).mutable_value(); + if (*current_key == kOffset && !current_value->empty()) { + offset = std::stoll(*current_value); + } else if (*current_key == kLength && !current_value->empty()) { + length = std::stoul(*current_value); + } + } + names.push_back(proto.name()); + bytes.push_back(static_cast(onnx_external_data_bytestream) + offset); + sizes.push_back(length); + } else { + return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, + "[TensorRT EP] Proto: " + proto_name + " has default as data location which is not supported"); + } + } else { + if (!proto.has_raw_data()) { + return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, + "[TensorRT EP] Proto: " + proto_name + " has no raw data"); + } + auto& raw_data = proto.raw_data(); + names.push_back(proto.name()); + bytes.push_back(raw_data.c_str()); + sizes.push_back(raw_data.size()); + } + } else { + LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Initializer with name: " << proto_name << " was not marked as refittable"; + } + } + } + + // Load extracted initializers into the parser + if (!names.empty()) { + LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Number of initializers submitted to refitter " << names.size(); + for (size_t i = 0; i < names.size(); i++) { + bool refloadInit = parser_refitter->loadInitializer(names[i].c_str(), bytes[i], sizes[i]); + if (!refloadInit) { + return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, + "TensorRT EP's IParserRefitter could not refit deserialized weight-stripped engine with weights contained in the provided bytestraem"); + } + } + } + // Perform refit. + if (!parser_refitter->refitModelProto()) { return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "TensorRT EP's IParserRefitter could not refit deserialized weight-stripped engine with weights contained in the provided bytestraem"); + "TensorRT EP's IParserRefitter refitModelProto() failed with the provided external data bytestream."); + } + refit_complete = true; + } +#endif // (NV_TENSORRT_MAJOR == 10 && NV_TENSORRT_MINOR > 12) || NV_TENSORRT_MAJOR > 10 + // If new refit flow was not completed, then fallback to refit_from_file. + if (!refit_complete) { + if (refit_from_file) { + LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Refitting from file on disk: " << onnx_model_path.string(); + if (!parser_refitter->refitFromFile(onnx_model_path.string().c_str())) { + return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, + "TensorRT EP's IParserRefitter could not refit deserialized weight-stripped engine with weights contained in: " + onnx_model_path.string()); + } + } else { + LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Refitting from byte array"; + if (!parser_refitter->refitFromBytes(onnx_model_bytestream, onnx_model_bytestream_size)) { + return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, + "TensorRT EP's IParserRefitter could not refit deserialized weight-stripped engine with weights contained in the provided bytestream"); + } } } if (refitter->refitCudaEngine()) { @@ -2926,11 +3084,26 @@ Status TensorrtExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphView auto model = graph_body_viewer.CreateModel(*GetLogger()); auto model_proto = model->ToProto(); + auto userWeights = std::make_unique>(); + + if (load_user_initializer_) { + auto allInitializers = graph_body_viewer.GetAllInitializedTensors(); + + for (auto entry : allInitializers) { + auto name = entry.first; + auto* tp = entry.second; + if (tp->has_raw_data()) { + userWeights->push_back( + TensorrtUserWeights{tp->name(), tp->raw_data(), (int64_t)tp->raw_data().size()}); + } + } + } + // ORT's default topological sort is using reversed DFS. // When creating model proto from graph viewer, let ORT use priority-based topological sort based on node index. // The reason is, in some cases, for example ResNet50, using default topological sort will end up with generating // the model proto that has different node ordering compared to original onnx model. - graph_body_viewer.ToProto(*model_proto->mutable_graph(), true, true, 1 /*priority-based topological sort*/); + graph_body_viewer.ToProto(*model_proto->mutable_graph(), true, true, 1 /*priority-based topological sort*/, !load_user_initializer_ /*include_initializer_data*/); model_proto->set_ir_version(ONNX_NAMESPACE::Version::IR_VERSION); std::string string_buf; model_proto->SerializeToString(string_buf); @@ -2952,7 +3125,20 @@ Status TensorrtExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphView auto trt_network = std::unique_ptr(trt_builder->createNetworkV2(network_flags)); auto trt_config = std::unique_ptr(trt_builder->createBuilderConfig()); auto trt_parser = tensorrt_ptr::unique_pointer(nvonnxparser::createParser(*trt_network, trt_logger)); + +#if (NV_TENSORRT_MAJOR == 10 && NV_TENSORRT_MINOR > 12) || NV_TENSORRT_MAJOR > 10 + if (load_user_initializer_) { + trt_parser->loadModelProto(string_buf.data(), string_buf.size(), model_path_); + for (auto const& userWeight : *userWeights) { + trt_parser->loadInitializer(userWeight.name.c_str(), static_cast(userWeight.data.c_str()), userWeight.size); + } + trt_parser->parseModelProto(); + } else { + trt_parser->parse(string_buf.data(), string_buf.size(), model_path_); + } +#else trt_parser->parse(string_buf.data(), string_buf.size(), model_path_); +#endif if (max_workspace_size_ > 0) { trt_config->setMemoryPoolLimit(nvinfer1::MemoryPoolType::kWORKSPACE, max_workspace_size_); } @@ -3489,6 +3675,8 @@ Status TensorrtExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphView false /* path check for security */, onnx, onnx_size, + onnx_external_data_bytestream_, + onnx_external_data_bytestream_size_, trt_engine.get(), true /* serialize refitted engine to disk */, detailed_build_log_); @@ -3545,6 +3733,7 @@ Status TensorrtExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphView engines_.emplace(fused_node.Name(), std::move(trt_engine)); contexts_.emplace(fused_node.Name(), std::move(trt_context)); networks_.emplace(fused_node.Name(), std::move(trt_network)); + weights_.emplace(fused_node.Name(), std::move(userWeights)); input_info_[fused_node.Name()].push_back(input_indexes); output_info_[fused_node.Name()].push_back(output_indexes); output_info_[fused_node.Name()].push_back(output_types); @@ -3597,7 +3786,7 @@ Status TensorrtExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphView engine_decryption_, engine_encryption_, timing_cache_enable_, global_cache_path_, force_timing_cache_match_, detailed_build_log_, build_heuristics_enable_, sparsity_enable_, builder_optimization_level_, auxiliary_streams_, !tactic_sources_.empty(), tactics, cuda_graph_enable_, cache_prefix_, cache_suffix, engine_hw_compatible_, - preview_features_}; + preview_features_, &weights_[context->node_name]}; *state = p.release(); return 0; }; @@ -3975,6 +4164,8 @@ Status TensorrtExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphView false /* path check for security */, onnx_model_bytestream_, onnx_model_bytestream_size_, + onnx_external_data_bytestream_, + onnx_external_data_bytestream_size_, trt_engine, true /* serialize refitted engine to disk */, detailed_build_log_); @@ -4208,6 +4399,8 @@ Status TensorrtExecutionProvider::CreateNodeComputeInfoFromPrecompiledEngine(con onnx_model_folder_path_, onnx_model_bytestream_, onnx_model_bytestream_size_, + onnx_external_data_bytestream_, + onnx_external_data_bytestream_size_, detailed_build_log_); auto status = trt_cache_model_handler.GetEpContextFromGraph(graph_body_viewer); if (status != Status::OK()) { diff --git a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.h b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.h index 7e02cf7590f66..dba17f7822eac 100644 --- a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.h +++ b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.h @@ -157,6 +157,13 @@ class OutputAllocator : public nvinfer1::IOutputAllocator { */ using ShapeRangesMap = std::unordered_map>>>; +// Struct to hold user weights when ModelProtos are serialized with data. +struct TensorrtUserWeights { + std::string name{}; + std::string data{}; + int64_t size{}; +}; + // Information to construct kernel function state. struct TensorrtFuncState { AllocateFunc test_allocate_func = nullptr; @@ -205,6 +212,7 @@ struct TensorrtFuncState { std::string cache_suffix; bool engine_hw_compatible = false; std::vector preview_features; + std::unique_ptr>* userWeights = nullptr; }; // Minimum information to construct kernel function state for direct engine load code path @@ -287,6 +295,8 @@ class TensorrtExecutionProvider : public IExecutionProvider { bool path_check, const void* onnx_model_bytestream, size_t onnx_model_bytestream_size, + const void* onnx_external_data_bytestream, + size_t onnx_external_data_bytestream_size, nvinfer1::ICudaEngine* trt_engine, bool serialize_refitted_engine, bool detailed_build_log); @@ -314,6 +324,8 @@ class TensorrtExecutionProvider : public IExecutionProvider { std::string onnx_model_folder_path_; const void* onnx_model_bytestream_; size_t onnx_model_bytestream_size_; + const void* onnx_external_data_bytestream_ = nullptr; + size_t onnx_external_data_bytestream_size_ = 0; bool build_heuristics_enable_ = false; bool sparsity_enable_ = false; int builder_optimization_level_ = 3; @@ -340,6 +352,7 @@ class TensorrtExecutionProvider : public IExecutionProvider { bool engine_hw_compatible_ = false; std::string op_types_to_exclude_; std::vector preview_features_; + bool load_user_initializer_ = false; // The format is as for TENSORRT_VERSION: (MAJOR * 100 + MINOR) * 100 + PATCH int32_t trt_version_; @@ -380,6 +393,7 @@ class TensorrtExecutionProvider : public IExecutionProvider { std::unordered_map input_shape_ranges_; // The profile shape ranges that the engine is built with std::unordered_map> profiles_; std::unordered_map dds_output_allocator_maps_; + std::unordered_map>> weights_; // User provided weights. // for external stream, we need to create its cudnn/cublass handle before cuda EP enable cuda graph capture cudnnHandle_t external_cudnn_handle_ = nullptr; diff --git a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider_info.cc b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider_info.cc index 1a515c37f7ecb..17457b11f8cef 100644 --- a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider_info.cc +++ b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider_info.cc @@ -57,8 +57,11 @@ constexpr const char* kDumpEpContextModel = "trt_dump_ep_context_model"; constexpr const char* kEngineHwCompatible = "trt_engine_hw_compatible"; constexpr const char* kONNXBytestream = "trt_onnx_bytestream"; constexpr const char* kONNXBytestreamSize = "trt_onnx_bytestream_size"; +constexpr const char* kExternalDataBytestream = "trt_external_data_bytestream"; +constexpr const char* kExternalDataBytestreamSize = "trt_external_data_bytestream_size"; constexpr const char* kOpTypesToExclude = "trt_op_types_to_exclude"; constexpr const char* kPreviewFeatures = "trt_preview_features"; +constexpr const char* kGraphIncludeInitializer = "trt_load_user_initializer"; } // namespace provider_option_names } // namespace tensorrt @@ -67,6 +70,7 @@ TensorrtExecutionProviderInfo TensorrtExecutionProviderInfo::FromProviderOptions TensorrtExecutionProviderInfo info{}; void* user_compute_stream = nullptr; void* onnx_bytestream = nullptr; + void* external_data_bytestream = nullptr; ORT_THROW_IF_ERROR( ProviderOptionsParser{} .AddValueParser( @@ -138,13 +142,24 @@ TensorrtExecutionProviderInfo TensorrtExecutionProviderInfo::FromProviderOptions return Status::OK(); }) .AddAssignmentToReference(tensorrt::provider_option_names::kONNXBytestreamSize, info.onnx_bytestream_size) + .AddValueParser( + tensorrt::provider_option_names::kExternalDataBytestream, + [&external_data_bytestream](const std::string& value_str) -> Status { + size_t address; + ORT_RETURN_IF_ERROR(ParseStringWithClassicLocale(value_str, address)); + external_data_bytestream = reinterpret_cast(address); + return Status::OK(); + }) + .AddAssignmentToReference(tensorrt::provider_option_names::kExternalDataBytestreamSize, info.external_data_bytestream_size) .AddAssignmentToReference(tensorrt::provider_option_names::kOpTypesToExclude, info.op_types_to_exclude) .AddAssignmentToReference(tensorrt::provider_option_names::kPreviewFeatures, info.preview_features) + .AddAssignmentToReference(tensorrt::provider_option_names::kGraphIncludeInitializer, info.load_user_initializer) .Parse(options)); // add new provider option here. info.user_compute_stream = user_compute_stream; info.has_user_compute_stream = (user_compute_stream != nullptr); info.onnx_bytestream = onnx_bytestream; + info.external_data_bytestream = external_data_bytestream; return info; } @@ -195,8 +210,11 @@ ProviderOptions TensorrtExecutionProviderInfo::ToProviderOptions(const TensorrtE {tensorrt::provider_option_names::kEngineHwCompatible, MakeStringWithClassicLocale(info.engine_hw_compatible)}, {tensorrt::provider_option_names::kONNXBytestream, MakeStringWithClassicLocale(info.onnx_bytestream)}, {tensorrt::provider_option_names::kONNXBytestreamSize, MakeStringWithClassicLocale(info.onnx_bytestream_size)}, + {tensorrt::provider_option_names::kExternalDataBytestream, MakeStringWithClassicLocale(info.external_data_bytestream)}, + {tensorrt::provider_option_names::kExternalDataBytestreamSize, MakeStringWithClassicLocale(info.external_data_bytestream_size)}, {tensorrt::provider_option_names::kOpTypesToExclude, MakeStringWithClassicLocale(info.op_types_to_exclude)}, {tensorrt::provider_option_names::kPreviewFeatures, MakeStringWithClassicLocale(info.preview_features)}, + {tensorrt::provider_option_names::kGraphIncludeInitializer, MakeStringWithClassicLocale(info.load_user_initializer)}, }; return options; } @@ -262,7 +280,10 @@ ProviderOptions TensorrtExecutionProviderInfo::ToProviderOptions(const OrtTensor {tensorrt::provider_option_names::kEngineHwCompatible, MakeStringWithClassicLocale(info.trt_engine_hw_compatible)}, {tensorrt::provider_option_names::kONNXBytestream, MakeStringWithClassicLocale(reinterpret_cast(info.trt_onnx_bytestream))}, {tensorrt::provider_option_names::kONNXBytestreamSize, MakeStringWithClassicLocale(info.trt_onnx_bytestream_size)}, + {tensorrt::provider_option_names::kExternalDataBytestream, MakeStringWithClassicLocale(reinterpret_cast(info.trt_external_data_bytestream))}, + {tensorrt::provider_option_names::kExternalDataBytestreamSize, MakeStringWithClassicLocale(info.trt_external_data_bytestream_size)}, {tensorrt::provider_option_names::kOpTypesToExclude, kOpTypesToExclude_}, + {tensorrt::provider_option_names::kGraphIncludeInitializer, MakeStringWithClassicLocale(info.trt_load_user_initializer)}, }; return options; } @@ -368,7 +389,10 @@ void TensorrtExecutionProviderInfo::UpdateProviderOptions(void* provider_options trt_provider_options_v2.trt_engine_hw_compatible = internal_options.engine_hw_compatible; trt_provider_options_v2.trt_onnx_bytestream = internal_options.onnx_bytestream; trt_provider_options_v2.trt_onnx_bytestream_size = internal_options.onnx_bytestream_size; + trt_provider_options_v2.trt_external_data_bytestream = internal_options.external_data_bytestream; + trt_provider_options_v2.trt_external_data_bytestream_size = internal_options.external_data_bytestream_size; trt_provider_options_v2.trt_op_types_to_exclude = copy_string_if_needed(internal_options.op_types_to_exclude); trt_provider_options_v2.trt_preview_features = copy_string_if_needed(internal_options.preview_features); + trt_provider_options_v2.trt_load_user_initializer = internal_options.load_user_initializer; } } // namespace onnxruntime diff --git a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider_info.h b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider_info.h index a7c3624674dc6..f0bd653de471a 100644 --- a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider_info.h +++ b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider_info.h @@ -37,6 +37,8 @@ struct TensorrtExecutionProviderInfo { std::string onnx_model_folder_path{""}; const void* onnx_bytestream{nullptr}; size_t onnx_bytestream_size{0}; + const void* external_data_bytestream{nullptr}; + size_t external_data_bytestream_size{0}; bool engine_decryption_enable{false}; std::string engine_decryption_lib_path{""}; bool force_sequential_engine_build{false}; @@ -63,6 +65,7 @@ struct TensorrtExecutionProviderInfo { bool engine_hw_compatible{false}; std::string op_types_to_exclude{""}; std::string preview_features{""}; + bool load_user_initializer{false}; static TensorrtExecutionProviderInfo FromProviderOptions(const ProviderOptions& options); static ProviderOptions ToProviderOptions(const TensorrtExecutionProviderInfo& info); diff --git a/onnxruntime/core/providers/tensorrt/tensorrt_provider_factory.cc b/onnxruntime/core/providers/tensorrt/tensorrt_provider_factory.cc index da1c2514bf6a2..71ea66b0be89f 100644 --- a/onnxruntime/core/providers/tensorrt/tensorrt_provider_factory.cc +++ b/onnxruntime/core/providers/tensorrt/tensorrt_provider_factory.cc @@ -119,8 +119,11 @@ struct Tensorrt_Provider : Provider { info.engine_hw_compatible = options.trt_engine_hw_compatible != 0; info.onnx_bytestream = options.trt_onnx_bytestream; info.onnx_bytestream_size = options.trt_onnx_bytestream_size; + info.external_data_bytestream = options.trt_external_data_bytestream; + info.external_data_bytestream_size = options.trt_external_data_bytestream_size; info.op_types_to_exclude = options.trt_op_types_to_exclude == nullptr ? "" : options.trt_op_types_to_exclude; info.preview_features = options.trt_preview_features == nullptr ? "" : options.trt_preview_features; + info.load_user_initializer = options.trt_load_user_initializer != 0; return std::make_shared(info); } From c7360f19f92114b22547c990f0d056df1ddd259a Mon Sep 17 00:00:00 2001 From: Kevin Chen Date: Wed, 16 Jul 2025 09:10:31 -0700 Subject: [PATCH 2/2] Address review comments, fix build Signed-off-by: Kevin Chen --- .../tensorrt/tensorrt_execution_provider.cc | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc index 31354ef543cb9..64be445b4c15c 100644 --- a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc +++ b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc @@ -2852,6 +2852,7 @@ common::Status TensorrtExecutionProvider::RefitEngine(std::string onnx_model_fil bool detailed_build_log) { #if NV_TENSORRT_MAJOR >= 10 bool refit_from_file = onnx_model_bytestream == nullptr && onnx_model_bytestream_size == 0; + bool refit_with_external_data = onnx_external_data_bytestream != nullptr && onnx_external_data_bytestream_size != 0; bool refit_complete = false; std::filesystem::path onnx_model_path{onnx_model_folder_path}; if (refit_from_file) { @@ -2891,8 +2892,6 @@ common::Status TensorrtExecutionProvider::RefitEngine(std::string onnx_model_fil nvonnxparser::createParserRefitter(*refitter, trt_logger)); #if (NV_TENSORRT_MAJOR == 10 && NV_TENSORRT_MINOR > 12) || NV_TENSORRT_MAJOR > 10 - bool refit_with_external_data = onnx_external_data_bytestream != nullptr && onnx_external_data_bytestream_size != 0; - // New refit APIs if (refit_with_external_data) { // A valid model bytestream must be passed. @@ -2963,7 +2962,7 @@ common::Status TensorrtExecutionProvider::RefitEngine(std::string onnx_model_fil sizes.push_back(length); } else { return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "[TensorRT EP] Proto: " + proto_name + " has default as data location which is not supported"); + "[TensorRT EP] Proto: " + proto_name + " expected to have external datalocation, but default datalocation was provided instead."); } } else { if (!proto.has_raw_data()) { @@ -2988,7 +2987,7 @@ common::Status TensorrtExecutionProvider::RefitEngine(std::string onnx_model_fil bool refloadInit = parser_refitter->loadInitializer(names[i].c_str(), bytes[i], sizes[i]); if (!refloadInit) { return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "TensorRT EP's IParserRefitter could not refit deserialized weight-stripped engine with weights contained in the provided bytestraem"); + "TensorRT EP's IParserRefitter could not refit deserialized weight-stripped engine with weights contained in the provided bytestream"); } } } @@ -2999,6 +2998,11 @@ common::Status TensorrtExecutionProvider::RefitEngine(std::string onnx_model_fil } refit_complete = true; } +#else + // Refitting with external data is not supported prior to TensorRT 10.13. Log a warning in this case for the user. + if (refit_with_external_data) { + LOGS_DEFAULT(WARNING) << "[TensorRT EP] Refitting with an onnx_external_data_bytestream is only supported on TensorRT versions >= 10.13! This parameter will be ignored for refitting, and the resulting refitted engine may be incorrect."; + } #endif // (NV_TENSORRT_MAJOR == 10 && NV_TENSORRT_MINOR > 12) || NV_TENSORRT_MAJOR > 10 // If new refit flow was not completed, then fallback to refit_from_file. if (!refit_complete) {