onnxruntime

ad2d14b5 - [CUDA] Accelerate external weight loading with pipelined pinned staging and plugin EP support (#33839)

[CUDA] Accelerate external weight loading with pipelined pinned staging and plugin EP support (#33839) ### Description <!-- Describe your changes. --> This PR accelerates external initializer loading to GPU memory for both the built-in CUDA Execution Provider and the CUDA Plugin Execution Provider without introducing external storage dependencies (`libcufile` / GPUDirect Storage or DirectStorage): 1. **Lock-Free Multi-Worker Double-Buffered Pinned Staging & Async H2D Pipeline**: - Replaces the per-chunk global thread-pool barrier and per-tensor `cudaStreamSynchronize` barrier with a shared lock-free work queue (`ExecuteExternalDataLoadTasks` in `onnxruntime/core/framework/external_data_loader_thread_pool.h`). - Each reader worker owns a 4 KiB-aligned slice in both double-buffered pinned staging buffers and a pair of non-blocking CUDA streams, pulling chunks (`2 MiB` default) via atomic index and overlapping multi-threaded file reads (`pread` on POSIX / `OVERLAPPED` `ReadFile` on Windows) with `cudaMemcpyAsync` across tensor boundaries. 2. **CUDA Plugin EP (`PluginExecutionProvider`) Support**: - Implements `PluginExecutionProvider::GetExternalDataLoader()` using `host_accessible_memory_info` (`Reserve`/`Free` for pinned memory), lightweight copy-only `OrtSyncStreamImpl` streams (`copy_only=1`, skipping unnecessary cuBLAS/cuDNN handle creation), and `OrtDataTransferImpl` (`CopyTensorAsync`). - Ensures `CudaDataTransfer::CopyTensorsImpl` and `CudaSyncStream::FlushImpl` select and restore the target CUDA device when invoked from reader worker threads. 3. **Cross-Tensor Batch Lifecycle, Offset Sorting, and File-Handle Caching**: - Extends `IExternalDataLoader` and `ExternalDataLoaderManager` with `BeginLoad()`, `FinalizeLoad()`, and `EndLoad()` so external tensors in `SaveInitializedTensors` are pipelined across tensor boundaries and synchronized once in `FinalizeLoad()`. - Sorts initializers in `SaveInitializedTensors` by `(external_location, external_offset, ort_value_index)` for sequential file access instead of hash-bucket order, and caches validated `RandomAccessFile` handles per `SaveInitializedTensors` call. 4. **Conservative Defaults, Model-Size Gating, Session Options, and Minimal Build Exclusion**: - Computes the default reader thread count from physical CPU cores (`clamp(physical_cores / 2, 2, 8)`). - Uses a conservative **32 MiB** default total pinned staging allocation (two `16 MiB` buffers, 4x smaller than the previous 128 MiB default). - Defers pinned buffer, stream, and thread-pool allocation until `FinalizeLoad()` and gates it on a **16 MiB** minimum total external initializer size on the device (`session.external_data_loader.min_size_bytes`), so small models incur zero pinned-memory or stream-creation overhead. - Adds session configuration keys (`session.external_data_loader.threads`, `session.external_data_loader.pinned_memory_bytes`, `session.external_data_loader.chunk_bytes`, `session.external_data_loader.min_size_bytes`) alongside matching CUDA EP provider options (`external_data_loader_reading_threads`, `external_data_loader_pinned_memory_bytes`, `external_data_loader_chunk_bytes`). - Excludes the feature from minimal builds (`!defined(ORT_MINIMAL_BUILD)` and `!defined(USE_CUDA_MINIMAL)`). #### Summary of Changes | File | Change | |---|---| | `include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h` | Add `session.external_data_loader.{threads,pinned_memory_bytes,chunk_bytes,min_size_bytes}` keys | | `include/onnxruntime/core/providers/cuda/cuda_provider_options.h` | Add `external_data_loader_{pinned_memory_bytes,chunk_bytes}` and conservative defaults to `OrtCUDAProviderOptionsV2` | | `onnxruntime/core/framework/external_data_loader.h` | Add `ExternalDataLoaderOptions`, validation/parsing helpers, and `Configure`/`BeginLoad`/`FinalizeLoad`/`EndLoad` hooks | | `onnxruntime/core/framework/external_data_loader_thread_pool.h` | Shared portable `std::thread` worker pool and lock-free double-buffered `ExecuteExternalDataLoadTasks` pipeline | | `onnxruntime/core/framework/external_data_loader_manager.{h,cc}` | Add `BeginLoad`, `FinalizeLoad`, and `EndLoad` batch lifecycle management | | `onnxruntime/core/framework/session_state_utils.cc` | Sort external initializers by `(location, offset)`, cache validated file handles, and bracket loading with `BeginLoad`/`FinalizeLoad`/`EndLoad` | | `onnxruntime/core/framework/tensorprotoutils.{h,cc}` | Add `ExternalDataFileCache` to validate and open each external `.data` file once per `SaveInitializedTensors` pass | | `onnxruntime/core/providers/cuda/cuda_external_data_loader.{h,cc}` | Use shared lock-free pipeline, support configurable threads/pinned memory/chunk size/min model size, and guard minimal builds | | `onnxruntime/core/providers/cuda/cuda_execution_provider_info.{h,cc}`, `cuda_provider_factory.cc`, `cuda_execution_provider.cc`, `provider_bridge_ort.cc` | Parse, validate, and forward new external data loader options | | `onnxruntime/core/session/plugin_ep/ep_plugin_provider_interfaces.{h,cc}`, `onnxruntime/core/providers/cuda/plugin/*` | Implement `PluginExternalDataLoader`, copy-only sync streams, device-safe worker H2D copies, and option validation | | `docs/Model_Loading_Performance.md` | Document new session options, provider options, defaults, and model-size threshold | | `onnxruntime/test/framework/external_data_loader_test.cc`, `cuda_external_data_loader_test.cc`, `test_cuda_plugin_ep.py` | Add C++ and Python unit tests for session options, provider options, min-size gating, and invalid-option rejection | ### Motivation and Context <!-- - Why is this change required? What problem does it solve? - If it fixes an open issue, please link to the issue here. --> Loading large ONNX models with external weights (e.g., `qwen_flash_next_int4_int8_rc5_contrib_dflash2/model.onnx`, 68.46 GiB across 899 external tensors in 7 `.data` files) into GPU memory during `InferenceSession` creation had several bottlenecks: 1. `PluginExecutionProvider` (used when building the CUDA EP as a plugin) did not implement `GetExternalDataLoader()`, falling back to per-tensor `mmap` + single-threaded pageable `cudaMemcpy`. 2. In the existing pinned-buffer loader (#32437 / #32713 fallback when `nvidia-fs` is unavailable), a global thread-pool barrier per 64 MiB chunk and a `cudaStreamSynchronize` barrier at the end of every `LoadTensor()` call prevented CPU/GPU overlap across tensors `<= 64 MiB` (705 of 899 tensors), forced all 613 tensors `< 16 MiB` onto a single thread, iterated over initializers in arbitrary hash-bucket order, and reopened/revalidated `.data` files for every tensor. #### Benchmark Comparison (`qwen_flash_next_int4_int8_rc5_contrib_dflash2/model.onnx`, 68.46 GiB, 899 External Tensors, NVIDIA H200) | Configuration | Threads | Pinned Memory | End-to-End `InferenceSession` Time | Effective Throughput | Speedup vs Baseline | |---|---:|---:|---:|---:|---:| | **Baseline (`origin/main` CUDA Plugin EP)** | 1 (pageable) | 0 MiB | 14.481 s | 4.73 GiB/s | 1.00x | | **PR #32713 (`use_gds=1` -> 4-thread barrier fallback)** | 4 | 128 MiB | 5.110 s | 13.40 GiB/s | 2.83x | | **This PR (same 4 threads, 4x less pinned memory)** | **4** | **32 MiB** | **3.730 s** | **18.35 GiB/s** | **3.88x** | | **This PR (same 4 threads, same 128 MiB pinned memory)** | **4** | **128 MiB** | **3.657 s** | **18.72 GiB/s** | **3.96x** | | **This PR (Default: 8 threads from CPU cores, conservative 32 MiB pinned)** | **8** | **32 MiB** | **2.191 s** | **31.25 GiB/s** | **6.61x** | | **This PR (Opt-in: 16 threads, 64 MiB pinned)** | **16** | **64 MiB** | **1.967 s** | **34.81 GiB/s** | **7.36x** | Even at the **exact same 4-thread count and 4x smaller pinned memory (32 MiB vs 128 MiB)**, the lock-free cross-tensor pipeline + offset sorting + file-handle caching reduces session creation time from **5.110 s to 3.730 s (1.37x faster)**. With the CPU-core-derived default (`8` threads) and conservative `32 MiB` pinned memory, session creation completes in **2.191 s (31.25 GiB/s)**. #### Validation - `./onnxruntime_test_all --gtest_filter="*ExternalDataLoader*:*Initializer*"` — **127 tests ran, 126 passed, 1 skipped (32-bit platform check), 0 failed**. - `ORT_TEST_CUDA_PLUGIN_EP=1 python onnxruntime/test/python/transformers/test_cuda_plugin_ep.py` — **94 tests ran, 92 passed, 2 skipped, 0 failed** (including `test_external_data_loader_modes_and_validation`). - `lintrunner -a` — **passed (0 lint issues)**. --------- Co-authored-by: Xavier Dupré <xadupre@users.noreply.github.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>