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 5, 08:30 UTC
[Model][Perf] Enable checkpoints prefetching for Lustre FS by default (#39422)
Signed-off-by: Artem Perevedentsev <aperevedents@nvidia.com>
Artem Perevedentsev committed
Apr 10, 2026 at 03:47 UTC
bb6047db132045a33c8d11d698b7cf3cc8e68706
1 file changed
+71
-27
vllm/model_executor/model_loader/weight_utils.py
+71
-27
@@ -763,31 +763,24 @@ def np_cache_weights_iterator(
763
yield name, torch.from_numpy(param)
764
765
766
-def _checkpoints_fit_in_ram(files: list[str], threshold: float = 0.9) -> bool:
767
- """Return True if total size of *files* fits within *threshold* of available RAM."""
766
+def _get_checkpoints_size_bytes(files: list[str]) -> int:
767
+ """Return the total size of the checkpoint files in bytes."""
768
if not files:
769
- return True
769
+ return 0
770
+ return sum(os.path.getsize(f) for f in files)
771
+
772
+
773
+def _get_available_ram_bytes() -> int:
774
+ """Return the available RAM in bytes."""
775
import psutil
776
772
- total_size = sum(os.path.getsize(f) for f in files)
773
- available_ram = psutil.virtual_memory().available
774
- fits = total_size <= threshold * available_ram
775
- if not fits:
776
- logger.warning(
777
- "NFS detected but checkpoint total size (%.2f GiB) exceeds "
778
- "%.0f%% of available RAM (%.2f GiB). Skipping prefetching checkpoints.",
779
- total_size / (1024**3),
780
- threshold * 100,
781
- available_ram / (1024**3),
782
- )
783
- return fits
777
+ return psutil.virtual_memory().available
778
779
786
-def _is_nfs_path(files: list[str]) -> bool:
787
- """Check whether the first file in *files* resides on an NFS
788
- filesystem (Linux only)."""
780
+def _get_fs_type(files: list[str]) -> str:
781
+ """Get the filesystem type of the first file in *files* (Linux only)."""
782
if not files:
790
- return False
783
+ return ""
784
try:
785
# Only the first file is checked — all checkpoint shards reside
786
# in the same directory and therefore on the same filesystem.
@@ -810,12 +803,11 @@ def _is_nfs_path(files: list[str]) -> bool:
803
) and len(mount_point) > len(best_mount):
804
best_mount = mount_point
805
best_fstype = fstype
813
- return best_fstype in ("nfs", "nfs4")
806
+ return best_fstype
807
except Exception:
808
# /proc/mounts is Linux-specific; on other OSes (or if the read
816
- # fails for any reason) we fall back to "not NFS" rather than
817
- # crashing model loading.
818
- return False
809
+ # fails for any reason) we fall back to an empty string.
810
+ return ""
811
812
813
def _prefetch_checkpoint(file_path: str) -> None:
@@ -901,11 +893,63 @@ def safetensors_weights_iterator(
893
894
sorted_files = sorted(hf_weights_files, key=_natural_sort_key)
895
904
- should_prefetch = safetensors_load_strategy == "prefetch" or (
905
- safetensors_load_strategy is None
906
- and _is_nfs_path(sorted_files)
907
- and _checkpoints_fit_in_ram(sorted_files)
896
+ fs_type = _get_fs_type(sorted_files)
897
+ is_net_fs = fs_type in ("nfs", "nfs4", "lustre")
898
+ total_bytes = _get_checkpoints_size_bytes(sorted_files)
899
+ avail_bytes = _get_available_ram_bytes()
900
+ ram_threshold_pct = 90
901
+ fits_in_ram = total_bytes <= (ram_threshold_pct / 100.0) * avail_bytes
902
+ fs_name = fs_type.upper() if fs_type else "unknown"
903
+
904
+ logger.info_once(
905
+ "Filesystem type for checkpoints: %s. Checkpoint size: %.2f GiB. "
906
+ "Available RAM: %.2f GiB.",
907
+ fs_name,
908
+ total_bytes / 1024**3,
909
+ avail_bytes / 1024**3,
910
)
911
+
912
+ should_prefetch = safetensors_load_strategy == "prefetch"
913
+ if safetensors_load_strategy is None:
914
+ if is_net_fs and fits_in_ram:
915
+ should_prefetch = True
916
+ elif is_net_fs and not fits_in_ram:
917
+ logger.warning_once(
918
+ "Network filesystem (%s) detected but checkpoint total size "
919
+ "(%.2f GiB) exceeds %d%% of available RAM (%.2f GiB). "
920
+ "Skipping auto-prefetch.",
921
+ fs_name,
922
+ total_bytes / 1024**3,
923
+ ram_threshold_pct,
924
+ avail_bytes / 1024**3,
925
+ )
926
+ elif not is_net_fs and fits_in_ram:
927
+ logger.info_once(
928
+ "Auto-prefetch is disabled because the filesystem (%s) is not a "
929
+ "recognized network FS (NFS/Lustre). If you want to force "
930
+ "prefetching, start vLLM with --safetensors-load-strategy=prefetch.",
931
+ fs_name,
932
+ )
933
+ elif not is_net_fs and not fits_in_ram:
934
+ logger.info_once(
935
+ "Auto-prefetch is disabled because the filesystem (%s) is not a "
936
+ "recognized network FS (NFS/Lustre) and the checkpoint size "
937
+ "(%.2f GiB) exceeds %d%% of available RAM (%.2f GiB).",
938
+ fs_name,
939
+ total_bytes / 1024**3,
940
+ ram_threshold_pct,
941
+ avail_bytes / 1024**3,
942
+ )
943
+ elif should_prefetch and not fits_in_ram:
944
+ logger.warning_once(
945
+ "safetensors_load_strategy='prefetch' was explicitly specified, but "
946
+ "checkpoint total size (%.2f GiB) exceeds %d%% of available RAM "
947
+ "(%.2f GiB). This may cause out-of-memory errors.",
948
+ total_bytes / 1024**3,
949
+ ram_threshold_pct,
950
+ avail_bytes / 1024**3,
951
+ )
952
+
953
if should_prefetch:
954
_prefetch_all_checkpoints(sorted_files)
955