Support direct usage of ORT format model flatbuffer for initializers (#12465)

* Add ability to use ORT format model flatbuffer directly for intiializers by leveraging the TensorProto external data infrastructure.

Requires user to provide ORT format model bytes when creating the session, and set both `session.use_ort_model_bytes_directly` and `session.use_ort_model_bytes_for_initializers` to 1 in SessionOptions config entries (AddSessionConfigEntry in C API).

include/onnxruntime/core/graph/graph.h

@@ -485,9 +485,12 @@ class Node {
 #endif
 
   static Status LoadFromOrtFormat(const onnxruntime::fbs::Node& fbs_node, Graph& graph,
+                                  bool can_use_flatbuffer_for_initializers,
                                   const logging::Logger& logger, std::unique_ptr<Node>& node);
 
-  Status LoadFromOrtFormat(const onnxruntime::fbs::Node& fbs_node, const logging::Logger& logger);
+  Status LoadFromOrtFormat(const onnxruntime::fbs::Node& fbs_node,
+                           bool can_use_flatbuffer_for_initializers,
+                           const logging::Logger& logger);
   Status LoadEdgesFromOrtFormat(const onnxruntime::fbs::NodeEdge& fbs_node_edgs, const Graph& graph);
 
   /**
@@ -1314,17 +1317,18 @@ class Graph {
 
   virtual ~Graph();
 
-  static common::Status LoadFromOrtFormat(
-      const onnxruntime::fbs::Graph& fbs_graph, const Model& owning_model,
-      const std::unordered_map<std::string, int>& domain_to_version,
+  static common::Status LoadFromOrtFormat(const onnxruntime::fbs::Graph& fbs_graph, const Model& owning_model,
+                                          const std::unordered_map<std::string, int>& domain_to_version,
 #if !defined(ORT_MINIMAL_BUILD)
-      IOnnxRuntimeOpSchemaCollectionPtr schema_registry,
+                                          IOnnxRuntimeOpSchemaCollectionPtr schema_registry,
 #endif
-      const logging::Logger& logger, std::unique_ptr<Graph>& graph);
+                                          bool can_use_flatbuffer_for_initializers,
+                                          const logging::Logger& logger, std::unique_ptr<Graph>& graph);
 
   // deserialize a subgraph
   static Status LoadFromOrtFormat(const onnxruntime::fbs::Graph& fbs_graph,
                                   Graph& parent_graph, const Node& parent_node,
+                                  bool can_use_flatbuffer_for_initializers,
                                   const logging::Logger& logger, std::unique_ptr<Graph>& graph);
 
 #if !defined(ORT_MINIMAL_BUILD) || defined(ORT_EXTENDED_MINIMAL_BUILD)
@@ -1368,7 +1372,8 @@ class Graph {
         bool strict_shape_type_inference);
 
   // Populate Graph instance from ORT format serialized data.
-  common::Status LoadFromOrtFormat(const onnxruntime::fbs::Graph& fbs_graph);
+  common::Status LoadFromOrtFormat(const onnxruntime::fbs::Graph& fbs_graph,
+                                   bool can_use_flatbuffer_for_initializers);
 
 #if !defined(ORT_MINIMAL_BUILD)
   // Constructor: Given a <GraphProto> loaded from model file, construct

include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h

@@ -79,6 +79,16 @@ static const char* const kOrtSessionOptionsConfigAllowIntraOpSpinning = "session
 // has to guarantee that the model bytes are valid until the ORT session using the model bytes is destroyed.
 static const char* const kOrtSessionOptionsConfigUseORTModelBytesDirectly = "session.use_ort_model_bytes_directly";
 
+/// <summary>
+/// Key for using the ORT format model flatbuffer bytes directly for initializers.
+/// This avoids copying the bytes and reduces peak memory usage during model loading and initialization. 
+/// Requires `session.use_ort_model_bytes_directly` to be true. 
+/// If set, the flatbuffer bytes provided when creating the InferenceSession MUST remain valid for the entire 
+/// duration of the InferenceSession.
+/// </summary>
+static const char* const kOrtSessionOptionsConfigUseORTModelBytesForInitializers =
+    "session.use_ort_model_bytes_for_initializers";
+
 // This should only be specified when exporting an ORT format model for use on a different platform.
 // If the ORT format model will be used on ARM platforms set to "1". For other platforms set to "0"
 // Available since version 1.11.

onnxruntime/core/framework/session_state_utils.cc

@@ -61,12 +61,15 @@ struct ExtDataValueDeleter {
   OrtCallback ext_delete_cb;
   Tensor* p_tensor;
   void operator()(void*) noexcept {
-    ext_delete_cb.f(ext_delete_cb.param);
+    if (ext_delete_cb.f) {
+      ext_delete_cb.f(ext_delete_cb.param);
+    }
+
     delete p_tensor;
   }
 };
 
-// given a tensor proto with externdal data return an OrtValue with a tensor for
+// given a tensor proto with external data return an OrtValue with a tensor for
 // that data; the pointers for the tensor data and the tensor itself are owned
 // by the OrtValue's deleter
 static inline common::Status ExtDataTensorProtoToTensor(const Env& env,
@@ -74,7 +77,6 @@ static inline common::Status ExtDataTensorProtoToTensor(const Env& env,
                                                         const ONNX_NAMESPACE::TensorProto& tensor_proto,
                                                         Tensor& tensor, OrtCallback& ext_data_deleter) {
   ORT_ENFORCE(utils::HasExternalData(tensor_proto));
-  ORT_ENFORCE(!proto_path.empty());
 
   void* ext_data_buf = nullptr;
   SafeInt<size_t> ext_data_len = 0;
@@ -88,8 +90,7 @@ static inline common::Status ExtDataTensorProtoToTensor(const Env& env,
   std::vector<int64_t> tensor_shape_vec = utils::GetTensorShapeFromTensorProto(tensor_proto);
   TensorShape tensor_shape{tensor_shape_vec};
 
-  tensor = Tensor(type, tensor_shape, ext_data_buf, OrtMemoryInfo(CPU,
-                                           OrtAllocatorType::OrtDeviceAllocator));
+  tensor = Tensor(type, tensor_shape, ext_data_buf, OrtMemoryInfo(CPU, OrtAllocatorType::OrtDeviceAllocator));
 
   return common::Status::OK();
 }

onnxruntime/core/framework/tensor_external_data_info.h

@@ -9,20 +9,13 @@
 
 namespace onnxruntime {
 class ExternalDataInfo {
- private:
+ public:
 #ifdef _WIN32
   using OFFSET_TYPE = int64_t;
 #else
   using OFFSET_TYPE = off_t;
 #endif
-  std::basic_string<ORTCHAR_T> rel_path_;
-  OFFSET_TYPE offset_ = 0;
 
-  // 0 means the whole file
-  size_t length_ = 0;
-  std::string checksum_;
-
- public:
   const std::basic_string<ORTCHAR_T>& GetRelPath() const { return rel_path_; }
 
   OFFSET_TYPE GetOffset() const { return offset_; }
@@ -32,7 +25,16 @@ class ExternalDataInfo {
 
   // If the value of 'offset' or 'length' field is larger the max value of ssize_t, this function will treat it as a
   // wrong value and return FAIL.
-  static common::Status Create(const ::google::protobuf::RepeatedPtrField<::ONNX_NAMESPACE::StringStringEntryProto>& input,
-                               std::unique_ptr<ExternalDataInfo>& out);
+  static common::Status Create(
+      const ::google::protobuf::RepeatedPtrField<::ONNX_NAMESPACE::StringStringEntryProto>& input,
+      std::unique_ptr<ExternalDataInfo>& out);
+
+ private:
+  std::basic_string<ORTCHAR_T> rel_path_;
+  OFFSET_TYPE offset_ = 0;
+
+  // 0 means the whole file
+  size_t length_ = 0;
+  std::string checksum_;
 };
-}  // namespace onnxruntime
+}  // namespace onnxruntime

onnxruntime/core/framework/tensorprotoutils.cc

@@ -131,20 +131,26 @@ static Status GetExternalDataInfo(const ONNX_NAMESPACE::TensorProto& tensor_prot
   std::unique_ptr<onnxruntime::ExternalDataInfo> external_data_info;
   ORT_RETURN_IF_ERROR(onnxruntime::ExternalDataInfo::Create(tensor_proto.external_data(), external_data_info));
 
-  if (tensor_proto_dir != nullptr) {
-    external_file_path = onnxruntime::ConcatPathComponent<ORTCHAR_T>(tensor_proto_dir, external_data_info->GetRelPath());
-  } else {
-    external_file_path = external_data_info->GetRelPath();
-  }
+  const auto& location = external_data_info->GetRelPath();
 
-  file_offset = external_data_info->GetOffset();
+  if (location == onnxruntime::utils::kTensorProtoMemoryAddressTag) {
+    external_file_path = location;
+  } else {
+    if (tensor_proto_dir != nullptr) {
+      external_file_path = onnxruntime::ConcatPathComponent<ORTCHAR_T>(tensor_proto_dir,
+                                                                       external_data_info->GetRelPath());
+    } else {
+      external_file_path = external_data_info->GetRelPath();
+    }
+  }
 
   ORT_RETURN_IF_ERROR(onnxruntime::utils::GetSizeInBytesFromTensorProto<0>(tensor_proto, &tensor_byte_size));
   const size_t external_data_length = external_data_info->GetLength();
-
   ORT_RETURN_IF_NOT(external_data_length == 0 || external_data_length == tensor_byte_size,
-                    "TensorProto: ", tensor_proto.name(), " external data size mismatch. Computed size: ", *&tensor_byte_size,
-                    ", external_data.length: ", external_data_length);
+                    "TensorProto: ", tensor_proto.name(), " external data size mismatch. Computed size: ",
+                    *&tensor_byte_size, ", external_data.length: ", external_data_length);
+
+  file_offset = external_data_info->GetOffset();
 
   return Status::OK();
 }
@@ -597,18 +603,31 @@ Status GetExtDataFromTensorProto(const Env& env, const ORTCHAR_T* model_path,
   ORT_RETURN_IF_ERROR(GetExternalDataInfo(tensor_proto, t_prot_dir_s, external_data_file_path, file_offset,
                                           raw_data_safe_len));
 
-  size_t file_length;
-  ORT_RETURN_IF_ERROR(env.GetFileLength(external_data_file_path.c_str(), file_length));
+  if (external_data_file_path == onnxruntime::utils::kTensorProtoMemoryAddressTag) {
+    // the value in location is the memory address of the data
+    ext_data_buf = reinterpret_cast<void*>(file_offset);
+    ext_data_len = raw_data_safe_len;
+    ext_data_deleter = OrtCallback{nullptr, nullptr};
+  } else {
+    size_t file_length;
+    // error reporting is inconsistent across platforms. Make sure the full path we attempted to open is included.
+    auto status = env.GetFileLength(external_data_file_path.c_str(), file_length);
+    if (!status.IsOK()) {
+      return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "GetFileLength for ", ToUTF8String(external_data_file_path),
+                             " failed:", status.ErrorMessage());
+    }
 
-  SafeInt<FileOffsetType> end_of_read(file_offset);
-  end_of_read += raw_data_safe_len;
-  ORT_RETURN_IF(file_offset < 0 || end_of_read > gsl::narrow<FileOffsetType>(file_length),
+    SafeInt<FileOffsetType> end_of_read(file_offset);
+    end_of_read += raw_data_safe_len;
+    ORT_RETURN_IF(file_offset < 0 || end_of_read > gsl::narrow<FileOffsetType>(file_length),
                   "External initializer: ", tensor_proto.name(),
                   " offset: ", file_offset, " size to read: ", static_cast<size_t>(raw_data_safe_len),
                   " given file_length: ", file_length, " are out of bounds or can not be read in full.");
-  ORT_RETURN_IF_ERROR(GetFileContent(env, external_data_file_path.c_str(), file_offset, raw_data_safe_len,
-                                     ext_data_buf, ext_data_deleter));
-  ext_data_len = raw_data_safe_len;
+    ORT_RETURN_IF_ERROR(GetFileContent(env, external_data_file_path.c_str(), file_offset, raw_data_safe_len,
+                                       ext_data_buf, ext_data_deleter));
+    ext_data_len = raw_data_safe_len;
+  }
+
   return Status::OK();
 }
 

onnxruntime/core/framework/tensorprotoutils.h

@@ -76,6 +76,13 @@ ONNXTensorElementDataType GetTensorElementType(const ONNX_NAMESPACE::TensorProto
 template <size_t alignment>
 common::Status GetSizeInBytesFromTensorProto(const ONNX_NAMESPACE::TensorProto& tensor_proto, size_t* out);
 
+/**
+Special marker used to indicate an existing memory buffer contains the TensorProto external data.
+If the 'location' field of the external data info is set to this marker, the 'offset' field should contain the
+address of the memory containing the data.
+*/
+constexpr const ORTCHAR_T* kTensorProtoMemoryAddressTag = ORT_TSTR("*/_ORT_MEM_ADDR_/*");
+
 // Given a tensor proto with external data obtain a pointer to the data and its length.
 // The ext_data_deleter argument is updated with a callback that owns/releases the data.
 common::Status GetExtDataFromTensorProto(const Env& env, const ORTCHAR_T* model_path,

onnxruntime/core/graph/graph.cc

@@ -714,12 +714,15 @@ flatbuffers::Offset<fbs::NodeEdge> Node::SaveEdgesToOrtFormat(flatbuffers::FlatB
 #endif  // !defined(ORT_MINIMAL_BUILD)
 
 Status Node::LoadFromOrtFormat(const onnxruntime::fbs::Node& fbs_node, Graph& graph,
+                               bool can_use_flatbuffer_for_initializers,
                                const logging::Logger& logger, std::unique_ptr<Node>& node) {
   node = std::make_unique<Node>(fbs_node.index(), graph);
-  return node->LoadFromOrtFormat(fbs_node, logger);
+  return node->LoadFromOrtFormat(fbs_node, can_use_flatbuffer_for_initializers, logger);
 }
 
-Status Node::LoadFromOrtFormat(const onnxruntime::fbs::Node& fbs_node, const logging::Logger& logger) {
+Status Node::LoadFromOrtFormat(const onnxruntime::fbs::Node& fbs_node,
+                               bool can_use_flatbuffer_for_initializers,
+                               const logging::Logger& logger) {
   auto LoadNodeArgsFromOrtFormat =
       [&](const flatbuffers::Vector<flatbuffers::Offset<flatbuffers::String>>* fbs_node_arg_names,
           std::vector<NodeArg*>& node_args,
@@ -758,7 +761,8 @@ Status Node::LoadFromOrtFormat(const onnxruntime::fbs::Node& fbs_node, const log
       AttributeProto attr_proto;
       std::unique_ptr<Graph> subgraph;
       ORT_RETURN_IF_ERROR(
-          fbs::utils::LoadAttributeOrtFormat(*fbs_attr, attr_proto, subgraph, *graph_, *this, logger));
+          fbs::utils::LoadAttributeOrtFormat(*fbs_attr, attr_proto, subgraph, *graph_, *this,
+                                             can_use_flatbuffer_for_initializers, logger));
 
       // If we have a sub graph in this attributes, it will be loaded into subgraph ptr
       // while the attribute proto contains the sub graph will have the empty g() field
@@ -4187,6 +4191,7 @@ Status Graph::LoadFromOrtFormat(const onnxruntime::fbs::Graph& fbs_graph,
 #if !defined(ORT_MINIMAL_BUILD)
                                 IOnnxRuntimeOpSchemaCollectionPtr schema_registry,
 #endif
+                                bool can_use_flatbuffer_for_initializers,
                                 const logging::Logger& logger, std::unique_ptr<Graph>& graph) {
   graph = std::make_unique<Graph>(owning_model, domain_to_version,
 #if !defined(ORT_MINIMAL_BUILD)
@@ -4196,7 +4201,7 @@ Status Graph::LoadFromOrtFormat(const onnxruntime::fbs::Graph& fbs_graph,
                                   // Assume anything in ORT format has already been validated.
                                   false);
 
-  ORT_RETURN_IF_ERROR(graph->LoadFromOrtFormat(fbs_graph));
+  ORT_RETURN_IF_ERROR(graph->LoadFromOrtFormat(fbs_graph, can_use_flatbuffer_for_initializers));
 
 #if !defined(ORT_MINIMAL_BUILD)
   // in a full build we need to run Resolve to fully populate ResolveContext and Node::op_,
@@ -4212,6 +4217,7 @@ Status Graph::LoadFromOrtFormat(const onnxruntime::fbs::Graph& fbs_graph,
 
 Status Graph::LoadFromOrtFormat(const onnxruntime::fbs::Graph& fbs_graph,
                                 Graph& parent_graph, const Node& parent_node,
+                                bool can_use_flatbuffer_for_initializers,
                                 const logging::Logger& logger, std::unique_ptr<Graph>& graph) {
   graph = std::make_unique<Graph>(parent_graph.owning_model_,
                                   parent_graph.domain_to_version_,
@@ -4223,7 +4229,7 @@ Status Graph::LoadFromOrtFormat(const onnxruntime::fbs::Graph& fbs_graph,
                                   // Assume anything in ORT format has already been validated.
                                   false);
 
-  return graph->LoadFromOrtFormat(fbs_graph);
+  return graph->LoadFromOrtFormat(fbs_graph, can_use_flatbuffer_for_initializers);
 }
 
 Graph::Graph(const Model& owning_model,
@@ -4252,7 +4258,8 @@ Graph::Graph(const Model& owning_model,
       is_loaded_from_model_file_(true) {  // true as the Graph isn't manually constructed from scratch
 }
 
-common::Status Graph::LoadFromOrtFormat(const onnxruntime::fbs::Graph& fbs_graph) {
+common::Status Graph::LoadFromOrtFormat(const onnxruntime::fbs::Graph& fbs_graph,
+                                        bool can_use_flatbuffer_for_initializers) {
   // We deserialize the graph from ORT format in the following order:
   // 1. Deserialize the initializers and sparse initializers. Convert sparse to dense.
   // 2. Deserialize the NodeArgs
@@ -4282,7 +4289,8 @@ common::Status Graph::LoadFromOrtFormat(const onnxruntime::fbs::Graph& fbs_graph
     for (const auto* fbs_tensor : *fbs_initializers) {
       ORT_RETURN_IF(nullptr == fbs_tensor, "Initializer tensor is missing. Invalid ORT format model.");
       TensorProto* initializer = deserialized_proto_data_.add_initializer();
-      ORT_RETURN_IF_ERROR(fbs::utils::LoadInitializerOrtFormat(*fbs_tensor, *initializer));
+      ORT_RETURN_IF_ERROR(fbs::utils::LoadInitializerOrtFormat(*fbs_tensor, *initializer,
+                                                               can_use_flatbuffer_for_initializers));
       auto p = name_to_initial_tensor_.emplace(initializer->name(), initializer);
       if (!p.second) {
         LOGS(logger_, WARNING) << "Duplicate initializer (dense or ConstantNode): '" << initializer->name()
@@ -4342,7 +4350,8 @@ common::Status Graph::LoadFromOrtFormat(const onnxruntime::fbs::Graph& fbs_graph
     for (const auto* fbs_node : *fbs_nodes) {
       ORT_RETURN_IF(nullptr == fbs_node, "Node is missing. Invalid ORT format model.");
       std::unique_ptr<Node> node;
-      ORT_RETURN_IF_ERROR(Node::LoadFromOrtFormat(*fbs_node, *this, logger_, node));
+      ORT_RETURN_IF_ERROR(Node::LoadFromOrtFormat(*fbs_node, *this, can_use_flatbuffer_for_initializers, logger_,
+                                                  node));
       ORT_RETURN_IF(node->Index() >= fbs_graph.max_node_index(), "Node index is out of range");
       nodes_[node->Index()] = std::move(node);
       ++num_of_nodes_;

onnxruntime/core/graph/graph_flatbuffers_utils.cc

@@ -5,6 +5,7 @@
 #include "core/flatbuffers/flatbuffers_utils.h"
 #include "core/flatbuffers/schema/ort.fbs.h"
 #include "core/framework/tensorprotoutils.h"
+#include "core/framework/tensor_external_data_info.h"
 #include "graph_flatbuffers_utils.h"
 #include "flatbuffers/flatbuffers.h"
 
@@ -171,8 +172,8 @@ Status SaveAttributeOrtFormat(flatbuffers::FlatBufferBuilder& builder,
 
 #endif
 
-Status LoadInitializerOrtFormat(const fbs::Tensor& fbs_tensor,
-                                TensorProto& initializer) {
+Status LoadInitializerOrtFormat(const fbs::Tensor& fbs_tensor, TensorProto& initializer,
+                                bool can_use_flatbuffer_for_initializers) {
   initializer.Clear();
 
   LOAD_STR_FROM_ORT_FORMAT(initializer, name, fbs_tensor.name());
@@ -196,8 +197,29 @@ Status LoadInitializerOrtFormat(const fbs::Tensor& fbs_tensor,
     const auto* fbs_raw_data = fbs_tensor.raw_data();
     ORT_RETURN_IF(nullptr == fbs_raw_data, "Missing raw data for initializer. Invalid ORT format model.");
 
-    // fbs_raw_data is uint8_t vector, so the size is byte size
-    initializer.set_raw_data(fbs_raw_data->Data(), fbs_raw_data->size());
+    if (can_use_flatbuffer_for_initializers && fbs_raw_data->size() > 127) {
+      initializer.set_data_location(ONNX_NAMESPACE::TensorProto_DataLocation_EXTERNAL);
+
+      static_assert(sizeof(void*) <= sizeof(ExternalDataInfo::OFFSET_TYPE));
+      const void* data_offset = fbs_raw_data->Data();
+      // we reinterpret_cast this back to void* in tensorprotoutils.cc:GetExtDataFromTensorProto.
+      // use intptr_t as OFFSET_TYPE is signed. in theory you could get a weird looking value if the address uses the
+      // high bit, but that should be unlikely in a scenario where we care about memory usage enough to use this path.
+      auto offset = gsl::narrow<ExternalDataInfo::OFFSET_TYPE>(reinterpret_cast<intptr_t>(data_offset));
+
+      ONNX_NAMESPACE::StringStringEntryProto* entry = initializer.mutable_external_data()->Add();
+      entry->set_key("location");
+      entry->set_value(ToUTF8String(onnxruntime::utils::kTensorProtoMemoryAddressTag));
+      entry = initializer.mutable_external_data()->Add();
+      entry->set_key("offset");
+      entry->set_value(std::to_string(offset));
+      entry = initializer.mutable_external_data()->Add();
+      entry->set_key("length");
+      entry->set_value(std::to_string(fbs_raw_data->size()));
+    } else {
+      // fbs_raw_data is uint8_t vector, so the size is byte size
+      initializer.set_raw_data(fbs_raw_data->Data(), fbs_raw_data->size());
+    }
   }
 
   return Status::OK();
@@ -231,6 +253,7 @@ Status LoadAttributeOrtFormat(const fbs::Attribute& fbs_attr,
                               ONNX_NAMESPACE::AttributeProto& attr_proto,
                               std::unique_ptr<onnxruntime::Graph>& sub_graph,
                               onnxruntime::Graph& graph, onnxruntime::Node& node,
+                              bool can_use_flatbuffer_for_initializers,
                               const logging::Logger& logger) {
   attr_proto.Clear();
   LOAD_STR_FROM_ORT_FORMAT(attr_proto, name, fbs_attr.name());
@@ -253,7 +276,8 @@ Status LoadAttributeOrtFormat(const fbs::Attribute& fbs_attr,
     case AttributeProto_AttributeType_TENSOR: {
       auto fbs_tensor = fbs_attr.t();
       ORT_RETURN_IF(nullptr == fbs_tensor, "Null tensor attribute. Invalid ORT format model.");
-      ORT_RETURN_IF_ERROR(LoadInitializerOrtFormat(*fbs_tensor, *attr_proto.mutable_t()));
+      ORT_RETURN_IF_ERROR(LoadInitializerOrtFormat(*fbs_tensor, *attr_proto.mutable_t(),
+                                                   can_use_flatbuffer_for_initializers));
     } break;
     case AttributeProto_AttributeType_GRAPH: {
       // If the attribute type is a graph, we will create an empty graph in attr_proto so that the ONNX checker
@@ -261,7 +285,9 @@ Status LoadAttributeOrtFormat(const fbs::Attribute& fbs_attr,
       auto fbs_graph = fbs_attr.g();
       ORT_RETURN_IF(nullptr == fbs_graph, "Null graph attribute. Invalid ORT format model.");
       attr_proto.mutable_g()->set_name("Empty graph proto from deserialization of ORT format model");
-      ORT_RETURN_IF_ERROR(onnxruntime::Graph::LoadFromOrtFormat(*fbs_graph, graph, node, logger, sub_graph));
+      ORT_RETURN_IF_ERROR(onnxruntime::Graph::LoadFromOrtFormat(*fbs_graph, graph, node,
+                                                                can_use_flatbuffer_for_initializers,
+                                                                logger, sub_graph));
     } break;
     case AttributeProto_AttributeType_FLOATS: {
       auto fbs_floats = fbs_attr.floats();
@@ -294,7 +320,8 @@ Status LoadAttributeOrtFormat(const fbs::Attribute& fbs_attr,
       tensors->Reserve(fbs_tensors->size());
       for (const auto* fbs_tensor : *fbs_tensors) {
         ORT_RETURN_IF(nullptr == fbs_tensor, "Null tensor in tensors attribute. Invalid ORT format model.");
-        ORT_RETURN_IF_ERROR(LoadInitializerOrtFormat(*fbs_tensor, *tensors->Add()));
+        ORT_RETURN_IF_ERROR(LoadInitializerOrtFormat(*fbs_tensor, *tensors->Add(),
+                                                     can_use_flatbuffer_for_initializers));
       }
     } break;
 

onnxruntime/core/graph/graph_flatbuffers_utils.h

@@ -49,8 +49,19 @@ onnxruntime::common::Status SaveAttributeOrtFormat(
     flatbuffers::Offset<fbs::Attribute>& fbs_attr, const Path& model_path,
     const onnxruntime::Graph* subgraph);
 
-onnxruntime::common::Status LoadInitializerOrtFormat(
-    const fbs::Tensor& fbs_tensor, ONNX_NAMESPACE::TensorProto& initializer);
+/// <summary>
+/// Load an initializer from an ORT format flatbuffer.
+/// </summary>
+/// <param name="fbs_tensor">Flatbuffer Tensor</param>
+/// <param name="initializer">TensorProto to load data into</param>
+/// <param name="can_use_flatbuffer_for_initializers">
+/// If true, set the TensorProto to point to the memory in the flatbuffer instead of copying data.
+/// This requires the buffer to remain valid for the entire duration of the InferenceSession.
+/// </param>
+/// <returns>Status</returns>
+onnxruntime::common::Status LoadInitializerOrtFormat(const fbs::Tensor& fbs_tensor,
+                                                     ONNX_NAMESPACE::TensorProto& initializer,
+                                                     bool can_use_flatbuffer_for_initializers = false);
 
 onnxruntime::common::Status LoadSparseInitializerOrtFormat(const fbs::SparseTensor& fbs_sparse_tensor,
                                                            ONNX_NAMESPACE::SparseTensorProto& initializer);
@@ -62,6 +73,7 @@ onnxruntime::common::Status LoadAttributeOrtFormat(const fbs::Attribute& fbs_att
                                                    ONNX_NAMESPACE::AttributeProto& attr_proto,
                                                    std::unique_ptr<onnxruntime::Graph>& sub_graph,
                                                    onnxruntime::Graph& graph, onnxruntime::Node& node,
+                                                   bool can_use_flatbuffer_for_initializers,
                                                    const logging::Logger& logger);
 
 }  // namespace utils

onnxruntime/core/graph/model.cc

@@ -770,6 +770,7 @@ common::Status Model::LoadFromOrtFormat(const fbs::Model& fbs_model,
 #if !defined(ORT_MINIMAL_BUILD)
                                         const IOnnxRuntimeOpSchemaRegistryList* local_registries,
 #endif
+                                        bool can_use_flatbuffer_for_initializers,
                                         const logging::Logger& logger,
                                         std::unique_ptr<Model>& model) {
   model = std::make_unique<Model>();
@@ -827,10 +828,11 @@ common::Status Model::LoadFromOrtFormat(const fbs::Model& fbs_model,
   ORT_RETURN_IF(nullptr == fbs_graph, "Graph is null. Invalid ORT format model.");
 
 #if !defined(ORT_MINIMAL_BUILD)
-  ORT_RETURN_IF_ERROR(Graph::LoadFromOrtFormat(*fbs_graph, *model, domain_to_version, schema_registry, logger,
-                                               model->graph_));
+  ORT_RETURN_IF_ERROR(Graph::LoadFromOrtFormat(*fbs_graph, *model, domain_to_version, schema_registry,
+                                               can_use_flatbuffer_for_initializers, logger, model->graph_));
 #else
-  ORT_RETURN_IF_ERROR(Graph::LoadFromOrtFormat(*fbs_graph, *model, domain_to_version, logger, model->graph_));
+  ORT_RETURN_IF_ERROR(Graph::LoadFromOrtFormat(*fbs_graph, *model, domain_to_version,
+                                               can_use_flatbuffer_for_initializers, logger, model->graph_));
 #endif
   return Status::OK();
 }

onnxruntime/core/graph/model.h

@@ -40,7 +40,8 @@ struct ModelOptions {
   bool strict_shape_type_inference;
 
   ModelOptions(bool allow_released_opsets_only, bool strict_shape_type_inference)
-      : allow_released_opsets_only(allow_released_opsets_only), strict_shape_type_inference(strict_shape_type_inference) {}
+      : allow_released_opsets_only(allow_released_opsets_only),
+        strict_shape_type_inference(strict_shape_type_inference) {}
 
   ModelOptions() : ModelOptions(true, false) {}
 };
@@ -297,6 +298,7 @@ class Model {
 #if !defined(ORT_MINIMAL_BUILD)
                                           const IOnnxRuntimeOpSchemaRegistryList* local_registries,
 #endif
+                                          bool can_use_flatbuffer_for_initializers,
                                           const logging::Logger& logger,
                                           std::unique_ptr<Model>& model);
 

onnxruntime/core/session/inference_session.cc

@@ -972,9 +972,11 @@ Status InferenceSession::LoadOrtModel(const PathString& model_uri) {
 
 Status InferenceSession::LoadOrtModel(const void* model_data, int model_data_len) {
   return LoadOrtModelWithLoader([&]() {
+    const auto& config_options = GetSessionOptions().config_options;
     const auto use_ort_model_bytes_directly =
-        GetSessionOptions().config_options.GetConfigOrDefault(kOrtSessionOptionsConfigUseORTModelBytesDirectly, "0");
-    if (use_ort_model_bytes_directly != "1") {
+        config_options.GetConfigOrDefault(kOrtSessionOptionsConfigUseORTModelBytesDirectly, "0") == "1";
+
+    if (!use_ort_model_bytes_directly) {
       // copy bytes as we need them to be available when InferenceSession::Initialize is called later.
       ort_format_model_bytes_data_holder_.resize(model_data_len);
       std::copy_n(reinterpret_cast<const uint8_t*>(model_data), model_data_len,
@@ -1025,15 +1027,25 @@ Status InferenceSession::LoadOrtModelWithLoader(std::function<Status()> load_ort
   const auto* fbs_model = fbs_session->model();
   ORT_RETURN_IF(nullptr == fbs_model, "Missing Model. Invalid ORT format model.");
 
+  // if we're using the bytes directly because kOrtSessionOptionsConfigUseORTModelBytesDirectly was set and the user
+  // provided an existing buffer of bytes when creating the InferenceSession, ort_format_model_bytes_data_holder_
+  // will be empty.
+  // if that is the case we also allow creating initializers that directly use those bytes.
+  const auto& config_options = session_options_.config_options;
+  using_ort_model_bytes_for_initializers_ =
+      ort_format_model_bytes_data_holder_.empty() &&
+      config_options.GetConfigOrDefault(kOrtSessionOptionsConfigUseORTModelBytesForInitializers, "0") == "1";
+
   // need to go from unique_ptr to shared_ptr when moving into model_
   std::unique_ptr<Model> tmp_model;
 #if !defined(ORT_MINIMAL_BUILD)
   ORT_RETURN_IF_ERROR(Model::LoadFromOrtFormat(*fbs_model,
                                                HasLocalSchema() ? &custom_schema_registries_ : nullptr,
+                                               using_ort_model_bytes_for_initializers_,
                                                *session_logger_, tmp_model));
-
 #else
-  ORT_RETURN_IF_ERROR(Model::LoadFromOrtFormat(*fbs_model, *session_logger_, tmp_model));
+  ORT_RETURN_IF_ERROR(Model::LoadFromOrtFormat(*fbs_model, using_ort_model_bytes_for_initializers_, *session_logger_,
+                                               tmp_model));
 #endif
 
   ORT_RETURN_IF_ERROR(SaveModelMetadata(*tmp_model));
@@ -1264,9 +1276,9 @@ common::Status InferenceSession::Initialize() {
     //
     // NOTE: UpdateProvidersWithSharedAllocators is replace-only and will not insert a new allocator into the EP, so
     // it must be called after RegisterAllocator.
-    std::string use_env_allocators = session_options_.config_options.GetConfigOrDefault(kOrtSessionOptionsConfigUseEnvAllocators,
-                                                                                        "0");
-    if (use_env_allocators == "1") {
+    bool use_env_allocators =
+        session_options_.config_options.GetConfigOrDefault(kOrtSessionOptionsConfigUseEnvAllocators, "0") == "1";
+    if (use_env_allocators) {
       LOGS(*session_logger_, INFO) << "This session will use the allocator registered with the environment.";
       UpdateProvidersWithSharedAllocators();
     }
@@ -1446,10 +1458,10 @@ common::Status InferenceSession::Initialize() {
 
     is_inited_ = true;
 
-    // we don't directly use the ORT format bytes currently, so free those now
-    // TODO, we may need to keep the bytes if we are using the offset directly in the initializers
-    ort_format_model_bytes_ = gsl::span<const uint8_t>();
-    std::vector<uint8_t>().swap(ort_format_model_bytes_data_holder_);
+    if (!using_ort_model_bytes_for_initializers_) {
+      ort_format_model_bytes_ = gsl::span<const uint8_t>();
+      std::vector<uint8_t>().swap(ort_format_model_bytes_data_holder_);
+    }
 
     // and log telemetry
     bool model_has_fp16_inputs = ModelHasFP16Inputs(graph);

onnxruntime/core/session/inference_session.h

@@ -791,6 +791,8 @@ class InferenceSession {
   // "session.use_ort_model_bytes_directly" to "1", this will be empty
   std::vector<uint8_t> ort_format_model_bytes_data_holder_;
 
+  bool using_ort_model_bytes_for_initializers_{false};
+
   // Container to store pre-packed weights to share between sessions.
   // The life-cycle of the cache itself is maintained by the user and the user will ensure
   // the cache is valid until any session reliant on it is still in scope.

onnxruntime/test/framework/ort_model_only_test.cc

@@ -36,6 +36,7 @@ struct OrtModelTestInfo {
   std::vector<std::pair<std::string, std::string>> configs;
   bool run_use_buffer{false};
   bool disable_copy_ort_buffer{false};
+  bool use_buffer_for_initializers{false};
 };
 
 static void RunOrtModel(const OrtModelTestInfo& test_info) {
@@ -47,6 +48,10 @@ static void RunOrtModel(const OrtModelTestInfo& test_info) {
 
   if (test_info.disable_copy_ort_buffer) {
     ASSERT_STATUS_OK(so.config_options.AddConfigEntry(kOrtSessionOptionsConfigUseORTModelBytesDirectly, "1"));
+
+    if (test_info.use_buffer_for_initializers) {
+      ASSERT_STATUS_OK(so.config_options.AddConfigEntry(kOrtSessionOptionsConfigUseORTModelBytesForInitializers, "1"));
+    }
   }
 
   std::vector<char> model_data;
@@ -460,6 +465,15 @@ TEST(OrtModelOnlyTests, LoadOrtFormatModelFromBufferNoCopy) {
   RunOrtModel(test_info);
 }
 
+// Load the model from a buffer instead of a file path, and not copy the buffer in session creation
+TEST(OrtModelOnlyTests, LoadOrtFormatModelFromBufferNoCopyInitializersUseBuffer) {
+  OrtModelTestInfo test_info = GetTestInfoForLoadOrtFormatModel();
+  test_info.run_use_buffer = true;
+  test_info.disable_copy_ort_buffer = true;
+  test_info.use_buffer_for_initializers = true;
+  RunOrtModel(test_info);
+}
+
 #if !defined(DISABLE_ML_OPS)
 // test that we can deserialize and run a previously saved ORT format model
 // for a model with sequence and map outputs
@@ -543,6 +557,5 @@ TEST(OrtModelOnlyTests, TestBackwardsCompat) {
 }
 
 #endif  // !defined(DISABLE_ML_OPS)
-
 }  // namespace test
 }  // namespace onnxruntime