perf(model_loader): multi-threaded safetensors weight loading#287
perf(model_loader): multi-threaded safetensors weight loading#287yuanqingz wants to merge 5 commits into
Conversation
When prefetch_num_threads > 1, dispatch shard load_file() calls to a ThreadPoolExecutor and yield tensors as soon as each shard completes. Submissions are gated lazily by a generator expression so at most max_workers shards are resident in CPU memory at any time. Motivation: on NFS-backed checkpoints with many shards, the existing serial path is bottlenecked by per-file open/read latency. For Kimi K2.5-NVFP4 (119 safetensors shards, 553 GB) we measured ~50 min for model loading on B200 4-GPU; vLLM completes the equivalent path in ~6 min via a similar multi-threaded iterator (vllm/model_executor/ model_loader/weight_utils.py:multi_thread_safetensors_weights_iterator). The existing prefetch_checkpoint_files() helper that warms OS page cache is left in place for callers that need it explicitly; the serial path still calls it when prefetch=True. The new parallel path supersedes it for the common case. Wired into ServerArgs.weight_loader_prefetch_num_threads (default 4), so existing users automatically pick up the speedup with no flag change required. To opt out: --weight-loader-prefetch-num-threads 1. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 1f25e0fae5
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| # Lazy submission cap: only ``max_workers`` files in flight at once. | ||
| futures = (executor.submit(_load_file, st_file) for st_file in hf_weights_files) | ||
| futures_iter = tqdm( | ||
| concurrent.futures.as_completed(futures), |
There was a problem hiding this comment.
With the default weight_loader_prefetch_num_threads=4, any checkpoint with multiple safetensors shards now takes this path, but concurrent.futures.as_completed() materializes its input (set(fs)) before yielding, so this generator is consumed immediately and submits every shard at once. The executor only runs max_workers tasks concurrently, but completed futures keep their full state_dict results until the iterator reaches them, so large multi-shard checkpoints can still accumulate many loaded shards in CPU memory despite the intended cap of max_workers resident files.
Useful? React with 👍 / 👎.
|
@codex review |
|
Codex Review: Didn't find any major issues. More of your lovely PRs please. ℹ️ About Codex in GitHubCodex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback". |
…rs load
safetensors.torch.load_file(device="cpu") returns mmap-backed tensors.
With the serial loader, the consumer's per-tensor access overlaps with
NFS prefetch on the same file descriptor → no observable issue.
With the multi-threaded loader, by the time the consumer reaches file
B's tensors, the prefetch heuristic for B's fd is cold (a different
worker thread opened it), so each per-tensor access triggers an NFS
page fault. Net effect: load_file() parallelizes fine (3x speedup on
the safetensors phase alone for K2.5-NVFP4), but the downstream
DeepseekV3.load_weights() per-expert loop becomes ~13x slower than on
RAM-resident storage (>16 min vs 75 sec on /dev/shm).
Cloning each tensor inside the worker thread forces the page faults
to happen there, while we already pay NFS read cost as part of the
parallel I/O phase. Trade-off:
- +1 RAM copy of each shard while in flight (memory peak: 2x max_workers
shard size, still well under host RAM for typical large MoE models).
- Per-worker wall time goes up (mmap-touch → full read + copy), but the
parallelism still pays.
- Downstream consumer now only touches RAM-backed tensors.
Measured on Kimi K2.5-NVFP4 (553 GB, 119 shards) on Computelab B200
node with 4-thread loader against /home (NFS):
- Before clone(): safetensors 100% at +71s, model.load_weights()
still hung at +16 min — never reaches process_weights_after_loading.
- After clone(): <to be measured in the next bench run>
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
For a 553 GB / 119-shard model (Kimi K2.5-NVFP4) on a slow NFS backend, the multi-threaded safetensors loader needs more than 4 worker threads to saturate the aggregate read bandwidth. Measured on Computelab B200 node (against /home NFS, ~750 MB/s per-file ceiling): - 4 workers: ~19 min for 553 GB (~485 MB/s aggregate) - 8 workers: <to be measured next> Going from 4 → 8 mostly affects very large checkpoints; small models (< 10 GB) finish in either configuration in seconds. Peak CPU RAM during load doubles vs the 4-worker default, but is still bounded by ~2 × num_workers × shard_size = ~80 GB for typical 5 GB shards. That's well under the host RAM of any production-class server. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: f8c6219c93
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| decryption_key_file: str | None = None | ||
| weight_loader_prefetch_checkpoints: bool = False | ||
| weight_loader_prefetch_num_threads: int = 4 | ||
| weight_loader_prefetch_num_threads: int = 8 |
There was a problem hiding this comment.
With this default changed to 8, the checked-in runtime test test_load_config_defaults_keep_prefetch_disabled still asserts LoadConfig().weight_loader_prefetch_num_threads == 4 in test/runtime/test_weight_loader_prefetch.py:44. In an environment with the project dependencies installed, this commit will fail that test before exercising the new threaded loader, so either keep the default at 4 or update the expected default consistently.
Useful? React with 👍 / 👎.
1 participant
Parallelize the safetensors loader in
weight_utils.safetensors_weights_iterator. Whenprefetch_num_threads > 1, dispatch shardload_file()calls to aThreadPoolExecutorand yield tensors as soon as each shard completes. A sliding window keeps at mostmax_workersfiles resident in CPU memory at any time. After loading, each tensor is cloned in the worker thread so the bytes are RAM-materialized before being handed to the consumer — without this, the per-tensor access in the downstreammodel.load_weights()loop page-faults back through NFS and the parallel I/O is silently serialized.Mirrors vLLM's
multi_thread_safetensors_weights_iterator(invllm/model_executor/model_loader/weight_utils.py).Motivation
The existing serial path is bottlenecked by per-file open / read latency. On Kimi K2.5-NVFP4 (119 safetensors shards, 553 GiB) on a 4× B200 host backed by a shared NFS appliance, the serial loader takes ~20 min to materialize the checkpoint on cold client cache. With this change the same workload finishes in ~10 min on the same hardware.
Implementation
Four commits on this branch:
perf(model_loader): multi-threaded safetensors weight loading— adds_multi_thread_safetensors_weights_iterator, wires it intosafetensors_weights_iteratorwhenprefetch_num_threads > 1.fix(model_loader): cap concurrent safetensors shard loads— switches the loader fromas_completed(generator)(whichconcurrent.futuresdrains eagerly into a queue) to an explicit sliding window: submit the initialmax_workersfiles, then_submit_next()after each completion. Cancel pending futures onBaseException, shutdown executor infinally. Guarantees ≤max_workers × shard_sizeresident at any time.fix(model_loader): force materialize tensors after parallel safetensors load—safetensors.torch.load_file(device="cpu")returns mmap-backed tensors. With the serial loader the consumer accesses pages on the same fd that just opened them and NFS prefetch hides the latency. With the parallel loader the consumer reaches file B's tensors after worker thread B's prefetch heuristic has gone cold, so each per-tensor access triggers an NFS page fault. Cloning each tensor inside the worker thread moves the page faults to the parallel I/O phase. Trade-off: peak CPU memory ≈ 2 × max_workers × shard_size (still well under any production host RAM budget — measured 61 GiB peak resident in the 8-worker test, see below).perf(model_loader): bump default weight_loader_prefetch_num_threads to 8— for very large checkpoints (>500 GiB / >100 shards) 4 workers leaves throughput on the table.Existing users automatically pick up the speedup with no flag change required. Opt out with
--weight-loader-prefetch-num-threads 1.Test Plan — measured speedup (replicated)
Standalone loader-only timing harness, Kimi-K2.5-NVFP4 read directly from NFS — no tmpfs /
/dev/shmstaging — exercisingweight_utils.safetensors_weights_iteratorend-to-end with.clone()per tensor inside the consumer loop so we time actual NFS bandwidth, not lazy mmap setup. Each run is a fresh host allocation = cold client page cache. Server-side NFS cache state was similar across all four runs (submitted within minutes of each other against the same NFS volume).Observations:
resource.getrusage). Well within any production host RAM budget.Correctness end-to-end on the engine
Verified with
ts serve(--weight-loader-prefetch-num-threads 8, attn_tp4_moe_tp4 + EAGLE3 spec-dec, weights served from NFS). On all 4 TP ranks the full weight-loading pipeline completes on RAM-materialized tensors:Multi-thread loading shards (workers=8): 100% Completed | 119/119 [12:42<00:00](this PR's tqdm).model.load_weights()returns on all ranks — the per-expert loop inDeepseekV3.load_weights → moe/checkpoint/loader.load → load_model_weight → _load_w13/_load_w2proceeds normally on RAM-materialized tensors. Without the clone in commit (3) this phase hangs indefinitely (verified by cancelling a control run at ~28 min withmodel.load_weights()still not returned despite the iterator reporting 100% at +71 s — that 71 s was lazy mmap, not real I/O).process_weights_after_loadingloop completes (10 quant/self method calls forattn_tp4_moe_tp4).Load weight end. type=Eagle3DeepseekV2ForCausalLM, dtype=torch.bfloat16, avail mem=32.51 GB (weight_loader.py:116)on all 4 ranks.A separate engine-side crash in
profile_available_cache_memory_bytes → torch.distributed.all_reduce(Gloo "Connection closed by peer") reproduced on both 4- and 8-worker full-engine runs — that's post-Load weight endand does not involve the loader code path; it is being investigated separately (suspect: gateway startup timeout or peer rank exit during long warmup).Unit test
Added
test/runtime/test_weight_loader_prefetch.py(covers sliding-window correctness, exception cancellation, max_workers RSS bound).Notes for reviewers
prefetch_checkpoint_files()(the older page-cache warmer) is left in place forprefetch=Truecallers on the serial path. The new parallel path supersedes it for the common case.