Repository navigation
Conversation
The Run:ai streamer reads every tensor into one reusable pageable host buffer, and the weight loaders then copy out of it with param.copy_(), i.e. a driver-staged pageable memcpy per tensor. On an 8xB300 loading a 433 GiB NVFP4 checkpoint (232k tensors) as four TP=2 replicas that copy path ran at about 2.4 GB/s per worker and was the longest span of a tuned cold start. Add a `pinned_staging` key to the loader's extra config. When set, the streamer's staging buffer is page-locked right after stream_files() (pre-faulted with MADV_POPULATE_WRITE from several threads, then cudaHostRegister), the streamer copies each tensor to the current device out of that pinned buffer, and the loaders only ever see device tensors so every TP narrow is a device-to-device copy. The buffer is bounded to 16 GB unless `memory_limit` is set, is unregistered when the iterator is exhausted, and any failure to pin falls back to the stock path with a warning. Nothing changes when the key is unset. Measured on vLLM 0.20.2 (GLM-5.2-NVFP4, runai-model-streamer 0.16.1, 4x TP=2 on 8xB300): "Model loading took" 39-42 s vs 75-79 s (-48 %), four replicas serving at 133 s instead of 170 s from script start, every GPU parameter bit-identical to the stock loader, throughput unchanged. At TP=8 the gain is small (every rank reads the whole checkpoint through the page cache, so the ceiling is a memcpy). Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Tom <tom@getimpala.ai>
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment Once the PR is approved or has the If you have any questions, please reach out to us on Slack at https://slack.vllm.ai. Agent GuidelinesIMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban. 🚀 |
|
Documentation preview: https://vllm--61303.org.readthedocs.build/en/61303/ |
|
This pull request has merge conflicts that must be resolved before it can be |
This branch has not been deployed
1 participant
Overview
Adds a
pinned_stagingkey to the Run:ai streamer loader's--model-loader-extra-config. When set, the streamer's CPU staging buffer is page-locked once per load and the streamer copies every tensor to the device out of it, so the weight loaders only see device tensors and the host-to-device traffic is DMA out of pinned memory rather than pageable memcpy. Off by default; nothing changes when the key is unset.Claims
main, see Validation): GLM-5.2-NVFP4 (433 GiB, 232k tensors) as four TP=2 replicas on an 8xB300,"Model loading took"75–79 s → 39–42 s (−48 %), four replicas serving at 133 s instead of 170 s from script start.process_weights_after_loadingon both paths), and throughput is unchanged.cudaHostRegistererror, streamer internals changed) falls back to the existing path with a warning.Validation
tests/model_executor/model_loader/runai_streamer_loader/: new tests for the extra-config validation, the bounded defaultRUNAI_STREAMER_MEMORY_LIMIT, and the CPU fallback of the iterator withpinned_staging=True; the existing buffer-reuse and loader tests still pass (13 passed locally on CPU,--noconftest).pre-commitruff-check,ruff-format,typospass on the changed files.main. The numbers above come from the same change on 0.20.2 (patch on the loader invllm/model_executor/model_loader/weight_utils.pywithrunai-model-streamer0.16.1). Since thenmainstarted cloning tensors out of the streamer (Fix RunAI streamer tensor buffer reuse during weight loading #43464); on the pinned path the streamer's own.to(device)is that copy, so the clone is skipped there (and kept for the distributed streamer, whose device buffer is reused). I will re-measure onmainas soon as I have the node again and post the numbers here; until then this stays a draft.Repro of the measurement:
vllm serve nvidia/GLM-5.2-NVFP4 -tp 2 --load-format runai_streamer --model-loader-extra-config '{"pinned_staging": true}'vs. without the key, same node, same page-cache state; compare theModel loading tooklog line.Details
The streamer reads each tensor into one reusable pageable host buffer (
requests_iterator.buffer), andparam.copy_()out of it runs as a driver-staged pageable memcpy per tensor: about 2.4 GB/s per worker with eight workers on this node, 4.5 GB/s alone. With the buffer registered (cudaHostRegister) the copies run at about 15 GB/s per worker.Registration of a large buffer is single-threaded and page-fault-bound, so the buffer is first faulted in with
MADV_POPULATE_WRITEfrom 16 threads (and advisedMADV_HUGEPAGE), then registered: about 0.2 s/GB, i.e. 3–4 s for the 16 GB default. Because that cost is per byte of buffer, the streamer's 40 GB defaultRUNAI_STREAMER_MEMORY_LIMITis bounded to 16 GB whenpinned_stagingis set and the user did not setmemory_limit(16 GB showed no streaming penalty on this model). The buffer is unregistered in afinallywhen the iterator is exhausted.Where it does not help: at TP=8 on the same node the gain was 2–4 s on 47 s, because every rank reads the whole checkpoint through the page cache and the stage is a memcpy, not the copy path.
Stacked PR using this: the sleep-mode reload fix, #61304.
Pull Request Checklist
I used vLLM's
/pr-checklistskill. (Mandatory for agents, optional for humans).AI assistance was used during the creation of this PR.
Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.
Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI. (CPU unit tests added; hardware re-measure on
mainpending.)Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.
Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.
🤖 Generated with Claude Code