Skip to content

Save much memory at model loading time by converting weights to OrtValues early - #26345

Merged
yuslepukhin merged 8 commits into
mainfrom
yuslepukhin/convert_weights_early
Oct 30, 2025
Merged

Save much memory at model loading time by converting weights to OrtValues early#26345
yuslepukhin merged 8 commits into
mainfrom
yuslepukhin/convert_weights_early

Conversation

@yuslepukhin

@yuslepukhin yuslepukhin commented Oct 17, 2025

Copy link
Copy Markdown
Contributor

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
Before change DEBUG 2025-10-16 144819

After

After change DEBUG 2025-10-16 144819

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 ConstantFolding where 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.

image

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.
Sanpshot 3 Peak on ConstantFolding Transpose Optimizer

Snapshot 4 Peak AddInitializer from ConstantFolding image

Motivation and Context

Reduce memory usage.

This reverts commit 3ca49d8.

It also makes adjustments for the current source code state.

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 AddInitializerWithExternalData instead of AddInitializer
  • 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.

@yuslepukhin
yuslepukhin marked this pull request as ready for review October 17, 2025 21:43
Comment thread onnxruntime/core/framework/session_state_utils.cc
Comment thread onnxruntime/core/graph/graph.cc Outdated
Comment thread onnxruntime/core/optimizer/attention_fusion.cc Outdated
Comment thread onnxruntime/test/framework/session_state_test.cc Outdated
Comment thread onnxruntime/test/providers/tensorrt/tensorrt_basic_test.cc
Comment thread onnxruntime/core/graph/graph_utils.cc
Comment thread onnxruntime/core/graph/graph.cc
Comment thread onnxruntime/core/graph/graph.cc Outdated
Comment thread onnxruntime/core/graph/graph.cc Outdated
Comment thread onnxruntime/core/graph/graph.cc
@yuslepukhin
yuslepukhin merged commit b860988 into main Oct 30, 2025
92 checks passed
@yuslepukhin
yuslepukhin deleted the yuslepukhin/convert_weights_early branch October 30, 2025 17:44
titaiwangms added a commit that referenced this pull request Jun 23, 2026
…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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants