Disable multi-threaded load by default when prefetch is on (#30146)

This commit is contained in:
Mohammad Miadh Angkad
2026-07-09 00:28:53 -07:00
committed by GitHub
parent 64e2a73c80
commit 666a09fe2a
5 changed files with 224 additions and 5 deletions
+28
View File
@@ -576,6 +576,34 @@ class DefaultModelLoader(BaseModelLoader):
server_args.weight_loader_drop_cache_after_load
)
# Prefetch and multi-threaded loading both read the same shards,
# competing for I/O on shared/network storage. When prefetch is
# active (mmap path, not FASTSAFETENSORS) and the user didn't
# explicitly request multi-threaded loading, fall back to the
# single-threaded loader and let prefetch feed the page cache.
# Setting enable_multithread_load or num_threads in
# --model-loader-extra-config opts out (the latter is consumed
# only by the multi-threaded iterator, so it signals intent);
# e.g. local NVMe, where prefetch is a no-op and multi-threading
# helps.
if (
weight_loader_prefetch
and not weight_loader_disable_mmap
and self.load_config.load_format != LoadFormat.FASTSAFETENSORS
and use_multithread
and not (
{"enable_multithread_load", "num_threads"} & extra_config.keys()
)
):
logger.warning(
"--weight-loader-prefetch-checkpoints is enabled; falling "
"back to single-threaded weight loading to avoid I/O "
"oversubscription with the prefetch threads. Set "
"enable_multithread_load=true in --model-loader-extra-config "
"to keep multi-threaded loading."
)
use_multithread = False
if self.load_config.load_format == LoadFormat.FASTSAFETENSORS:
weights_iterator = fastsafetensors_weights_iterator(
hf_weights_files,
+1 -1
View File
@@ -2442,7 +2442,7 @@ class ServerArgs:
] = False
weight_loader_prefetch_checkpoints: A[
bool,
"Prefetch checkpoint files into OS page cache before loading. Each rank prefetches a fraction of the shards, reducing total network I/O on shared filesystems (NFS/Lustre) from N*checkpoint to 1*checkpoint. Recommended for models on network storage.",
"Prefetch checkpoint files into OS page cache before loading. Each rank prefetches a fraction of the shards, reducing total network I/O on shared filesystems (NFS/Lustre) from N*checkpoint to 1*checkpoint. Recommended for models on network storage. When enabled, multi-threaded safetensors loading is disabled by default to avoid I/O oversubscription with the prefetch threads; set enable_multithread_load=true in --model-loader-extra-config to keep multi-threaded loading (e.g. on local NVMe where prefetch is a no-op).",
] = False
weight_loader_prefetch_num_threads: A[
int,