# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Utilities for downloading and initializing model weights.""" import asyncio import concurrent.futures import fnmatch import glob import hashlib import json import os import tempfile import threading import time from collections import defaultdict from collections.abc import Callable, Generator, Iterable from contextlib import contextmanager from pathlib import Path from typing import IO, Any import filelock import huggingface_hub.constants import numpy as np import regex as re import torch from safetensors.torch import load, load_file, safe_open, save_file from tqdm.auto import tqdm from transformers.utils import SAFE_WEIGHTS_INDEX_NAME from vllm import envs from vllm.config import ModelConfig from vllm.config.load import ( DEFAULT_SAFETENSORS_PREFETCH_BLOCK_SIZE, DEFAULT_SAFETENSORS_PREFETCH_NUM_THREADS, LoadConfig, ) from vllm.distributed import get_tensor_model_parallel_rank, get_world_group from vllm.logger import init_logger from vllm.model_executor.layers.quantization import ( QuantizationConfig, get_quantization_config, ) from vllm.model_executor.model_loader.ep_weight_filter import ( should_skip_weight, ) from vllm.platforms import current_platform from vllm.tracing import instrument from vllm.transformers_utils.repo_utils import hf_api, hf_fs from vllm.utils.import_utils import PlaceholderModule try: from runai_model_streamer import SafetensorsStreamer except ImportError: runai_model_streamer = PlaceholderModule("runai_model_streamer") # type: ignore[assignment] SafetensorsStreamer = runai_model_streamer.placeholder_attr("SafetensorsStreamer") try: from fastsafetensors import SingleGroup except ImportError: fastsafetensors = PlaceholderModule("fastsafetensors") SingleGroup = fastsafetensors.placeholder_attr("SingleGroup") from vllm.model_executor.layers.quantization.torchao import torchao_version_at_least logger = init_logger(__name__) # use system-level temp directory for file locks, so that multiple users # can share the same lock without error. # lock files in the temp directory will be automatically deleted when the # system reboots, so users will not complain about annoying lock files temp_dir = tempfile.gettempdir() def enable_xet_high_performance(): """automatically activates xet high performance mode""" if "HF_XET_HIGH_PERFORMANCE" not in os.environ: huggingface_hub.constants.HF_XET_HIGH_PERFORMANCE = True enable_xet_high_performance() class DisabledTqdm(tqdm): def __init__(self, *args, **kwargs): kwargs["disable"] = True super().__init__(*args, **kwargs) def get_lock(model_name_or_path: str | Path, cache_dir: str | None = None): lock_dir = cache_dir or temp_dir model_name_or_path = str(model_name_or_path) os.makedirs(os.path.dirname(lock_dir), exist_ok=True) model_name = model_name_or_path.replace("/", "-") hash_name = hashlib.sha256(model_name.encode()).hexdigest() # add hash to avoid conflict with old users' lock files lock_file_name = hash_name + model_name + ".lock" # mode 0o666 is required for the filelock to be shared across users lock = filelock.FileLock(os.path.join(lock_dir, lock_file_name), mode=0o666) return lock @contextmanager def atomic_writer( filepath: str | Path, mode: str = "w", encoding: str | None = None ) -> Generator[IO]: """ Context manager that provides an atomic file writing routine. The context manager writes to a temporary file and, if successful, atomically replaces the original file. Args: filepath (str or Path): The path to the file to write. mode (str): The file mode for the temporary file (e.g., 'w', 'wb'). encoding (str): The encoding for text mode. Yields: file object: A handle to the temporary file. """ # Create a temporary file in the same directory as the target file # to ensure it's on the same filesystem for an atomic replace. temp_dir = os.path.dirname(filepath) temp_fd, temp_path = tempfile.mkstemp(dir=temp_dir) try: # Open the temporary file for writing with os.fdopen(temp_fd, mode=mode, encoding=encoding) as temp_file: yield temp_file # If the 'with' block completes successfully, # perform the atomic replace. os.replace(temp_path, filepath) except Exception: logger.exception( "Error during atomic write. Original file '%s' not modified", filepath ) raise finally: # Clean up the temporary file if it still exists. if os.path.exists(temp_path): os.remove(temp_path) def _natural_sort_key(filepath: str) -> list: """Natural sort key for filenames with numeric components, such as model-00001-of-00005.safetensors -> ['model-', 1, '-of-', 5, '.safetensors']""" return [ int(s) if s.isdigit() else s for s in re.split(r"(\d+)", os.path.basename(filepath)) ] def maybe_download_from_modelscope( model: str, revision: str | None = None, download_dir: str | None = None, ignore_patterns: str | list[str] | None = None, allow_patterns: list[str] | str | None = None, ) -> str | None: """Download model from ModelScope hub if VLLM_USE_MODELSCOPE is True. Returns the path to the downloaded model, or None if the model is not downloaded from ModelScope.""" if envs.VLLM_USE_MODELSCOPE: # download model from ModelScope hub, # lazy import so that modelscope is not required for normal use. # pylint: disable=C. from modelscope.hub.snapshot_download import snapshot_download # Use file lock to prevent multiple processes from # downloading the same model weights at the same time. with get_lock(model, download_dir): if not os.path.exists(model): model_path = snapshot_download( model_id=model, cache_dir=download_dir, local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE, revision=revision, ignore_file_pattern=ignore_patterns, allow_patterns=allow_patterns, ) else: model_path = model return model_path return None def _shared_pointers(tensors): ptrs = defaultdict(list) for k, v in tensors.items(): ptrs[v.data_ptr()].append(k) failing = [] for _, names in ptrs.items(): if len(names) > 1: failing.append(names) return failing def convert_bin_to_safetensor_file( pt_filename: str, sf_filename: str, ) -> None: loaded = torch.load(pt_filename, map_location="cpu", weights_only=True) if "state_dict" in loaded: loaded = loaded["state_dict"] shared = _shared_pointers(loaded) for shared_weights in shared: for name in shared_weights[1:]: loaded.pop(name) # For tensors to be contiguous loaded = {k: v.contiguous() for k, v in loaded.items()} dirname = os.path.dirname(sf_filename) os.makedirs(dirname, exist_ok=True) save_file(loaded, sf_filename, metadata={"format": "pt"}) # check file size sf_size = os.stat(sf_filename).st_size pt_size = os.stat(pt_filename).st_size if (sf_size - pt_size) / pt_size > 0.01: raise RuntimeError( f"""The file size different is more than 1%: - {sf_filename}: {sf_size} - {pt_filename}: {pt_size} """ ) # check if the tensors are the same reloaded = load_file(sf_filename) for k in loaded: pt_tensor = loaded[k] sf_tensor = reloaded[k] if not torch.equal(pt_tensor, sf_tensor): raise RuntimeError(f"The output tensors do not match for key {k}") # TODO(woosuk): Move this to other place. def get_quant_config( model_config: ModelConfig, load_config: LoadConfig ) -> QuantizationConfig: if model_config.quantization is None: raise ValueError("Model quantization method is not specified in the config.") quant_cls = get_quantization_config(model_config.quantization) # Read the quantization config from the HF model config, if available. hf_quant_config = getattr(model_config.hf_config, "quantization_config", None) # some vision model may keep quantization_config in their text_config hf_text_config = getattr(model_config.hf_config, "text_config", None) if hf_quant_config is None and hf_text_config is not None: hf_quant_config = getattr(hf_text_config, "quantization_config", None) if hf_quant_config is None: # compressed-tensors uses a compressions_config hf_quant_config = getattr(model_config.hf_config, "compression_config", None) # Pipe information about heads to enable TP-aware loading of attn_head scales if ( hf_quant_config is not None and hf_quant_config.get("quant_method") == "compressed-tensors" and "config_groups" in hf_quant_config ): if hf_text_config is not None: n_heads = getattr(hf_text_config, "num_attention_heads", None) n_kv_heads = getattr(hf_text_config, "num_key_value_heads", None) else: n_heads = getattr(model_config.hf_config, "num_attention_heads", None) n_kv_heads = getattr(model_config.hf_config, "num_key_value_heads", None) hf_quant_config["total_num_heads"] = n_heads hf_quant_config["total_num_kv_heads"] = ( n_kv_heads if n_kv_heads is not None else n_heads ) if hf_quant_config is not None: # `model_config.quantization_config` may be set alongside a checkpoint # quant config: the checkpoint determines `quant_cls`, and the user's # QuantizationConfigArgs is consulted by individual quant methods # (e.g. for activation overrides via the MXFP4 oracle). # For modelopt_mixed, config.json's quantization_config may or may # not contain the per-layer quantized_layers map. Newer checkpoints # embed it directly; older ones keep it only in hf_quant_config.json. # If it is missing, fall through to the file-based loading path. if ( model_config.quantization == "modelopt_mixed" and "quantized_layers" not in hf_quant_config ): pass # fall through to file-based loading below else: return quant_cls.from_config(hf_quant_config) # if hf_quant_config is None, we will try to get config from # hf_overrides hf_overrides = model_config.hf_overrides if not isinstance(hf_overrides, dict): raise ValueError( "hf_overrides must be a dict for get_quant_config " "to get the quantization config from it." ) quantization_config_file = hf_overrides.get("quantization_config_file", None) if quantization_config_file is not None: if hasattr(quant_cls, "from_config_file"): return quant_cls.from_config_file(quantization_config_file) else: raise NotImplementedError( "from_config_file is specified in hf_override config, " "but quant_cls.from_config_file is not implemented in " f"{quant_cls}" ) quantization_config_json = hf_overrides.get("quantization_config_dict_json", None) if quantization_config_json is not None: if hasattr(quant_cls, "from_config_dict_json"): return quant_cls.from_config_dict_json(quantization_config_json) else: raise NotImplementedError( "from_config_dict_json is specified in hf_override config, " "but quant_cls.from_config_dict_json is not implemented in " f"{quant_cls}" ) # Online quantization doesn't read from checkpoint configs - it quantizes # fp16/bf16 weights on the fly during loading. if model_config.quantization_config is not None: from vllm.config.quantization import QuantizationConfigArgs from vllm.model_executor.layers.quantization.online.base import ( OnlineQuantizationConfig, ) assert isinstance(model_config.quantization_config, QuantizationConfigArgs) return OnlineQuantizationConfig(args=model_config.quantization_config) # Inflight BNB quantization if model_config.quantization == "bitsandbytes": return quant_cls.from_config({}) model_name_or_path = ( maybe_download_from_modelscope( model_config.model, revision=model_config.revision, download_dir=load_config.download_dir, allow_patterns=["*.json"], ) or model_config.model ) is_local = os.path.isdir(model_name_or_path) if not is_local: # Download the config files. with get_lock(model_config.model, load_config.download_dir): hf_folder = hf_api().snapshot_download( model_config.model, revision=model_config.revision, allow_patterns="*.json", cache_dir=load_config.download_dir, local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE, tqdm_class=DisabledTqdm, ) else: hf_folder = model_name_or_path possible_config_filenames = quant_cls.get_config_filenames() # If the quantization config is not found, use the default config. if not possible_config_filenames: return quant_cls() config_files = glob.glob(os.path.join(hf_folder, "*.json")) quant_config_files = [ f for f in config_files if any(f.endswith(x) for x in possible_config_filenames) ] if len(quant_config_files) == 0: raise ValueError(f"Cannot find the config file for {model_config.quantization}") if len(quant_config_files) > 1: raise ValueError( f"Found multiple config files for {model_config.quantization}: " f"{quant_config_files}" ) quant_config_file = quant_config_files[0] with open(quant_config_file) as f: config = json.load(f) if model_config.quantization == "bitsandbytes": config["adapter_name_or_path"] = model_config.model elif model_config.quantization in ("modelopt", "modelopt_mixed"): if config.get("producer", {}).get("name") == "modelopt": return quant_cls.from_config(config) else: raise ValueError( f"Unsupported quantization config" f" found for {model_config.quantization} in {f}." ) return quant_cls.from_config(config) def get_sparse_attention_config( model_config: ModelConfig, load_config: LoadConfig, sparse_attention_config_filename: str = "sparse_attention_config.json", ) -> dict[str, Any]: model_name_or_path = model_config.model is_local = os.path.isdir(model_name_or_path) if not is_local: # Download the config files. with get_lock(model_name_or_path, load_config.download_dir): hf_folder = hf_api().snapshot_download( model_name_or_path, revision=model_config.revision, allow_patterns="*.json", cache_dir=load_config.download_dir, local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE, tqdm_class=DisabledTqdm, ) else: hf_folder = model_name_or_path config_file = os.path.join(hf_folder, sparse_attention_config_filename) if not os.path.exists(config_file): return {} # Load the sparse attention config. with open(config_file) as f: config = json.load(f) logger.info("Loaded sparse attention config from %s", config_file) return config @instrument(span_name="Download weights - HF") def download_weights_from_hf( model_name_or_path: str, cache_dir: str | None, allow_patterns: list[str], revision: str | None = None, subfolder: str | None = None, ignore_patterns: str | list[str] | None = None, ) -> str: """Download model weights from Hugging Face Hub. Args: model_name_or_path (str): The model name or path. cache_dir (Optional[str]): The cache directory to store the model weights. If None, will use HF defaults. allow_patterns (list[str]): The allowed patterns for the weight files. Files matched by any of the patterns will be downloaded. revision (Optional[str]): The revision of the model. subfolder (Optional[str]): The subfolder within the model repository to download weights from. ignore_patterns (Optional[Union[str, list[str]]]): The patterns to filter out the weight files. Files matched by any of the patterns will be ignored. Returns: str: The path to the downloaded model weights. """ assert len(allow_patterns) > 0 local_only = huggingface_hub.constants.HF_HUB_OFFLINE if not local_only: # Attempt to reduce allow_patterns to a single pattern # so we only have to call snapshot_download once. try: fs = hf_fs() file_list = fs.ls( os.path.join(model_name_or_path, subfolder or ""), detail=False, revision=revision, ) # If downloading safetensors and an index file exists, use the # specific file names from the index to avoid downloading # unnecessary files (e.g., from subdirectories like "original/"). index_file = f"{model_name_or_path}/{SAFE_WEIGHTS_INDEX_NAME}" if "*.safetensors" in allow_patterns and index_file in file_list: index_path = hf_api().hf_hub_download( repo_id=model_name_or_path, filename=SAFE_WEIGHTS_INDEX_NAME, cache_dir=cache_dir, revision=revision, subfolder=subfolder, ) with open(index_path) as f: weight_map = json.load(f)["weight_map"] if weight_map: # Extra [] so that weight_map files are treated as a # single allow_pattern in the loop below allow_patterns = [list(set(weight_map.values()))] # type: ignore[list-item] else: allow_patterns = ["*.safetensors"] else: # Use the first pattern found in the HF repo's files. for pattern in allow_patterns: if fnmatch.filter(file_list, pattern): allow_patterns = [pattern] break except Exception as e: logger.warning( "Failed to get file list for '%s'. Trying each pattern in " "allow_patterns individually until weights have been " "downloaded. Error: %s", model_name_or_path, e, ) logger.debug("Using model weights format %s", allow_patterns) # Use file lock to prevent multiple processes from # downloading the same model weights at the same time. with get_lock(model_name_or_path, cache_dir): start_time = time.perf_counter() for allow_pattern in allow_patterns: hf_folder = hf_api().snapshot_download( model_name_or_path, allow_patterns=allow_pattern, ignore_patterns=ignore_patterns, cache_dir=cache_dir, tqdm_class=DisabledTqdm, revision=revision, local_files_only=local_only, ) # If we have downloaded weights for this allow_pattern, # we don't need to check the rest. # allow_pattern can be a list (from weight_map) or str (glob) if isinstance(allow_pattern, list): break if any(Path(hf_folder).glob(allow_pattern)): break time_taken = time.perf_counter() - start_time if time_taken > 0.5: logger.info( "Time spent downloading weights for %s: %.6f seconds", model_name_or_path, time_taken, ) return hf_folder def download_safetensors_index_file_from_hf( model_name_or_path: str, index_file: str, cache_dir: str | None, subfolder: str | None = None, revision: str | None = None, ) -> None: """Download hf safetensors index file from Hugging Face Hub. Args: model_name_or_path (str): The model name or path. index_file (str): The safetensors index file name cache_dir (Optional[str]): The cache directory to store the model weights. If None, will use HF defaults. subfolder (Optional[str]): The subfolder within the model repository to download weights from. revision (Optional[str]): The revision of the model. """ # Use file lock to prevent multiple processes from # downloading the same model weights at the same time. with get_lock(model_name_or_path, cache_dir): try: # Download the safetensors index file. hf_api().hf_hub_download( repo_id=model_name_or_path, filename=index_file, cache_dir=cache_dir, revision=revision, subfolder=subfolder, local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE, ) # If file not found on remote or locally, we should not fail since # only some models will have index_file. except huggingface_hub.utils.LocalEntryNotFoundError: logger.info("No %s found in local cache.", index_file) except huggingface_hub.utils.EntryNotFoundError: logger.info("No %s found in remote.", index_file) # For models like Mistral-7B-v0.3, there are both sharded # safetensors files and a consolidated safetensors file. # Passing both of these to the weight loader functionality breaks. # So, we use the index_file to # look up which safetensors files should be used. def filter_duplicate_safetensors_files( hf_weights_files: list[str], hf_folder: str, index_file: str ) -> list[str]: # model.safetensors.index.json is a mapping from keys in the # torch state_dict to safetensors file holding that weight. index_file_name = os.path.join(hf_folder, index_file) if not os.path.isfile(index_file_name): return hf_weights_files # Iterate through the weight_map (weight_name: safetensors files) # to identify weights that we should use. with open(index_file_name) as f: weight_map = json.load(f)["weight_map"] weight_files_in_index = set() for weight_name in weight_map: weight_files_in_index.add(os.path.join(hf_folder, weight_map[weight_name])) # Check if files referenced in model.safetensors.index.json actually exist. # Raise error if any file is missing. hf_weights_files_set = set(hf_weights_files) missing_files = weight_files_in_index - hf_weights_files_set if missing_files: raise FileNotFoundError( f"Weight files referenced in index but missing: {missing_files}" ) # Filter out any fields that are not found in the index file. hf_weights_files = [f for f in hf_weights_files if f in weight_files_in_index] return hf_weights_files def filter_files_not_needed_for_inference(hf_weights_files: list[str]) -> list[str]: """ Exclude files that are not needed for inference. See https://github.com/huggingface/transformers/blob/v4.34.0/src/transformers/trainer.py#L227-L233 """ blacklist = [ "training_args.bin", "optimizer.bin", "optimizer.pt", "scheduler.pt", "scaler.pt", ] hf_weights_files = [ f for f in hf_weights_files if not any(f.endswith(x) for x in blacklist) ] return hf_weights_files # explicitly use pure text format, with a newline at the end # this makes it impossible to see the animation in the progress bar # but will avoid messing up with ray or multiprocessing, which wraps # each line of output with some prefix. _BAR_FORMAT = "{desc}: {percentage:3.0f}% Completed | {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]\n" # noqa: E501 def enable_tqdm(use_tqdm_on_load: bool): return use_tqdm_on_load and ( not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0 ) def np_cache_weights_iterator( model_name_or_path: str, cache_dir: str | None, hf_folder: str, hf_weights_files: list[str], use_tqdm_on_load: bool, ) -> Generator[tuple[str, torch.Tensor], None, None]: """Iterate over the weights in the model np files. Will dump the model weights to numpy files if they are not already dumped. """ # Convert the model weights from torch tensors to numpy arrays for # faster loading. np_folder = os.path.join(hf_folder, "np") os.makedirs(np_folder, exist_ok=True) weight_names_file = os.path.join(np_folder, "weight_names.json") # Use file lock to prevent multiple processes from # dumping the same model weights to numpy at the same time. with get_lock(model_name_or_path, cache_dir): if not os.path.exists(weight_names_file): weight_names: list[str] = [] for bin_file in tqdm( hf_weights_files, desc="Loading np_cache checkpoint shards", disable=not enable_tqdm(use_tqdm_on_load), bar_format=_BAR_FORMAT, ): state = torch.load(bin_file, map_location="cpu", weights_only=True) for name, param in state.items(): param_path = os.path.join(np_folder, name) with open(param_path, "wb") as f: np.save(f, param.cpu().detach().numpy()) weight_names.append(name) with open(weight_names_file, "w") as f: json.dump(weight_names, f) with open(weight_names_file) as f: weight_names = json.load(f) for name in weight_names: param_path = os.path.join(np_folder, name) with open(param_path, "rb") as f: param = np.load(f) yield name, torch.from_numpy(param) def _get_checkpoints_size_bytes(files: list[str]) -> int: """Return the total size of the checkpoint files in bytes.""" if not files: 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 return psutil.virtual_memory().available def _get_fs_type(files: list[str]) -> str: """Get the filesystem type of the first file in *files* (Linux only).""" if not files: return "" try: # Only the first file is checked — all checkpoint shards reside # in the same directory and therefore on the same filesystem. resolved = os.path.realpath(files[0]) best_mount = "" best_fstype = "" # /proc/mounts may contain nested mount points (e.g. "/" -> ext4, # "/data" -> nfs4, "/data/local" -> ext4). We pick the entry with # the longest matching mount_point — the same "longest prefix match" # rule the kernel uses to decide which filesystem serves a path. with open("/proc/mounts") as f: for line in f: parts = line.split() if len(parts) < 3: continue mount_point, fstype = parts[1], parts[2] if ( resolved == mount_point or resolved.startswith(os.path.join(mount_point, "")) ) and len(mount_point) > len(best_mount): best_mount = mount_point best_fstype = fstype 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 an empty string. return "" def _prefetch_checkpoint( file_path: str, block_size: int = DEFAULT_SAFETENSORS_PREFETCH_BLOCK_SIZE, ) -> None: """Prefetch a checkpoint file into the OS page cache. Reads the file in blocks so the kernel caches its pages before workers load the same file. """ if block_size < 1: raise ValueError("safetensors prefetch block size must be >= 1") with open(file_path, "rb") as f: while f.read(block_size): pass def _prefetch_all_checkpoints( sorted_files: list[str], num_prefetch_threads: int = DEFAULT_SAFETENSORS_PREFETCH_NUM_THREADS, block_size: int = DEFAULT_SAFETENSORS_PREFETCH_BLOCK_SIZE, ) -> None: """Start prefetching checkpoint files into page cache in a background thread.""" if num_prefetch_threads < 1: raise ValueError("safetensors prefetch num threads must be >= 1") if block_size < 1: raise ValueError("safetensors prefetch block size must be >= 1") if torch.distributed.is_initialized(): rank = torch.distributed.get_rank() world_size = torch.distributed.get_world_size() else: rank = 0 world_size = 1 paths_to_prefetch = sorted_files[rank::world_size] total_for_rank = len(paths_to_prefetch) async def _prefetch_all() -> None: loop = asyncio.get_running_loop() completed = 0 next_log_pct = 10 async def prefetch_one( path: str, executor: concurrent.futures.ThreadPoolExecutor, ) -> None: nonlocal completed, next_log_pct try: await loop.run_in_executor( executor, _prefetch_checkpoint, path, block_size ) completed += 1 if total_for_rank > 0 and next_log_pct <= 100: pct = 100 * completed / total_for_rank if pct >= next_log_pct: logger.info( "Prefetching checkpoint files: %d%% (%d/%d)", next_log_pct, completed, total_for_rank, ) next_log_pct += 10 except Exception: logger.warning( "Failed to prefetch checkpoint file %r.", path, exc_info=True ) with concurrent.futures.ThreadPoolExecutor( max_workers=num_prefetch_threads ) as executor: await asyncio.gather( *(prefetch_one(p, executor) for p in paths_to_prefetch) ) def _run_prefetch() -> None: start = time.perf_counter() asyncio.run(_prefetch_all()) elapsed = time.perf_counter() - start logger.info( "Prefetching checkpoint files into page cache finished in %.2fs", elapsed, ) logger.info( "Prefetching checkpoint files into page cache started " "(in background, num_threads=%d, block_size=%d bytes)", num_prefetch_threads, block_size, ) threading.Thread(target=_run_prefetch, daemon=True).start() def safetensors_weights_iterator( hf_weights_files: list[str], use_tqdm_on_load: bool, safetensors_load_strategy: str | None = None, local_expert_ids: set[int] | None = None, *, safetensors_prefetch_num_threads: int = DEFAULT_SAFETENSORS_PREFETCH_NUM_THREADS, safetensors_prefetch_block_size: int = DEFAULT_SAFETENSORS_PREFETCH_BLOCK_SIZE, ) -> Generator[tuple[str, torch.Tensor], None, None]: """Iterate over the weights in the model safetensor files. When *local_expert_ids* is provided, expert weights not belonging to this rank are skipped **before** reading from disk, which drastically reduces storage I/O for MoE models under EP. """ loading_desc = "Loading safetensors checkpoint shards" if safetensors_load_strategy == "eager": loading_desc += " (eager)" sorted_files = sorted(hf_weights_files, key=_natural_sort_key) 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, num_prefetch_threads=safetensors_prefetch_num_threads, block_size=safetensors_prefetch_block_size, ) leftover_state_dict: dict[str, torch.Tensor] = {} for st_file in tqdm( sorted_files, desc=loading_desc, disable=not enable_tqdm(use_tqdm_on_load), bar_format=_BAR_FORMAT, ): if safetensors_load_strategy == "eager": with open(st_file, "rb") as f: state_dict = load(f.read()) for name, param in state_dict.items(): if not should_skip_weight(name, local_expert_ids): yield name, param elif safetensors_load_strategy == "torchao": # we can't load flattened torchao tensor subclasses directly into the model # instead we reconstruct the subclasses here before returning if not torchao_version_at_least("0.15.0"): raise ValueError( "Please use torchao version >= 0.15.0 " "to load torchao safetensors checkpoint" ) from torchao.prototype.safetensors.safetensors_support import ( unflatten_tensor_state_dict, ) with safe_open(st_file, framework="pt") as f: state_dict = {} for name in f.keys(): # noqa: SIM118 if should_skip_weight(name, local_expert_ids): continue state_dict[name] = f.get_tensor(name) # update with leftover tensor data from previous iteration, if any state_dict.update(leftover_state_dict) metadata = f.metadata() # due to sharded checkpoints, we are not guaranteed that we have all # tensor subclass data on one file # state_dict has the leftover data from this step and we wait for # missing information to be provided in a future iteration unflattened_state_dict, leftover_state_dict = ( unflatten_tensor_state_dict(state_dict, metadata) ) yield from unflattened_state_dict.items() else: with safe_open(st_file, framework="pt") as f: for name in f.keys(): # noqa: SIM118 if should_skip_weight(name, local_expert_ids): continue param = f.get_tensor(name) yield name, param def multi_thread_safetensors_weights_iterator( hf_weights_files: list[str], use_tqdm_on_load: bool, max_workers: int = 4, ) -> Generator[tuple[str, torch.Tensor], None, None]: """Multi-Thread iterate over the weights in the model safetensor files.""" def _load_file(st_file: str): result = load_file(st_file, device="cpu") return result with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor: # Note to use generator here so we do not store all the loaded files in memory # at the same time, which can cause OOM for large models. futures = (executor.submit(_load_file, st_file) for st_file in hf_weights_files) futures_iter = tqdm( concurrent.futures.as_completed(futures), total=len(hf_weights_files), desc="Multi-thread loading shards", disable=not enable_tqdm(use_tqdm_on_load), bar_format=_BAR_FORMAT, ) for future in futures_iter: state_dict = future.result() del future for key in list(state_dict): yield key, state_dict.pop(key) def runai_safetensors_weights_iterator( hf_weights_files: list[str], use_tqdm_on_load: bool, is_distributed: bool = False, ) -> Generator[tuple[str, torch.Tensor], None, None]: """Iterate over the weights in the model safetensor files.""" with SafetensorsStreamer() as streamer: is_cuda_alike = current_platform.is_cuda_alike() device = ( f"cuda:{current_platform.current_device()}" if is_distributed and is_cuda_alike else "cpu" ) streamer.stream_files( hf_weights_files, device=device, is_distributed=is_distributed, ) total_tensors = sum( len(tensors_meta) for tensors_meta in streamer.files_to_tensors_metadata.values() ) tensor_iter = tqdm( streamer.get_tensors(), total=total_tensors, desc="Loading safetensors using Runai Model Streamer", bar_format=_BAR_FORMAT, disable=not enable_tqdm(use_tqdm_on_load), mininterval=2, ) for name, tensor in tensor_iter: yield name, tensor.clone() def fastsafetensors_weights_iterator( hf_weights_files: list[str], use_tqdm_on_load: bool, ) -> Generator[tuple[str, torch.Tensor], None, None]: """Iterate over the weights in the model safetensor files using fastsafetensor library. Uses ParallelLoader for pipelined loading: the producer thread prepares metadata for the next shard while the consumer yields tensors from the current shard. """ from fastsafetensors.parallel_loader import ParallelLoader if torch.distributed.is_initialized(): pg = torch.distributed.group.WORLD else: pg = SingleGroup() device = torch.device(f"cuda:{current_platform.current_device()}") hf_weights_files = sorted(hf_weights_files, key=_natural_sort_key) # Use nogds=True for TP > 1 to avoid cuFileDriverOpen() which # initializes the GDS DMA subsystem for all visible GPUs, creating # unwanted CUDA contexts on every device. nogds = pg.size() > 1 queue_size = envs.VLLM_FASTSAFETENSORS_QUEUE_SIZE tqdm_enabled = enable_tqdm(use_tqdm_on_load) def _make_loader(nogds: bool) -> "ParallelLoader": return ParallelLoader( pg=pg, hf_weights_files=hf_weights_files, queue_size=queue_size, use_tqdm_on_load=tqdm_enabled, device=str(device), nogds=nogds, ) # GDS can fail either at construction or lazily inside the producer # thread during iteration (e.g. cuFileHandleRegister returning # CU_FILE_HANDLE_NOT_REGISTERED on a filesystem without GDS support). # Catch both and fall back to nogds, but only before yielding any # tensor -- restarting mid-stream would reload earlier shards. pl = None yielded = False try: try: pl = _make_loader(nogds) for name, tensor in pl.iterate_weights(): yielded = True yield name, tensor except RuntimeError as e: if nogds or yielded or "gds" not in str(e): raise logger.warning_once( "GDS not enabled, setting `nogds=True`.\n" "For more information, see: https://github.com/foundation-model-stack/" "fastsafetensors?tab=readme-ov-file#basic-api-usages" ) if pl is not None: pl.close() pl = _make_loader(nogds=True) yield from pl.iterate_weights() finally: if pl is not None: pl.close() def instanttensor_weights_iterator( hf_weights_files: list[str], use_tqdm_on_load: bool, ) -> Generator[tuple[str, torch.Tensor], None, None]: """Iterate over the weights in the model safetensor files using instanttensor library.""" try: import instanttensor except ImportError as e: raise ImportError( "Please install instanttensor via `pip install instanttensor`" ) from e if not current_platform.is_cuda(): raise ValueError("InstantTensor requires NVIDIA GPUs") try: world_group = get_world_group() except AssertionError: # Entering here only in unit tests where the world group is not initialized. process_group = None else: process_group = world_group.device_group if world_group.world_size > 1 else None device = current_platform.current_device() with instanttensor.safe_open( hf_weights_files, framework="pt", device=device, process_group=process_group ) as f: yield from tqdm( f.tensors(), desc="Loading safetensors using InstantTensor loader", disable=not enable_tqdm(use_tqdm_on_load), bar_format=_BAR_FORMAT, position=tqdm._get_free_pos(), total=len(f.keys()), mininterval=1.0, ) def pt_weights_iterator( hf_weights_files: list[str], use_tqdm_on_load: bool, pt_load_map_location: str | dict[str, str] = "cpu", ) -> Generator[tuple[str, torch.Tensor], None, None]: """Iterate over the weights in the model bin/pt files.""" for bin_file in tqdm( hf_weights_files, desc="Loading pt checkpoint shards", disable=not enable_tqdm(use_tqdm_on_load), bar_format=_BAR_FORMAT, ): state = torch.load( bin_file, map_location=pt_load_map_location, weights_only=True ) yield from state.items() del state def multi_thread_pt_weights_iterator( hf_weights_files: list[str], use_tqdm_on_load: bool, pt_load_map_location: str | dict[str, str] = "cpu", max_workers: int = 4, ) -> Generator[tuple[str, torch.Tensor], None, None]: """Multi-Thread iterate over the weights in the model bin/pt files.""" def _load_file(bin_file: str): return torch.load( bin_file, map_location=pt_load_map_location, weights_only=True ) with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor: futures = [ executor.submit(_load_file, bin_file) for bin_file in hf_weights_files ] futures_iter = tqdm( concurrent.futures.as_completed(futures), total=len(hf_weights_files), desc="Multi-thread loading pt checkpoint shards", disable=not enable_tqdm(use_tqdm_on_load), bar_format=_BAR_FORMAT, ) for future in futures_iter: state = future.result() yield from state.items() del state def convert_pyslice_to_tensor(x: Any) -> torch.Tensor: """convert PySafeSlice object from safetensors to torch.Tensor PySafeSlice object supports indexing, which is done before loading the actual tensor and can reduce the amount of memory being read into the memory. However, it does not support more advanced functionalities like `.view()` or `.t()`. Therefore, if we need to modify the loaded tensor with these more complicated operators, we need to convert to tensor first. """ if not isinstance(x, torch.Tensor): x = x[:] return x def default_weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: """Default weight loader.""" try: if param.numel() == 1 and loaded_weight.numel() == 1: # Sometimes scalar values aren't considered tensors with shapes # so if both param and loaded_weight are a scalar, # reshape to match before copying param.data.copy_(loaded_weight.view(param.shape)) else: assert param.size() == loaded_weight.size(), ( f"Attempted to load weight ({loaded_weight.size()}) " f"into parameter ({param.size()})" ) param.data.copy_(loaded_weight) except Exception: # NOTE: This exception is added for the purpose of setting breakpoint to # debug weight loading issues. raise def row_parallel_weight_loader( param: torch.Tensor, loaded_weight: torch.Tensor ) -> None: """Load weights that are row-parallelized.""" tp_rank = get_tensor_model_parallel_rank() shard_dim = 0 if param.dim() != 1 else None if shard_dim is not None: shard_size = param.data.shape[shard_dim] start_idx = tp_rank * shard_size loaded_weight = loaded_weight.narrow(shard_dim, start_idx, shard_size) return default_weight_loader(param, loaded_weight) LoaderFunction = Callable[[torch.Tensor, torch.Tensor], None] def sharded_weight_loader(shard_axis: int) -> LoaderFunction: """Create a weight loader that shards the weights along the given axis""" def loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: tp_rank = get_tensor_model_parallel_rank() shard_size = param.data.shape[shard_axis] start_idx = tp_rank * shard_size loaded_weight = loaded_weight.narrow(shard_axis, start_idx, shard_size) return default_weight_loader(param, loaded_weight) return loader def composed_weight_loader( loader: LoaderFunction, fn: Callable[[torch.Tensor], torch.Tensor] ) -> LoaderFunction: """Create a weight loader that post-processes the weights after loading""" def composed_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: loader(param, loaded_weight) param.data.copy_(fn(param)) return return composed_loader def initialize_dummy_weights( model: torch.nn.Module, model_config: ModelConfig, low: float = -1e-3, high: float = 1e-3, seed: int = 1234, ) -> None: """Initialize model weights with random values. The model weights must be randomly initialized for accurate performance measurements. Additionally, the model weights should not cause NaNs in the forward pass. We empirically found that initializing the weights with values between -1e-3 and 1e-3 works well for most models. We use per-parameter random seed, so that dummy weights are consistent, even if the model is partitioned across multiple devices. When the seed is fixed, the random values generated by this function only depends on the parameter's number of elements and its data type. """ for param in model.state_dict().values(): initialize_single_dummy_weight(param, low, high, seed) @torch.no_grad() def initialize_single_dummy_weight( param: torch.Tensor, low: float = -1e-3, high: float = 1e-3, seed: int = 1234, ) -> None: if param.device.type == "meta": return # deferred to finalize_layerwise_processing (e.g. online quant) if not torch.is_floating_point(param): if current_platform.is_rocm(): # On ROCm, integer params (e.g. GPTQ qweight/qzeros) are left # as torch.empty() by default, giving non-deterministic values # across processes. Zero them for reproducibility. param.zero_() return if current_platform.is_tpu(): generator = torch.Generator(device="cpu") generator.manual_seed(seed) # Note: The param.uniform_ function cannot be used in this # context because it demands more TPU HBM than directly copying # from a CPU tensor. # Note: We avoid using torch.rank_like as it doesn't currently # support the generator argument. param.copy_( (high - low) * torch.rand( param.shape, generator=generator, dtype=param.dtype, layout=param.layout, requires_grad=param.requires_grad, device="cpu", ) + low ) torch._sync(param) return generator = torch.Generator(device=param.data.device) generator.manual_seed(seed) if torch.finfo(param.data.dtype).bits < 16: # uniform_ doesn't support < 16-bit datatypes (FP8) dtype = param.data.dtype tmp_param = param.data.to(torch.float16) tmp_param = tmp_param.uniform_(low, high, generator=generator).to(dtype) param.data.copy_(tmp_param) else: param.uniform_(low, high, generator=generator) def maybe_remap_kv_scale_name(name: str, params_dict: dict) -> str | None: """Remap the name of FP8 k/v_scale parameters. This function handles the remapping of FP8 k/v_scale parameter names. It detects if the given name ends with a suffix and attempts to remap it to the expected name format in the model. If the remapped name is not found in the params_dict, a warning is printed and None is returned. Args: name (str): The original loaded checkpoint parameter name. params_dict (dict): Dictionary containing the model's named parameters. Returns: str: The remapped parameter name if successful, or the original name if no remapping is needed. None: If the remapped name is not found in params_dict. """ # Already in vLLM's expected form (e.g. weights pre-renamed by a # `WeightsMapper` from the quant config). Skip the regex remap, which # would otherwise double-apply the `.attn` prefix and drop the weight. if name in params_dict: return name if name.endswith(".kv_scale"): logger.warning_once( "DEPRECATED. Found kv_scale in the checkpoint. " "This format is deprecated in favor of separate k_scale and " "v_scale tensors and will be removed in a future release. " "Functionally, we will remap kv_scale to k_scale and duplicate " "k_scale to v_scale" ) # NOTE: we remap the deprecated kv_scale to k_scale remapped_name = name.replace(".kv_scale", ".attn.k_scale") if remapped_name not in params_dict: logger.warning_once( "Found kv_scale in the checkpoint (e.g. %s), but not found the expected name in the model (e.g. %s). kv_scale is not loaded.", # noqa: E501 name, remapped_name, ) return None return remapped_name if any("mla_attn" in key for key in params_dict): attn_str = "mla_attn.mla_attn" logger.debug_once( f"Found mla_attn with k_scale and v_scale in " f"the checkpoint, using {attn_str} as attn_str" ) else: attn_str = "attn" # Define scale name mapping patterns in order of precedence scale_mapping_patterns = [ # ModelOpt format: .self_attn.{k,v}_proj.{k,v}_scale -> # .self_attn.attn.{k,v}_scale ( r"\.self_attn\.([kv])_proj\.([kv])_scale$", rf".self_attn.{attn_str}.\2_scale", ), # QKV proj format: .self_attn.qkv_proj.{k,v}_scale -> # .self_attn.attn.{k,v}_scale (r"\.self_attn\.qkv_proj\.([kv])_scale$", r".self_attn.attn.\1_scale"), # Qwen3 MoE format: .self_attn.qkqkv_proj.{k,v}_scale -> # .self_attn.attn.{k,v}_scale (r"\.self_attn\.qkqkv_proj\.([kv])_scale$", r".self_attn.attn.\1_scale"), # NemotronH format: .mixer.{k,v}_proj.{k,v}_scale -> # .mixer.attn.{k,v}_scale (r"\.mixer\.[kv]_proj\.([kv])_scale$", r".mixer.attn.\1_scale"), # HYV3 format: .self_attn.q.scale -> .self_attn.attn.q_scale (r"\.self_attn\.q\.scale$", r".self_attn.attn.q_scale"), # HYV3 format: .self_attn.{k,v}_cache.scale -> # .self_attn.attn.{k,v}_scale (r"\.self_attn\.([kv])_cache\.scale$", r".self_attn.attn.\1_scale"), # Default format: .{k,v}_scale -> .attn.{k,v}_scale (r"\.([qkv])_scale$", r".attn.\1_scale"), (r"\.([qkv])_zero_point$", r".attn.\1_zero_point"), ] # Check if name ends with k_scale or v_scale if name.endswith( ( ".k_scale", ".v_scale", ".q_scale", ".k_zero_point", ".v_zero_point", ".q_zero_point", ".q.scale", ".k_cache.scale", ".v_cache.scale", ) ): import regex as re for pattern, replacement in scale_mapping_patterns: if re.search(pattern, name): remapped_name = re.sub(pattern, replacement, name) if remapped_name not in params_dict: scale_type = name.split(".")[-1] logger.warning_once( "Found %s in the checkpoint (e.g. %s), but not found the expected name in the model (e.g. %s). %s is not loaded.", # noqa: E501 scale_type, name, remapped_name, scale_type, ) return None return remapped_name # If there were no matches, return the untouched param name return name def maybe_remap_moe_expert_param_name( name: str, params_dict: dict[str, torch.nn.Parameter], ) -> str: """ Remap MoE expert parameter names to account for routed_experts hierarchy. This handles the transition from the old FusedMoE structure where weights were directly in the experts module, to the new MoERunner → RoutedExperts structure. Checkpoint weights have names like: layers.0.mlp.experts.w13_weight layers.0.feed_forward.experts.w2_input_scale But actual parameters are now: layers.0.mlp.experts.routed_experts.w13_weight layers.0.feed_forward.experts.routed_experts.w2_input_scale This function inserts 'routed_experts.' into the path when needed. Args: name: Parameter name from checkpoint params_dict: Dictionary of model parameters (from named_parameters()) Returns: Remapped parameter name if routed_experts hierarchy exists, otherwise the original name """ # Only remap if this looks like an expert parameter if ".experts." not in name: return name # Skip if already has routed_experts if ".experts.routed_experts." in name: return name # Expert parameter patterns to check expert_param_suffixes = [ "w13_weight", "w2_weight", "w13_weight_scale", "w2_weight_scale", "w13_input_scale", "w2_input_scale", "w13_bias", "w2_bias", "w13_scale", "w2_scale", "w13_g_idx", "w2_g_idx", "w13_qweight", "w2_qweight", "w13_qzeros", "w2_qzeros", "w13_weight_shape", "w2_weight_shape", ] # Check if this is an expert weight parameter is_expert_param = any( f".{suffix}" in name or name.endswith(suffix) for suffix in expert_param_suffixes ) if not is_expert_param: return name # Try inserting routed_experts after .experts. new_name = name.replace(".experts.", ".experts.routed_experts.", 1) # Only use the new name if it exists in the model if new_name in params_dict: return new_name # Otherwise return original name (old checkpoint format or different structure) return name def remap_moe_expert_weights( weights: Iterable[tuple[str, torch.Tensor]], params_dict: dict[str, torch.nn.Parameter], ) -> Generator[tuple[str, torch.Tensor], None, None]: """ Wrapper generator that remaps MoE expert parameter names for backward compatibility. This allows models with custom weight loading to automatically handle both old and new checkpoint formats without needing model-specific remapping code. Usage: params_dict = dict(model.named_parameters()) for name, weight in remap_moe_expert_weights(weights, params_dict): # name is automatically remapped if needed param = params_dict[name] ... Args: weights: Iterator of (name, tensor) tuples from checkpoint params_dict: Dictionary of model parameters (from named_parameters()) Yields: (remapped_name, tensor) tuples """ for name, weight in weights: remapped_name = maybe_remap_moe_expert_param_name(name, params_dict) yield (remapped_name, weight)