Disable multi-threaded load by default when prefetch is on (#30146)
This commit is contained in:
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user