utils.py 11.6 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
"""Utilities for selecting and loading models."""
import contextlib
5
6
7
import inspect
import warnings
from contextlib import contextmanager
8
from dataclasses import dataclass, field
9
from typing import Optional
10
11
12

import torch
from torch import nn
13
from typing_extensions import assert_never
14

15
from vllm.attention import Attention
16
from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config
17
from vllm.logger import init_logger
18
from vllm.model_executor.layers.linear import QKVCrossParallelLinear
19
from vllm.model_executor.layers.quantization.base_config import (
20
    QuantizationConfig, QuantizeMethodBase)
21
22
23
24
25
from vllm.model_executor.models.adapters import (
    as_embedding_model, as_reward_model, as_seq_cls_model,
    try_create_mm_pooling_model_cls)
from vllm.model_executor.models.interfaces import (SupportsQuant,
                                                   supports_multimodal)
26
from vllm.utils import is_pin_memory_available
27

28
29
logger = init_logger(__name__)

30
31
32
33
34
35
36
37
38
39

@contextlib.contextmanager
def set_default_torch_dtype(dtype: torch.dtype):
    """Sets the default torch dtype to the given dtype."""
    old_dtype = torch.get_default_dtype()
    torch.set_default_dtype(dtype)
    yield
    torch.set_default_dtype(old_dtype)


40
41
42
43
44
def initialize_model(
    vllm_config: VllmConfig,
    *,
    prefix: str = "",
    model_class: Optional[type[nn.Module]] = None,
45
    model_config: Optional[ModelConfig] = None,
46
47
) -> nn.Module:
    """Initialize a model with the given configurations."""
48
49
    if model_config is None:
        model_config = vllm_config.model_config
50
51
52
53
54
55
56
57
58
59
    if model_class is None:
        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
60
61
62
        with set_current_vllm_config(vllm_config,
                                     check_compile=True,
                                     prefix=prefix):
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
            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
90
91
92
    with set_current_vllm_config(vllm_config,
                                 check_compile=True,
                                 prefix=prefix):
93
94
95
96
97
        return model_class(**kwargs)


def process_weights_after_loading(model: nn.Module, model_config: ModelConfig,
                                  target_device: torch.device) -> None:
98
99
100
101
102
103
104

    # to avoid circular dependency
    from vllm.model_executor.model_loader.online_quantization import (
        maybe_save_metadata_and_attributes_for_weight_reloading)
    maybe_save_metadata_and_attributes_for_weight_reloading(
        model, model_config)

105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
    for _, module in model.named_modules():
        if isinstance(module, QKVCrossParallelLinear):
            # NOTE(Isotr0py): special case for cross QKV layer because
            # q and kv proj aren't registered as submodules intentionally
            module.process_weights_after_loading()
            continue
        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)


@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

140
    original_device_states: dict[str, torch.device] = {}
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174

    # 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


175
_MODEL_ARCH_BY_HASH = dict[int, tuple[type[nn.Module], str]]()
176
177
178
179
"""Caches the outputs of `_get_model_architecture`."""


def _get_model_architecture(
180
        model_config: ModelConfig) -> tuple[type[nn.Module], str]:
181
    architectures = getattr(model_config.hf_config, "architectures", [])
182

183
184
185
186
187
188
    model_cls, arch = model_config.registry.resolve_model_cls(
        architectures,
        model_config=model_config,
    )

    if arch == model_config._get_transformers_backend_cls():
189
190
        assert model_config.model_impl != "vllm"
        if model_config.model_impl == "auto":
191
192
193
194
195
196
            logger.warning_once(
                "%s has no vLLM implementation, falling back to Transformers "
                "implementation. Some features may not be supported and "
                "performance may not be optimal.", arch)

    convert_type = model_config.convert_type
