A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 6, 00:45 UTC
[Performance] Auto-enable prefetch on NFS with RAM guard (#37673)
Signed-off-by: Artem Perevedentsev <aperevedents@nvidia.com>
Artem Perevedentsev committed
Mar 25, 2026 at 02:31 UTC
a93a53f8a1302c992ad185e70c6ab4affe43c4d7
3 files changed
+69
-5
vllm/config/load.py
+6
-2
@@ -54,10 +54,14 @@ class LoadConfig:
54
download_dir: str | None = None
55
"""Directory to download and load the weights, default to the default
56
cache directory of Hugging Face."""
57
- safetensors_load_strategy: str = "lazy"
57
+ safetensors_load_strategy: str | None = None
58
"""Specifies the loading strategy for safetensors weights.
59
- - "lazy" (default): Weights are memory-mapped from the file. This enables
59
+ - None (default): Uses memory-mapped (lazy) loading. When an NFS
60
+ filesystem is detected and the total checkpoint size fits within 90%%
61
+ of available RAM, prefetching is enabled automatically.
62
+ - "lazy": Weights are memory-mapped from the file. This enables
63
on-demand loading and is highly efficient for models on local storage.
64
+ Unlike the default (None), auto-prefetch on NFS is not performed.
65
- "eager": The entire file is read into CPU memory upfront before loading.
66
This is recommended for models on network filesystems (e.g., Lustre, NFS)
67
as it avoids inefficient random reads, significantly speeding up model
vllm/engine/arg_utils.py
+1
-1
@@ -388,7 +388,7 @@ class EngineArgs:
388
allowed_local_media_path: str = ModelConfig.allowed_local_media_path
389
allowed_media_domains: list[str] | None = ModelConfig.allowed_media_domains
390
download_dir: str | None = LoadConfig.download_dir
391
- safetensors_load_strategy: str = LoadConfig.safetensors_load_strategy
391
+ safetensors_load_strategy: str | None = LoadConfig.safetensors_load_strategy
392
load_format: str | LoadFormats = LoadConfig.load_format
393
config_format: str = ModelConfig.config_format
394
dtype: ModelDType = ModelConfig.dtype
vllm/model_executor/model_loader/weight_utils.py
+62
-2
@@ -729,6 +729,61 @@ def np_cache_weights_iterator(
729
yield name, torch.from_numpy(param)
730
731
732
+def _checkpoints_fit_in_ram(files: list[str], threshold: float = 0.9) -> bool:
733
+ """Return True if total size of *files* fits within *threshold* of available RAM."""
734
+ if not files:
735
+ return True
736
+ import psutil
737
+
738
+ total_size = sum(os.path.getsize(f) for f in files)
739
+ available_ram = psutil.virtual_memory().available
740
+ fits = total_size <= threshold * available_ram
741
+ if not fits:
742
+ logger.warning(
743
+ "NFS detected but checkpoint total size (%.2f GiB) exceeds "
744
+ "%.0f%% of available RAM (%.2f GiB). Skipping prefetching checkpoints.",
745
+ total_size / (1024**3),
746
+ threshold * 100,
747
+ available_ram / (1024**3),
748
+ )
749
+ return fits
750
+
751
+
752
+def _is_nfs_path(files: list[str]) -> bool:
753
+ """Check whether the first file in *files* resides on an NFS
754
+ filesystem (Linux only)."""
755
+ if not files:
756
+ return False
757
+ try:
758
+ # Only the first file is checked — all checkpoint shards reside
759
+ # in the same directory and therefore on the same filesystem.
760
+ resolved = os.path.realpath(files[0])
761
+ best_mount = ""
762
+ best_fstype = ""
763
+ # /proc/mounts may contain nested mount points (e.g. "/" -> ext4,
764
+ # "/data" -> nfs4, "/data/local" -> ext4). We pick the entry with
765
+ # the longest matching mount_point — the same "longest prefix match"
766
+ # rule the kernel uses to decide which filesystem serves a path.
767
+ with open("/proc/mounts") as f:
768
+ for line in f:
769
+ parts = line.split()
770
+ if len(parts) < 3:
771
+ continue
772
+ mount_point, fstype = parts[1], parts[2]
773
+ if (
774
+ resolved == mount_point
775
+ or resolved.startswith(os.path.join(mount_point, ""))
776
+ ) and len(mount_point) > len(best_mount):
777
+ best_mount = mount_point
778
+ best_fstype = fstype
779
+ return best_fstype in ("nfs", "nfs4")
780
+ except Exception:
781
+ # /proc/mounts is Linux-specific; on other OSes (or if the read
782
+ # fails for any reason) we fall back to "not NFS" rather than
783
+ # crashing model loading.
784
+ return False
785
+
786
+
787
def _prefetch_checkpoint(file_path: str) -> None:
788
"""Prefetch a checkpoint file into the OS page cache.
789
@@ -797,7 +852,7 @@ def _prefetch_all_checkpoints(sorted_files: list[str]) -> None:
852
def safetensors_weights_iterator(
853
hf_weights_files: list[str],
854
use_tqdm_on_load: bool,
800
- safetensors_load_strategy: str = "lazy",
855
+ safetensors_load_strategy: str | None = None,
856
local_expert_ids: set[int] | None = None,
857
) -> Generator[tuple[str, torch.Tensor], None, None]:
858
"""Iterate over the weights in the model safetensor files.
@@ -812,7 +867,12 @@ def safetensors_weights_iterator(
867
868
sorted_files = sorted(hf_weights_files, key=_natural_sort_key)
869
815
- if safetensors_load_strategy == "prefetch":
870
+ should_prefetch = safetensors_load_strategy == "prefetch" or (
871
+ safetensors_load_strategy is None
872
+ and _is_nfs_path(sorted_files)
873
+ and _checkpoints_fit_in_ram(sorted_files)
874
+ )
875
+ if should_prefetch:
876
_prefetch_all_checkpoints(sorted_files)
877
878
leftover_state_dict: dict[str, torch.Tensor] = {}