Skip to content

Use user-supplied external initializer in place when already on the planned device - #29013

Merged
Tianlei Wu (tianleiwu) merged 1 commit into
microsoft:mainfrom
tianleiwu:tlwu/fix_29009_shared_device_init
Jun 12, 2026
Merged

Tianlei Wu (tianleiwu) merged 1 commit into
microsoft:mainfrom
tianleiwu:tlwu/fix_29009_shared_device_init

Conversation

@tianleiwu

Copy link
Copy Markdown
Contributor

Description

SaveInitializedTensors only used a user-supplied initializer OrtValue (from
AddExternalInitializers) in place when the initializer was planned on CPU. For any non-CPU
(e.g. CUDA) initializer it always allocated a fresh device tensor and copied into it, even when
the supplied OrtValue already lived on the planned device.

This change uses the supplied OrtValue directly when its tensor's device matches the planned
device, mirroring the existing CPU case and the AddInitializer (initializers_to_share_map)
no-copy path:

const auto& graph_value_device = ort_value_from_graph.Get<Tensor>().Location().device;
if (memory_info.device == default_cpu_device || graph_value_device == memory_info.device) {
  // Planned on CPU, or the supplied initializer already lives on the planned device:
  // use it in place (no per-session allocation/copy; enables cross-session sharing).
  ort_value = std::move(ort_value_from_graph);
} else {
  // existing allocate-on-device + CopyTensorFromCPUToDevice fallback (true cross-device case)
}

Motivation

Two benefits:

  1. Avoids a redundant per-session device allocation + device copy for every externally-supplied
    initializer that is already on the target device.
  2. Enables cross-session device-memory sharing. Supplying the same device OrtValue to
    multiple sessions (e.g. a large token embedding + lm_head shared between a main decoder and an
    auxiliary speculative-decoding / multi-token-prediction head) now keeps a single device buffer
    instead of one copy per session. For a large-vocab model this saves ~2 GB of VRAM.

This brings AddExternalInitializers in line with AddInitializer, which already uses the supplied
OrtValue in place when its device matches the planned device.

Fixes #29009.

Behavior / compatibility

  • The CPU path is unchanged (memory_info.device == default_cpu_device still short-circuits first).
  • The true cross-device case (supplied tensor on a different device than planned) still falls back to
    allocate + CopyTensorFromCPUToDevice, so existing behavior is preserved there.
  • No public API change.

Testing

  • Existing TestExternalInitializersInjection (CPU) continues to pass (CPU path untouched).
  • Validated end-to-end with ONNX Runtime GenAI on CUDA: sharing a fp16 embedding + lm_head
    (1017 MB each) between two sessions that load separate graphs drops the second model's device
    footprint by ~2145 MB (≈2 GB), with identical inference output vs the non-shared baseline.

Note: a device-level unit test that asserts the shared buffer is reused (no copy) needs internal
session-state access plus a GPU EP harness; happy to add one under test/providers/cuda if
reviewers prefer.

…lanned device

SaveInitializedTensors only used a graph initializer OrtValue directly when the
planned device was CPU; for any non-CPU (e.g. CUDA) initializer it always
allocated a fresh device tensor and copied into it, even when the user-supplied
OrtValue (from AddExternalInitializers) already lived on the planned device.

Use the supplied OrtValue in place when its tensor device matches the planned
device, mirroring the CPU case and the AddInitializer (initializers_to_share_map)
path. This avoids a redundant per-session device allocation + copy and lets the
same device buffer be shared across sessions that supply the same OrtValue (e.g.
an embedding / lm_head shared between a main model and a speculative/MTP head).

Fixes microsoft#29009.

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 updates session initializer materialization so that externally supplied initializer OrtValues (via AddExternalInitializers) are reused in-place when they already live on the device that ORT planned for that initializer (e.g., CUDA), avoiding a redundant allocation+copy and enabling cross-session device-buffer sharing.

Changes:

  • In SaveInitializedTensors, detect the device of the graph-provided initializer OrtValue and reuse it directly when it matches the planned device.
  • Preserve the existing fallback behavior for true cross-device cases by allocating on the planned device and copying via the DataTransferManager.

Comment thread onnxruntime/core/framework/session_state_utils.cc

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.

Need unit test

@tianleiwu
Tianlei Wu (tianleiwu) enabled auto-merge (squash) June 12, 2026 17:24
@tianleiwu
Tianlei Wu (tianleiwu) merged commit be41f5c into microsoft:main Jun 12, 2026
89 of 91 checks passed
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.

AddExternalInitializers copies device (GPU) OrtValues per session instead of using them in place

3 participants