197
198
199
200
201
202
203
204
205
    if convert_type != "none" and supports_multimodal(model_cls):
        logger.debug_once("Detected conversion of Multi Modal model.")
        converted = try_create_mm_pooling_model_cls(model_cls)
        if converted is not None:
            logger.debug_once("Creating wrapper class to forward pooler.")
            return converted, arch
        else:
            logger.debug_once("Attempting direct conversion.")

206
207
208
209
    if convert_type == "none":
        pass
    elif convert_type == "embed":
        logger.debug_once("Converting to embedding model.")
210
        model_cls = as_embedding_model(model_cls)
211
212
    elif convert_type == "classify":
        logger.debug_once("Converting to sequence classification model.")
213
        model_cls = as_seq_cls_model(model_cls)
214
215
    elif convert_type == "reward":
        logger.debug_once("Converting to reward model.")
216
        model_cls = as_reward_model(model_cls)
217
218
    else:
        assert_never(convert_type)
219
220

    return model_cls, arch
221
222


223
224
def get_model_architecture(
        model_config: ModelConfig) -> tuple[type[nn.Module], str]:
225
226
227
228
229
230
231
232
    key = hash((
        model_config.model,
        model_config.convert_type,
        model_config.runner_type,
        model_config.trust_remote_code,
        model_config.model_impl,
        tuple(getattr(model_config.hf_config, "architectures", [])),
    ))
233
234
235
236
237
238
239
240
    if key in _MODEL_ARCH_BY_HASH:
        return _MODEL_ARCH_BY_HASH[key]

    model_arch = _get_model_architecture(model_config)
    _MODEL_ARCH_BY_HASH[key] = model_arch
    return model_arch


241
242
243
244
def get_model_cls(model_config: ModelConfig) -> type[nn.Module]:
    return get_model_architecture(model_config)[0]


245
246
def get_architecture_class_name(model_config: ModelConfig) -> str:
    return get_model_architecture(model_config)[1]
247
248
249
250
251
252


@dataclass
class ParamMapping:
    """
    A class to handle parameter mapping for model weight loading.
253
    It creates a bidirectional mapping between packed parameters and their
254
255
    constituent parts.
    """
256
257
    packed_mapping: dict[str, list[str]]
    inverse_packed_mapping: dict[str, tuple[str,
258
259
260
261
262
263
264
265
266
267
268
269
                                            int]] = field(default_factory=dict)

    def __post_init__(self):
        for packed_name, sub_params in self.packed_mapping.items():
            # Skip self-contained cases (e.g., {"W_pack": ["W_pack"]})
            if len(sub_params) == 1 and sub_params[0] == packed_name:
                continue
            for index, param_name in enumerate(sub_params):
                self.inverse_packed_mapping[param_name] = (
                    packed_name,
                    index,
                )
270
271

    def get_sub_modules(self,
272
                        module_name: str) -> Optional[tuple[str, list[str]]]:
273
274
275
276
        for key, value in self.packed_mapping.items():
            if module_name.endswith(key):
                return key, value
        return None
277
278
279


def configure_quant_config(quant_config: QuantizationConfig,
280
                           model_class: type[nn.Module]):
281
282
283
284
285
286
    """
    Pass packed_modules_mapping by reference to quant_config so that
    quant_config can properly match fused modules

    Note that model attributes are passed by reference to quant_config,
    enabling them to be updated by model_class.__new__ (ex. chatglm, qwen)
287
288
289

    Once the `SupportsQuant` mixin has been added to all models, this
    function can be removed
290
    """
291
292
293
294
295
296
297
298
299
    if not issubclass(model_class, SupportsQuant):
        hf_to_vllm_mapper = getattr(model_class, "hf_to_vllm_mapper", None)
        packed_mapping = getattr(model_class, "packed_modules_mapping", None)

        # pass mappings by reference to quant_config
        if hf_to_vllm_mapper is not None:
            quant_config.apply_vllm_mapper(hf_to_vllm_mapper)
        if packed_mapping is not None:
            quant_config.packed_modules_mapping = packed_mapping