Use user-supplied external initializer in place when already on the planned device - #29013
Merged
Tianlei Wu (tianleiwu) merged 1 commit intoJun 12, 2026
Conversation
…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.
Tianlei Wu (tianleiwu)
requested review from
Copilot and
Dmitri Smirnov (yuslepukhin)
and removed request for
Copilot
June 11, 2026 20:28
Copilot started reviewing on behalf of
Dmitri Smirnov (yuslepukhin)
June 11, 2026 20:48
View session
Contributor
There was a problem hiding this comment.
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 initializerOrtValueand 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.
Dmitri Smirnov (yuslepukhin)
approved these changes
Jun 11, 2026
Dmitri Smirnov (yuslepukhin)
left a comment
Contributor
There was a problem hiding this comment.
Need unit test
Tianlei Wu (tianleiwu)
enabled auto-merge (squash)
June 12, 2026 17:24
Tianlei Wu (tianleiwu)
merged commit Jun 12, 2026
be41f5c
into
microsoft:main
89 of 91 checks passed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
SaveInitializedTensorsonly used a user-supplied initializerOrtValue(fromAddExternalInitializers) 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
OrtValuealready lived on the planned device.This change uses the supplied
OrtValuedirectly when its tensor's device matches the planneddevice, mirroring the existing CPU case and the
AddInitializer(initializers_to_share_map)no-copy path:
Motivation
Two benefits:
initializer that is already on the target device.
OrtValuetomultiple sessions (e.g. a large token embedding +
lm_headshared between a main decoder and anauxiliary 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
AddExternalInitializersin line withAddInitializer, which already uses the suppliedOrtValuein place when its device matches the planned device.Fixes #29009.
Behavior / compatibility
memory_info.device == default_cpu_devicestill short-circuits first).allocate +
CopyTensorFromCPUToDevice, so existing behavior is preserved there.Testing
TestExternalInitializersInjection(CPU) continues to pass (CPU path untouched).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.