# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import dataclasses import glob import os import time from collections.abc import Generator, Iterable from typing import cast import torch from torch import nn from transformers.utils import SAFE_WEIGHTS_INDEX_NAME from vllm.config import ModelConfig from vllm.config.load import LoadConfig from vllm.logger import init_logger from vllm.model_executor.layers.quantization.torchao import torchao_version_at_least from vllm.model_executor.model_loader.base_loader import BaseModelLoader from vllm.model_executor.model_loader.ep_weight_filter import ( compute_local_expert_ids, ) from vllm.model_executor.model_loader.weight_utils import ( download_safetensors_index_file_from_hf, download_weights_from_hf, fastsafetensors_weights_iterator, filter_duplicate_safetensors_files, filter_files_not_needed_for_inference, get_quant_config, instanttensor_weights_iterator, maybe_download_from_modelscope, multi_thread_pt_weights_iterator, multi_thread_safetensors_weights_iterator, np_cache_weights_iterator, pt_weights_iterator, safetensors_weights_iterator, ) from vllm.tracing import instrument from vllm.transformers_utils.repo_utils import list_filtered_repo_files logger = init_logger(__name__) class DefaultModelLoader(BaseModelLoader): """Model loader that can load different file types from disk.""" # default number of thread when enable multithread weight loading DEFAULT_NUM_THREADS = 8 @dataclasses.dataclass class Source: """A source for weights.""" model_or_path: str """The model ID or path.""" revision: str | None """The optional model revision.""" subfolder: str | None = None """The subfolder inside the model repo.""" prefix: str = "" """A prefix to prepend to all weights.""" fall_back_to_pt: bool = True """Whether .pt weights can be used.""" allow_patterns_overrides: list[str] | None = None """If defined, weights will load exclusively using these patterns.""" counter_before_loading_weights: float = 0.0 counter_after_loading_weights: float = 0.0 def __init__(self, load_config: LoadConfig): super().__init__(load_config) self.local_expert_ids: set[int] | None = None extra_config = load_config.model_loader_extra_config if not isinstance(extra_config, dict): raise ValueError( f"model_loader_extra_config must be a dict for load format " f"{load_config.load_format}, got {type(extra_config).__name__}" ) allowed_keys = { "enable_multithread_load", "num_threads", "enable_weights_track", } unexpected_keys = set(extra_config.keys()) - allowed_keys if unexpected_keys: raise ValueError( f"Unexpected extra config keys for load format " f"{load_config.load_format}: " f"{unexpected_keys}" ) enable_multithread_load = extra_config.get("enable_multithread_load", False) if not isinstance(enable_multithread_load, bool): raise ValueError( f"enable_multithread_load must be a bool, got " f"{type(enable_multithread_load).__name__}" ) num_threads = extra_config.get("num_threads") if num_threads is not None and not ( isinstance(num_threads, int) and num_threads > 0 ): raise ValueError( f"num_threads must be a positive integer, got {num_threads!r}" ) self.enable_weights_track: bool | None = extra_config.get( "enable_weights_track", None ) # The multi-thread loader ignores safetensors_load_strategy, so reject # the combination instead of silently dropping the requested strategy. if extra_config.get("enable_multithread_load") and ( load_config.safetensors_load_strategy not in (None, "lazy") ): raise ValueError( "enable_multithread_load does not support " "safetensors_load_strategy=" f"{load_config.safetensors_load_strategy!r}; the multi-thread " "loader only implements the default lazy strategy." ) def _prepare_weights( self, model_name_or_path: str, subfolder: str | None, revision: str | None, fall_back_to_pt: bool, allow_patterns_overrides: list[str] | None, ) -> tuple[str, list[str], bool]: """Prepare weights for the model. If the model is not local, it will be downloaded.""" model_name_or_path = ( maybe_download_from_modelscope(model_name_or_path, revision) or model_name_or_path ) is_local = os.path.isdir(model_name_or_path) load_format = self.load_config.load_format use_safetensors = False index_file = SAFE_WEIGHTS_INDEX_NAME # First check for 'auto' format that mistral files format are present. # This is to load mistral models with official format by default. if load_format == "auto": load_format = ( "mistral" if len( list_filtered_repo_files( model_name_or_path=model_name_or_path, allow_patterns=["consolidated*.safetensors"], revision=revision, ) ) > 0 else "hf" ) # Some quantized models use .pt files for storing the weights. if load_format == "hf": allow_patterns = ["*.safetensors", "*.bin"] elif ( load_format == "safetensors" or load_format == "fastsafetensors" or load_format == "instanttensor" ): use_safetensors = True allow_patterns = ["*.safetensors"] elif load_format == "mistral": use_safetensors = True allow_patterns = ["consolidated*.safetensors"] index_file = "consolidated.safetensors.index.json" elif load_format == "pt": allow_patterns = ["*.pt"] elif load_format == "npcache": allow_patterns = ["*.bin"] else: raise ValueError(f"Unknown load_format: {load_format}") # Don't fall back to .pt for explicit safetensors formats; otherwise a # .pt file is matched and later opened as safetensors. if fall_back_to_pt and not use_safetensors: allow_patterns += ["*.pt"] if allow_patterns_overrides is not None: allow_patterns = allow_patterns_overrides if not is_local: hf_folder = download_weights_from_hf( model_name_or_path, self.load_config.download_dir, allow_patterns, revision, subfolder=subfolder, ignore_patterns=self.load_config.ignore_patterns, ) else: hf_folder = model_name_or_path if subfolder is not None: hf_folder = os.path.join(hf_folder, subfolder) hf_weights_files: list[str] = [] for pattern in allow_patterns: hf_weights_files += glob.glob(os.path.join(hf_folder, pattern)) if len(hf_weights_files) > 0: if pattern.endswith(".safetensors"): use_safetensors = True break if use_safetensors: # For models like Mistral-7B-Instruct-v0.3 # there are both sharded safetensors files and a consolidated # safetensors file. Using both breaks. # Here, we download the `model.safetensors.index.json` and filter # any files not found in the index. if not is_local: download_safetensors_index_file_from_hf( model_name_or_path, index_file, cache_dir=self.load_config.download_dir, subfolder=subfolder, revision=revision, ) hf_weights_files = filter_duplicate_safetensors_files( hf_weights_files, hf_folder, index_file ) else: hf_weights_files = filter_files_not_needed_for_inference(hf_weights_files) if len(hf_weights_files) == 0: raise RuntimeError( f"Cannot find any model weights with `{model_name_or_path}`" ) return hf_folder, hf_weights_files, use_safetensors def _get_weights_iterator( self, source: "Source" ) -> Generator[tuple[str, torch.Tensor], None, None]: """Get an iterator for the model weights based on the load format.""" extra_config = self.load_config.model_loader_extra_config hf_folder, hf_weights_files, use_safetensors = self._prepare_weights( source.model_or_path, source.subfolder, source.revision, source.fall_back_to_pt, source.allow_patterns_overrides, ) if self.load_config.load_format == "npcache": # Currently np_cache only support *.bin checkpoints assert use_safetensors is False weights_iterator = np_cache_weights_iterator( source.model_or_path, self.load_config.download_dir, hf_folder, hf_weights_files, self.load_config.use_tqdm_on_load, ) elif use_safetensors: if self.load_config.load_format == "fastsafetensors": weights_iterator = fastsafetensors_weights_iterator( hf_weights_files, self.load_config.use_tqdm_on_load, ) elif self.load_config.load_format == "instanttensor": weights_iterator = instanttensor_weights_iterator( hf_weights_files, self.load_config.use_tqdm_on_load, ) else: if extra_config.get("enable_multithread_load"): weights_iterator = multi_thread_safetensors_weights_iterator( hf_weights_files, self.load_config.use_tqdm_on_load, max_workers=extra_config.get( "num_threads", self.DEFAULT_NUM_THREADS ), ) else: weights_iterator = safetensors_weights_iterator( hf_weights_files, self.load_config.use_tqdm_on_load, self.load_config.safetensors_load_strategy, local_expert_ids=self.local_expert_ids, safetensors_prefetch_num_threads=( self.load_config.safetensors_prefetch_num_threads ), safetensors_prefetch_block_size=( self.load_config.safetensors_prefetch_block_size ), ) else: if extra_config.get("enable_multithread_load"): weights_iterator = multi_thread_pt_weights_iterator( hf_weights_files, self.load_config.use_tqdm_on_load, self.load_config.pt_load_map_location, max_workers=extra_config.get( "num_threads", self.DEFAULT_NUM_THREADS ), ) else: weights_iterator = pt_weights_iterator( hf_weights_files, self.load_config.use_tqdm_on_load, self.load_config.pt_load_map_location, ) if self.counter_before_loading_weights == 0.0: self.counter_before_loading_weights = time.perf_counter() # Apply the prefix. return ((source.prefix + name, tensor) for (name, tensor) in weights_iterator) def get_all_weights( self, model_config: ModelConfig, model: nn.Module, ) -> Generator[tuple[str, torch.Tensor], None, None]: primary_weights = DefaultModelLoader.Source( model_config.model, model_config.revision, prefix="", fall_back_to_pt=getattr(model, "fall_back_to_pt_during_load", True), allow_patterns_overrides=getattr(model, "allow_patterns_overrides", None), ) yield from self._get_weights_iterator(primary_weights) secondary_weights = cast( Iterable[DefaultModelLoader.Source], getattr(model, "secondary_weights", ()), ) for source in secondary_weights: yield from self._get_weights_iterator(source) def download_model(self, model_config: ModelConfig) -> None: self._prepare_weights( model_name_or_path=model_config.model, subfolder=None, revision=model_config.revision, fall_back_to_pt=True, allow_patterns_overrides=None, ) def _init_ep_weight_filter(self, model_config: ModelConfig) -> None: """Compute local expert ids for EP weight filtering. When expert parallelism is active, each rank only needs a subset of expert weights. By computing the set upfront we can skip non-local expert tensors *before* reading them from disk. """ from vllm.config import get_current_vllm_config vllm_config = get_current_vllm_config() parallel_config = vllm_config.parallel_config if not ( model_config.is_moe and parallel_config.enable_expert_parallel and parallel_config.enable_ep_weight_filter ): return # When EPLB is enabled, redundant physical expert slots may map to # logical experts that belong to other ranks in the default partition. # The weight loader needs to see ALL logical expert weights so it can # populate these redundant slots. Skip the filter entirely. if parallel_config.enable_eplb: return num_experts = model_config.get_num_experts() if num_experts <= 0: return # EP size/rank computation mirrors FusedMoEParallelConfig.make(): # ep_size = dp_size * pcp_size * tp_size (flattened) # ep_rank = dp_rank * pcp_size * tp_size + pcp_rank * tp_size + tp_rank from vllm.distributed import ( get_dp_group, get_pcp_group, get_tensor_model_parallel_rank, ) dp_size = parallel_config.data_parallel_size tp_size = parallel_config.tensor_parallel_size pcp_size = parallel_config.prefill_context_parallel_size dp_rank = get_dp_group().rank_in_group if dp_size > 1 else 0 tp_rank = get_tensor_model_parallel_rank() if tp_size > 1 else 0 pcp_rank = get_pcp_group().rank_in_group if pcp_size > 1 else 0 ep_size = dp_size * pcp_size * tp_size ep_rank = dp_rank * pcp_size * tp_size + pcp_rank * tp_size + tp_rank self.local_expert_ids = compute_local_expert_ids( num_experts, ep_size, ep_rank, placement=parallel_config.expert_placement_strategy, ) if self.local_expert_ids is not None: logger.info_once( "EP weight filter: ep_size=%d, ep_rank=%d, loading %d/%d experts", ep_size, ep_rank, len(self.local_expert_ids), num_experts, ) @instrument(span_name="Load weights") def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None: if model_config.quantization == "torchao": quant_config = get_quant_config(model_config, self.load_config) if ( hasattr(quant_config, "is_checkpoint_torchao_serialized") and quant_config.is_checkpoint_torchao_serialized and torchao_version_at_least("0.15.0") ): self.load_config.safetensors_load_strategy = "torchao" self._init_ep_weight_filter(model_config) loaded_weights = model.load_weights(self.get_all_weights(model_config, model)) self.counter_after_loading_weights = time.perf_counter() logger.info_once( "Loading weights took %.2f seconds", self.counter_after_loading_weights - self.counter_before_loading_weights, ) # We only enable strict check for non-quantized models # that have loaded weights tracking by default. default_enable_weights_track = ( model_config.quantization is None and loaded_weights is not None ) enable_weights_track = ( self.enable_weights_track if self.enable_weights_track is not None else default_enable_weights_track ) if enable_weights_track: self.track_weights_loading(model, loaded_weights) def track_weights_loading( self, model: nn.Module, loaded_weights: set[str] | None ) -> None: weights_to_load = {name for name, _ in model.named_parameters()} if loaded_weights is not None: # ignore online quantization scales for name, module in model.named_modules(): quant_method = getattr(module, "quant_method", None) has_online_quant = getattr(quant_method, "uses_meta_device", False) has_postprocess_quant = getattr( quant_method, "process_weights_after_loading", None ) # ignore kv_cache scale and online quant scale, # which can be missing in checkpoints if has_online_quant or has_postprocess_quant: for param_name, _ in module.named_parameters(): full_name = f"{name}.{param_name}" if name else param_name loaded_weights.add(full_name) weights_not_loaded = weights_to_load - loaded_weights if weights_not_loaded: raise ValueError( "Following weights were not initialized from " f"checkpoint: {weights_not_loaded}" )