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

Commit
88 days ago
Use user-supplied external initializer in place when already on the planned device (#29013) ### 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: ```cpp 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.
Author
Parents
Loading