452 lines
18 KiB
Python
452 lines
18 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# Adapted from vllm
|
|
# Copyright 2023 The vLLM Authors.
|
|
# Copyright 2025 The FastVideo Authors.
|
|
|
|
# ruff: noqa: SIM117
|
|
import collections
|
|
import copy
|
|
import dataclasses
|
|
import fnmatch
|
|
import glob
|
|
import inspect
|
|
import itertools
|
|
import math
|
|
import os
|
|
import time
|
|
import warnings
|
|
from abc import ABC, abstractmethod
|
|
from contextlib import contextmanager
|
|
from typing import (Any, Callable, Dict, Generator, Iterable, List, Optional,
|
|
Tuple, cast)
|
|
|
|
import gguf
|
|
import huggingface_hub
|
|
import numpy as np
|
|
import torch
|
|
from huggingface_hub import HfApi
|
|
from torch import nn
|
|
from transformers import AutoModelForCausalLM
|
|
from transformers.utils import SAFE_WEIGHTS_INDEX_NAME
|
|
|
|
from vllm.attention import Attention
|
|
from vllm.config import (LoadConfig, LoadFormat, ModelConfig, ParallelConfig,
|
|
VllmConfig, set_current_vllm_config)
|
|
from vllm.distributed import (get_tensor_model_parallel_rank,
|
|
get_tensor_model_parallel_world_size)
|
|
from vllm.envs import VLLM_USE_MODELSCOPE
|
|
from vllm.logger import init_logger
|
|
from vllm.model_executor.layers.linear import (LinearBase,
|
|
MergedColumnParallelLinear,
|
|
QKVParallelLinear,
|
|
ReplicatedLinear,
|
|
RowParallelLinear)
|
|
from vllm.model_executor.layers.quantization.base_config import (
|
|
QuantizeMethodBase)
|
|
from vllm.model_executor.model_loader.tensorizer import (
|
|
TensorizerConfig, is_vllm_tensorized, load_with_tensorizer,
|
|
serialize_vllm_model, tensorizer_weights_iterator)
|
|
from vllm.model_executor.model_loader.utils import (ParamMapping,
|
|
configure_quant_config,
|
|
get_model_architecture,
|
|
set_default_torch_dtype)
|
|
from vllm.model_executor.model_loader.weight_utils import (
|
|
download_safetensors_index_file_from_hf, download_weights_from_hf,
|
|
filter_duplicate_safetensors_files, filter_files_not_needed_for_inference,
|
|
get_gguf_extra_tensor_names, get_lock, gguf_quant_weights_iterator,
|
|
initialize_dummy_weights, np_cache_weights_iterator, pt_weights_iterator,
|
|
runai_safetensors_weights_iterator, safetensors_weights_iterator)
|
|
from vllm.model_executor.utils import set_weight_attrs
|
|
from vllm.platforms import current_platform
|
|
from vllm.transformers_utils.s3_utils import glob as s3_glob
|
|
from vllm.transformers_utils.utils import is_s3
|
|
from vllm.utils import is_pin_memory_available
|
|
|
|
|
|
@contextmanager
|
|
def device_loading_context(module: torch.nn.Module,
|
|
target_device: torch.device):
|
|
if target_device.type == "cpu":
|
|
# If target is CPU, no need to move anything
|
|
yield module
|
|
return
|
|
|
|
original_device_states: Dict[str, torch.device] = {}
|
|
|
|
# Store original device states and move parameters to GPU if they're on CPU
|
|
for name, p in module.named_parameters():
|
|
if p.device.type == "cpu":
|
|
original_device_states[name] = p.device
|
|
p.data = p.data.to(target_device)
|
|
# Parameters already on target device are not touched
|
|
|
|
try:
|
|
yield module
|
|
|
|
finally:
|
|
# Restore parameters to their original devices, ignoring new parameters
|
|
pin_memory = is_pin_memory_available()
|
|
for name, p in module.named_parameters():
|
|
if name in original_device_states:
|
|
original_device: torch.device = original_device_states[name]
|
|
if original_device.type == "cpu":
|
|
# `torch.empty_like` does not support `pin_memory` argument
|
|
cpu_data = torch.empty_strided(
|
|
size=p.data.size(),
|
|
stride=p.data.stride(),
|
|
dtype=p.data.dtype,
|
|
layout=p.data.layout,
|
|
device="cpu",
|
|
pin_memory=pin_memory,
|
|
)
|
|
cpu_data.copy_(p.data)
|
|
p.data = cpu_data
|
|
else:
|
|
p.data = p.data.to(original_device)
|
|
# New parameters or parameters already on target device are untouched
|
|
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
def _initialize_model(
|
|
vllm_config: VllmConfig,
|
|
*,
|
|
prefix: str = "",
|
|
) -> nn.Module:
|
|
"""Initialize a model with the given configurations."""
|
|
model_config = vllm_config.model_config
|
|
model_class, _ = get_model_architecture(model_config)
|
|
|
|
if vllm_config.quant_config is not None:
|
|
configure_quant_config(vllm_config.quant_config, model_class)
|
|
|
|
signatures = inspect.signature(model_class.__init__)
|
|
all_params = [param.name for param in signatures.parameters.values()]
|
|
if "vllm_config" in all_params and "prefix" in all_params:
|
|
# new-style model class
|
|
with set_current_vllm_config(vllm_config, check_compile=True):
|
|
return model_class(vllm_config=vllm_config, prefix=prefix)
|
|
|
|
msg = ("vLLM model class should accept `vllm_config` and `prefix` as "
|
|
"input arguments. Possibly you have an old-style model class"
|
|
" registered from out of tree and it is used for new vLLM version. "
|
|
"Check https://docs.vllm.ai/en/latest/design/arch_overview.html "
|
|
"for the design and update the model class accordingly.")
|
|
warnings.warn(msg, DeprecationWarning, stacklevel=2)
|
|
|
|
logger.warning(
|
|
"Trying to guess the arguments for old-style model class %s",
|
|
model_class,
|
|
)
|
|
# try to be compatible with old-style model class
|
|
kwargs = {}
|
|
if "prefix" in all_params:
|
|
kwargs["prefix"] = prefix
|
|
if "config" in all_params:
|
|
kwargs["config"] = model_config.hf_config
|
|
if "cache_config" in all_params:
|
|
kwargs["cache_config"] = vllm_config.cache_config
|
|
if "quant_config" in all_params:
|
|
kwargs["quant_config"] = vllm_config.quant_config
|
|
if "lora_config" in all_params:
|
|
kwargs["lora_config"] = vllm_config.lora_config
|
|
if "scheduler_config" in all_params:
|
|
kwargs["scheduler_config"] = vllm_config.scheduler_config
|
|
with set_current_vllm_config(vllm_config, check_compile=True):
|
|
return model_class(**kwargs)
|
|
|
|
|
|
def _process_weights_after_loading(model: nn.Module, model_config: ModelConfig,
|
|
target_device: torch.device) -> None:
|
|
for _, module in model.named_modules():
|
|
quant_method = getattr(module, "quant_method", None)
|
|
if isinstance(quant_method, QuantizeMethodBase):
|
|
# When quant methods need to process weights after loading
|
|
# (for repacking, quantizing, etc), they expect parameters
|
|
# to be on the global target device. This scope is for the
|
|
# case where cpu offloading is used, where we will move the
|
|
# parameters onto device for processing and back off after.
|
|
with device_loading_context(module, target_device):
|
|
quant_method.process_weights_after_loading(module)
|
|
|
|
# Currently only used by MLA.
|
|
# NOTE: This intentionally happens after other modules so we can easily
|
|
# decompress the weights for MLA.
|
|
for _, module in model.named_modules():
|
|
if isinstance(module, Attention) and \
|
|
hasattr(module, "process_weights_after_loading"):
|
|
# TODO(lucas): see if there is a way to unify the signatures
|
|
# of process_weights_after_loading
|
|
module.process_weights_after_loading(model_config.dtype)
|
|
|
|
|
|
class BaseModelLoader(ABC):
|
|
"""Base class for model loaders."""
|
|
|
|
def __init__(self, load_config: LoadConfig):
|
|
self.load_config = load_config
|
|
|
|
@abstractmethod
|
|
def download_model(self, model_config: ModelConfig) -> None:
|
|
"""Download a model so that it can be immediately loaded."""
|
|
raise NotImplementedError
|
|
|
|
@abstractmethod
|
|
def load_model(self, *, vllm_config: VllmConfig) -> nn.Module:
|
|
"""Load a model with the given configurations."""
|
|
raise NotImplementedError
|
|
|
|
|
|
class DefaultModelLoader(BaseModelLoader):
|
|
"""Model loader that can load different file types from disk."""
|
|
|
|
@dataclasses.dataclass
|
|
class Source:
|
|
"""A source for weights."""
|
|
|
|
model_or_path: str
|
|
"""The model ID or path."""
|
|
|
|
revision: Optional[str]
|
|
"""The optional model revision."""
|
|
|
|
prefix: str = ""
|
|
"""A prefix to prepend to all weights."""
|
|
|
|
fall_back_to_pt: bool = True
|
|
"""Whether .pt weights can be used."""
|
|
|
|
allow_patterns_overrides: Optional[list[str]] = 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)
|
|
if load_config.model_loader_extra_config:
|
|
raise ValueError(f"Model loader extra config is not supported for "
|
|
f"load format {load_config.load_format}")
|
|
|
|
def _maybe_download_from_modelscope(
|
|
self, model: str, revision: Optional[str]) -> Optional[str]:
|
|
"""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 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
|
|
|
|
if not os.path.exists(model):
|
|
# Use file lock to prevent multiple processes from
|
|
# downloading the same model weights at the same time.
|
|
with get_lock(model, self.load_config.download_dir):
|
|
model_path = snapshot_download(
|
|
model_id=model,
|
|
cache_dir=self.load_config.download_dir,
|
|
local_files_only=huggingface_hub.constants.
|
|
HF_HUB_OFFLINE,
|
|
revision=revision,
|
|
ignore_file_pattern=self.load_config.ignore_patterns,
|
|
)
|
|
else:
|
|
model_path = model
|
|
return model_path
|
|
return None
|
|
|
|
def _prepare_weights(
|
|
self,
|
|
model_name_or_path: str,
|
|
revision: Optional[str],
|
|
fall_back_to_pt: bool,
|
|
allow_patterns_overrides: Optional[list[str]],
|
|
) -> Tuple[str, List[str], bool]:
|
|
"""Prepare weights for the model.
|
|
|
|
If the model is not local, it will be downloaded."""
|
|
model_name_or_path = (self._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
|
|
# Some quantized models use .pt files for storing the weights.
|
|
if load_format == LoadFormat.AUTO:
|
|
allow_patterns = ["*.safetensors", "*.bin"]
|
|
elif load_format == LoadFormat.SAFETENSORS:
|
|
use_safetensors = True
|
|
allow_patterns = ["*.safetensors"]
|
|
elif load_format == LoadFormat.MISTRAL:
|
|
use_safetensors = True
|
|
allow_patterns = ["consolidated*.safetensors"]
|
|
index_file = "consolidated.safetensors.index.json"
|
|
elif load_format == LoadFormat.PT:
|
|
allow_patterns = ["*.pt"]
|
|
elif load_format == LoadFormat.NPCACHE:
|
|
allow_patterns = ["*.bin"]
|
|
else:
|
|
raise ValueError(f"Unknown load_format: {load_format}")
|
|
|
|
if fall_back_to_pt:
|
|
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,
|
|
ignore_patterns=self.load_config.ignore_patterns,
|
|
)
|
|
else:
|
|
hf_folder = model_name_or_path
|
|
|
|
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 == "*.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,
|
|
self.load_config.download_dir,
|
|
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."""
|
|
hf_folder, hf_weights_files, use_safetensors = self._prepare_weights(
|
|
source.model_or_path, source.revision, source.fall_back_to_pt,
|
|
source.allow_patterns_overrides)
|
|
if self.load_config.load_format == LoadFormat.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,
|
|
)
|
|
elif use_safetensors:
|
|
weights_iterator = safetensors_weights_iterator(hf_weights_files)
|
|
else:
|
|
weights_iterator = pt_weights_iterator(hf_weights_files)
|
|
|
|
if current_platform.is_tpu():
|
|
# In PyTorch XLA, we should call `xm.mark_step` frequently so that
|
|
# not too many ops are accumulated in the XLA program.
|
|
import torch_xla.core.xla_model as xm
|
|
|
|
def _xla_weights_iterator(iterator: Generator):
|
|
for weights in iterator:
|
|
yield weights
|
|
xm.mark_step()
|
|
|
|
weights_iterator = _xla_weights_iterator(weights_iterator)
|
|
|
|
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_config.model,
|
|
model_config.revision,
|
|
fall_back_to_pt=True,
|
|
allow_patterns_overrides=None)
|
|
|
|
def load_model(self, vllm_config: VllmConfig) -> nn.Module:
|
|
device_config = vllm_config.device_config
|
|
model_config = vllm_config.model_config
|
|
target_device = torch.device(device_config.device)
|
|
with set_default_torch_dtype(model_config.dtype):
|
|
with target_device:
|
|
model = _initialize_model(vllm_config=vllm_config)
|
|
|
|
weights_to_load = {name for name, _ in model.named_parameters()}
|
|
loaded_weights = model.load_weights(
|
|
self._get_all_weights(model_config, model))
|
|
self.counter_after_loading_weights = time.perf_counter()
|
|
logger.info(
|
|
"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 currently.
|
|
if model_config.quantization is None and loaded_weights is not None:
|
|
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}")
|
|
|
|
_process_weights_after_loading(model, model_config, target_device)
|
|
|
|
return model.eval()
|
|
|
|
|
|
|
|
|
|
|
|
def get_model_loader(load_config: LoadConfig) -> BaseModelLoader:
|
|
"""Get a model loader based on the load format."""
|
|
if isinstance(load_config.load_format, type):
|
|
return load_config.load_format(load_config)
|
|
|
|
return DefaultModelLoader(load_config) |