microsoft / microsoft/onnxruntime

[Feature Request] Streaming model I/O for large ONNX models

Open
#32,224 2 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

ep:WebNN feature request
Dominant language
C++
Stars
21.9k
Forks
4.2k
Avg merge
4d 8h
Merged PRs (30d)
179

Description

Describe the feature request

Motivation

The in-memory model I/O APIs like #29686, #31347 are useful, especially for applications that need to provide models from custom storage, encrypted containers, or other non-file-backed sources.

For large models, however, requiring the complete serialized model to be materialized in memory can become a significant limitation. This is particularly noticeable for multi-gigabyte models, where the temporary serialized representation may consume a substantial fraction of the memory available to the application.

It could therefore be useful to provide an additional streaming reader/writer abstraction alongside (or in place of) the existing in-memory APIs. This could cater all IO required by ORT: pure ONNX, external initialziers, EP context, EP context payloads.

Status quo

Area File-backed Memory-backed
ONNX model input CreateSession CreateSessionFromArray
Compiled model input ModelCompilationOptions_SetInputModelPath ModelCompilationOptions_SetInputModelFromBuffer
Compiled model output ModelCompilationOptions_SetOutputModelPath ModelCompilationOptions_SetOutputModelBuffer
ONNX external initializers Normal ONNX external-data files AddExternalInitializers / AddExternalInitializersFromFilesInMemory
EP Context, embed_mode=0 CreateSession CreateSessionFromArray + SetOutputModelWriteFunc

There is a lot of fragmentation on how one can interface with ORT for file IO. Looking at a regular compile path below we would have a full duplication of weights if one requires in-memory communication with ORT due to the original weights and a compiled set of weights being live at the same time. With a memory mapped external initializer file this looks better as we enable random access into files without loading the full file.

original.onnx
   |
   +--- external weights
   |
   ↓ compile
compiled_ep_context.onnx
   |
   +--- EP context binary
          |
          +--- potentially compiled weights (new external initializers or  baked into payload bytestream)

Considerations

For a stream reader and writer design 2 key things should be considered in my opinion:

  • block-based encryption and decryption without materializing the complete decrypted model (EP provided encryption + custom encryption)
  • direct file handle access by EPs to enable more efficient disk access leveraging e.g. DirectStorage or cuFile

The last point may become increasingly relevant for large models as that can significantly reduce setup time. A sufficiently flexible streaming/range-based model I/O abstraction could make it possible for execution providers or applications to take advantage of these kinds of I/O paths rather than requiring the complete serialized model to first exist as a host-memory buffer.

Describe scenario use case

For in memory communication with ORT for e.g. WebNN and ISVs with encrypted models the current API surface requires 2x the model size in memory when compiling a model. This limits the models that can be run on based on memory requirements.
The memory implications are also somewhat different depending on the system architecture:

  • On systems with discrete GPUs, staging model data in host memory before transferring it to the GPU is relatively common, and system memory capacity is often larger than GPU memory capacity.
  • On unified-memory systems, CPU and GPU workloads compete for the same physical memory pool. These systems are increasingly characterized by the maximum model size they can accommodate, so temporarily keeping both a serialized model representation and the instantiated model weights can significantly reduce the maximum practical model size.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by comparing the existing CreateSession, CreateSessionFromArray, ModelCompilationOptions, AddExternalInitializers, and SetOutputModelWriteFunc entry points, along with issues #29686 and #31347. Define the scope and design for a streaming or range-based abstraction that covers the listed model I/O paths, encryption needs, and execution-provider access before implementation; done requires an agreed API and validation plan.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp
Domain
api, backend-api-design, machine-learning
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Needs clarification
Newbie friendliness
30/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.