diff --git a/vllm/model_executor/model_loader/weight_utils.py b/vllm/model_executor/model_loader/weight_utils.py index e3689dcd869..3b961e8e143 100644 --- a/vllm/model_executor/model_loader/weight_utils.py +++ b/vllm/model_executor/model_loader/weight_utils.py @@ -763,31 +763,24 @@ def np_cache_weights_iterator( yield name, torch.from_numpy(param) -def _checkpoints_fit_in_ram(files: list[str], threshold: float = 0.9) -> bool: - """Return True if total size of *files* fits within *threshold* of available RAM.""" +def _get_checkpoints_size_bytes(files: list[str]) -> int: + """Return the total size of the checkpoint files in bytes.""" if not files: - return True + return 0 + return sum(os.path.getsize(f) for f in files) + + +def _get_available_ram_bytes() -> int: + """Return the available RAM in bytes.""" import psutil - total_size = sum(os.path.getsize(f) for f in files) - available_ram = psutil.virtual_memory().available - fits = total_size <= threshold * available_ram - if not fits: - logger.warning( - "NFS detected but checkpoint total size (%.2f GiB) exceeds " - "%.0f%% of available RAM (%.2f GiB). Skipping prefetching checkpoints.", - total_size / (1024**3), - threshold * 100, - available_ram / (1024**3), - ) - return fits + return psutil.virtual_memory().available -def _is_nfs_path(files: list[str]) -> bool: - """Check whether the first file in *files* resides on an NFS - filesystem (Linux only).""" +def _get_fs_type(files: list[str]) -> str: + """Get the filesystem type of the first file in *files* (Linux only).""" if not files: - return False + return "" try: # Only the first file is checked — all checkpoint shards reside # in the same directory and therefore on the same filesystem. @@ -810,12 +803,11 @@ def _is_nfs_path(files: list[str]) -> bool: ) and len(mount_point) > len(best_mount): best_mount = mount_point best_fstype = fstype - return best_fstype in ("nfs", "nfs4") + return best_fstype except Exception: # /proc/mounts is Linux-specific; on other OSes (or if the read - # fails for any reason) we fall back to "not NFS" rather than - # crashing model loading. - return False + # fails for any reason) we fall back to an empty string. + return "" def _prefetch_checkpoint(file_path: str) -> None: @@ -901,11 +893,63 @@ def safetensors_weights_iterator( sorted_files = sorted(hf_weights_files, key=_natural_sort_key) - should_prefetch = safetensors_load_strategy == "prefetch" or ( - safetensors_load_strategy is None - and _is_nfs_path(sorted_files) - and _checkpoints_fit_in_ram(sorted_files) + fs_type = _get_fs_type(sorted_files) + is_net_fs = fs_type in ("nfs", "nfs4", "lustre") + total_bytes = _get_checkpoints_size_bytes(sorted_files) + avail_bytes = _get_available_ram_bytes() + ram_threshold_pct = 90 + fits_in_ram = total_bytes <= (ram_threshold_pct / 100.0) * avail_bytes + fs_name = fs_type.upper() if fs_type else "unknown" + + logger.info_once( + "Filesystem type for checkpoints: %s. Checkpoint size: %.2f GiB. " + "Available RAM: %.2f GiB.", + fs_name, + total_bytes / 1024**3, + avail_bytes / 1024**3, ) + + should_prefetch = safetensors_load_strategy == "prefetch" + if safetensors_load_strategy is None: + if is_net_fs and fits_in_ram: + should_prefetch = True + elif is_net_fs and not fits_in_ram: + logger.warning_once( + "Network filesystem (%s) detected but checkpoint total size " + "(%.2f GiB) exceeds %d%% of available RAM (%.2f GiB). " + "Skipping auto-prefetch.", + fs_name, + total_bytes / 1024**3, + ram_threshold_pct, + avail_bytes / 1024**3, + ) + elif not is_net_fs and fits_in_ram: + logger.info_once( + "Auto-prefetch is disabled because the filesystem (%s) is not a " + "recognized network FS (NFS/Lustre). If you want to force " + "prefetching, start vLLM with --safetensors-load-strategy=prefetch.", + fs_name, + ) + elif not is_net_fs and not fits_in_ram: + logger.info_once( + "Auto-prefetch is disabled because the filesystem (%s) is not a " + "recognized network FS (NFS/Lustre) and the checkpoint size " + "(%.2f GiB) exceeds %d%% of available RAM (%.2f GiB).", + fs_name, + total_bytes / 1024**3, + ram_threshold_pct, + avail_bytes / 1024**3, + ) + elif should_prefetch and not fits_in_ram: + logger.warning_once( + "safetensors_load_strategy='prefetch' was explicitly specified, but " + "checkpoint total size (%.2f GiB) exceeds %d%% of available RAM " + "(%.2f GiB). This may cause out-of-memory errors.", + total_bytes / 1024**3, + ram_threshold_pct, + avail_bytes / 1024**3, + ) + if should_prefetch: _prefetch_all_checkpoints(sorted_files)