Save much memory at model loading time by converting weights to OrtValues early - #26345
Conversation
This reverts commit 3ca49d8. It also makes adjustments for the current source code state.
There was a problem hiding this comment.
Pull Request Overview
This PR saves significant memory during model loading by converting weight initializers to OrtValues early in the graph construction process, rather than later during graph transformation. The changes revert previous logic that deferred this conversion and implements early weight conversion at graph initialization time. The PR demonstrates dramatic memory savings during optimization phases (from 8.1GB to 3.5GB for Phi4 Instruct model) by enabling reuse of OrtValues during constant folding operations.
Key changes:
- Early conversion of large initializers to OrtValues during graph construction
- Update of all graph transformation code to use
AddInitializerWithExternalDatainstead ofAddInitializer - Removal of deferred initializer conversion logic from session inference flow
Reviewed Changes
Copilot reviewed 35 out of 35 changed files in this pull request and generated no comments.
Show a summary per file
| File | Description |
|---|---|
| onnxruntime/core/graph/graph.cc | Implements early initializer-to-OrtValue conversion during graph construction |
| include/onnxruntime/core/graph/graph.h | Removes ConvertInitializersIntoOrtValues method declaration |
| onnxruntime/core/session/inference_session.cc | Removes deferred initializer conversion call from transform pipeline |
| onnxruntime/core/optimizer/*.cc | Updates optimizer classes to use AddInitializerWithExternalData for new initializers |
| orttraining/orttraining/core/optimizer/*.cc | Updates training optimizer classes to use AddInitializerWithExternalData |
| onnxruntime/test/ir/graph_test.cc | Removes test code that called the now-removed conversion method |
| onnxruntime/test/framework/cuda/fence_cuda_test.cc | Updates test utility to use AddInitializerWithExternalData |
Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.
…x output in data propagation (#29084) ## Summary A spec-valid `Shape → Gather(1-D index [-1]) → TopK` model fails to load since ORT 1.25.0 with: ``` K input must be a one-dimensional tensor of size 1. ``` The model is valid: a rank-1 (single-element) Gather index produces a rank-1 Gather output, so the value feeding TopK's `K` input is a 1-D size-1 tensor — exactly what TopK requires. The failure was an **ORT rank-preservation bug in shape-inference data propagation**, not a problem with the model. **Root cause.** `GatherOpDataPropagation::infer()` routed by element **count** rather than index **rank**: it guarded on `indices.size() == 1`, which is true for *both* a 0-D scalar index and a 1-D single-element index, and then unconditionally called `SetInferredShapeScalarValue()`. That dropped the rank of the spec-valid 1-D size-1 case, so `Graph::getInputData()` emitted a 0-D (dimensionless) propagated value. ONNX TopK shape inference then correctly rejected the 0-D `K`. This path was introduced by #26269 (partial data propagation to enhance shape inference). This reproduces even at `GraphOptimizationLevel.ORT_DISABLE_ALL`, where constant folding never runs — confirming the cause is data propagation in shape inference, **not** constant folding (#26345 was an earlier mis-attribution; see the corrected analysis). Fixes the regression reported in #29072. Corrected root-cause analysis: #29072 (comment) ## The fix - **Gather — rank-based routing.** Distinguish the index rank instead of its element count. A genuine 0-D scalar index still stores a scalar value; a rank-1 single-element index now stores a **rank-1** value, so `getInputData()` emits a `TensorProto` with `dims=[1]` and downstream TopK sees a valid 1-D size-1 `K`. The index rank is taken from the same constant initializer the index value comes from (via `get_initialized_input_values` now reporting the initializer rank), rather than a second, independently-resolved `NodeArg` shape — removing a potential source-of-truth drift (EDGE #2). - **Rank-tolerant elementwise companion (Add/Sub/Mul/Div).** These ops were scalar-only and would silently stop propagating once an operand became a rank-1 value (e.g. a `Shape → Gather(1-D idx) → Mul → TopK` chain), because the custom-propagation result replaces ONNX's rank-correct fallback. They now accept a single element carried as either a rank-0 scalar or a rank-1 `[1]` value and keep the output rank consistent with ONNX broadcasting (rank-1 if any operand is rank-1, else scalar), so such chains keep propagating end-to-end. Div additionally guards against division by zero. - **Shared helper (`data_propagation_value_utils.h`).** Centralizes reading/writing a single-element shape value while preserving its rank, used by both the Gather producer and the elementwise consumers so they cannot disagree on rank. The reader **declines** a rank-1 multi-element value (it must never collapse to `element[0]`), so a multi-element value can never be mistaken for a single one. ## Testing Five `ShapeInferenceV2Test` cases (with fixtures + generators), all loading the model at **every** optimization level (including `ORT_DISABLE_ALL`): - `GatherToTopKRankPreservationTest` — the core `Shape → Gather([-1]) → TopK` regression; asserts the rank-1 `K` is preserved. - `GatherMulToTopKRankPreservationTest` — the `… → Gather(1-D idx) → Mul → TopK` chain; asserts propagation survives the elementwise op. - `SinglePropagatedShapeValueGuardTest` — a direct unit test pinning the shared reader's behavior on each channel (scalar, rank-1 single-element, rank-1 multi-element, symbolic, empty). **Mutation-proven**: relaxing the `dim_size()==1` guard makes this test fail, restoring it makes it pass — so the guard the whole fix hinges on is test-locked. - `ShapeMulMultiElementNoScalarCollapseTest` — end-to-end check that a multi-element `Shape → Mul → ConstantOfShape` chain still resolves to its full rank-2 shape (no bogus scalar collapse). - `PartialDataPropagationTest` — pre-existing scalar-index coverage, unchanged. Full `onnxruntime_test_all` suite passes (0 failures) on top of the current `main` (opset-27 / ONNX 1.22.0 integration). The constant-folding memory path (#26345) is untouched — the diff is confined to `data_propagation/`, a small `graph.cc` change, and tests. ## Follow-ups (intentionally out of scope for this PR) - Hardening for a rank ≥ 2 single-element index (e.g. shape `[1,1]`) to *decline* rather than route as rank-1 — needs its own discriminating unit test; pathological/non-exporter, worst case is degraded inference rather than a crash. - Explicit end-to-end coverage for Add/Sub/Div rank-1 chains (the shared-reader unit test already covers the read path for all four ops; only Mul is currently exercised end-to-end). - Minor readability nits. ## DCO Commit is DCO signed-off. --------- Signed-off-by: titaiwangms <titaiwang@microsoft.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Description
Converts weights early and revert "Properly remove in-memory references (#25652)"
This reverts commit 3ca49d8 and makes appropriate adjustments for the current state of the code.
This PR is made possible and on the heels of:
#26263
#25833.
Previous history:
#23979
#25320
#25626
#25652
The first change (#26263) allows us to convert initializers to OrtValues early and save lots of memory at model loading time.
Specifically, for Phi-4-mini-instruct-INT4 model before and after looks like this:
Before

After
The two peaks represent memory usage at optimization time (8.1Gb before) and after weights memory mapping (6.5Gb)
After this change corresponding numbers look 3.5Gb and 4.7Gb respectively.
Most of the savings during optimization phase come from
ConstantFoldingwhere we are able to reuse the resulting OrtValues directly for the new initializers.This PR concludes a series of PRs converting initializers to OrtValues.
Memory consumption before the conversion began was 9.3Gb and 6.7Gb respectively. We are saving almost 6Gb during optimization and 2Gb for the steady state.
The model also loads about 12 seconds faster.
Example of ConstantFolding being one of the top contributors where we duplicate memory for higher peak before Resolve takes care of no longer used initializers.

Motivation and Context
Reduce memory usage.