utils.py 26.4 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

4
import itertools
5
from collections.abc import Iterable, Mapping
6
from dataclasses import dataclass, field
7
from typing import Any, Literal, Optional, Protocol, Union, overload
8

9
import torch
10
import torch.nn as nn
11
from torch.func import functional_call
12
from transformers import PretrainedConfig
13
from typing_extensions import deprecated
14

15
import vllm.envs as envs
16
from vllm.config import VllmConfig
17
18
19
20
from vllm.distributed import (
    get_tensor_model_parallel_rank,
    get_tensor_model_parallel_world_size,
)
21
from vllm.logger import init_logger
22
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
23
from vllm.multimodal import NestedTensors
24
from vllm.sequence import IntermediateTensors
25
26
27
28
29
30
31
from vllm.utils import (
    cdiv,
    direct_register_custom_op,
    get_cuda_view_from_cpu_tensor,
    is_pin_memory_available,
    is_uva_available,
)
32
33

logger = init_logger(__name__)
34

35
36
WeightsMapping = Mapping[str, Optional[str]]
"""If a key maps to a value of `None`, the corresponding weight is ignored."""
37

38

39
40
41
@dataclass
class WeightsMapper:
    """Maps the name of each weight if they match the following patterns."""
42

43
44
45
    orig_to_new_substr: WeightsMapping = field(default_factory=dict)
    orig_to_new_prefix: WeightsMapping = field(default_factory=dict)
    orig_to_new_suffix: WeightsMapping = field(default_factory=dict)
46

47
48
49
50
51
    def _map_name(self, key: str) -> Optional[str]:
        for substr, new_key in self.orig_to_new_substr.items():
            if substr in key:
                if new_key is None:
                    return None
52

53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
                key = key.replace(substr, new_key, 1)

        for prefix, new_key in self.orig_to_new_prefix.items():
            if key.startswith(prefix):
                if new_key is None:
                    return None

                key = key.replace(prefix, new_key, 1)

        for suffix, new_key in self.orig_to_new_suffix.items():
            if key.endswith(suffix):
                if new_key is None:
                    return None

                key = new_key.join(key.rsplit(suffix, 1))

        return key
70

71
    def apply(
72
73
        self, weights: Iterable[tuple[str, torch.Tensor]]
    ) -> Iterable[tuple[str, torch.Tensor]]:
74
75
76
77
78
        return (
            (out_name, data)
            for name, data in weights
            if (out_name := self._map_name(name)) is not None
        )
79

80
81
    def apply_list(self, values: list[str]) -> list[str]:
        return [
82
83
            out_name
            for name in values
84
85
86
87
88
89
90
91
92
93
            if (out_name := self._map_name(name)) is not None
        ]

    def apply_dict(self, values: dict[str, Any]) -> dict[str, Any]:
        return {
            out_name: value
            for name, value in values.items()
            if (out_name := self._map_name(name)) is not None
        }

94
95

class AutoWeightsLoader:
96
    """
97
    Helper class to load weights into a [`torch.nn.Module`][]. It is able
98
99
100
101
102
103
104
105
    to automatically detect child modules and parameters while iterating over
    the weights only once.

    The weight loading logic for individual modules can be overridden
    by defining a ``load_weights`` method.

    Similarly, the weight loading logic for individual parameters can be
    overridden by defining a ``weight_loader`` method.
106
107
108

    Detailed weight loading information can be viewed by setting the
    environment variable ``VLLM_LOGGING_LEVEL=DEBUG``.
109
    """
110

111
112
113
114
115
116
117
118
    # Models trained using early version ColossalAI
    # may include these tensors in checkpoint. Skip them.
    ROTARY_EMBEDS_UNUSED_WEIGHTS = [
        "rotary_emb.inv_freq",
        "rotary_emb.cos_cached",
        "rotary_emb.sin_cached",
    ]

119
120
121
122
    def __init__(
        self,
        module: nn.Module,
        *,
123
        skip_prefixes: Optional[list[str]] = None,
124
        skip_substrs: Optional[list[str]] = None,
125
        ignore_unexpected_prefixes: Optional[list[str]] = None,
126
        ignore_unexpected_suffixes: Optional[list[str]] = None,
127
128
129
130
131
    ) -> None:
        super().__init__()

        self.module = module
        self.skip_prefixes = skip_prefixes or []
132
        self.skip_substrs = skip_substrs or []
133
        self.ignore_unexpected_prefixes = ignore_unexpected_prefixes or []
134
        self.ignore_unexpected_suffixes = ignore_unexpected_suffixes or []
135
136
        # update default skip_substrs
        self.skip_substrs += self.ROTARY_EMBEDS_UNUSED_WEIGHTS
137
138
139

    def _groupby_prefix(
        self,
140
141
        weights: Iterable[tuple[str, torch.Tensor]],
    ) -> Iterable[tuple[str, Iterable[tuple[str, torch.Tensor]]]]:
142
143
144
145
        weights_by_parts = (
            (weight_name.split(".", 1), weight_data)
            for weight_name, weight_data in weights
        )
146

147
        for prefix, group in itertools.groupby(weights_by_parts, key=lambda x: x[0][0]):
148
149
150
151
            yield (
                prefix,
                # Because maxsplit=1 in weight_name.split(...),
                # the length of `parts` must either be 1 or 2
152
153
154
155
                (
                    ("" if len(parts) == 1 else parts[1], weights_data)
                    for parts, weights_data in group
                ),
156
157
158
159
160
161
162
163
164
165
166
            )

    def _get_qualname(self, prefix: str, rest: str) -> str:
        if prefix == "":
            return rest
        if rest == "":
            return prefix

        return ".".join((prefix, rest))

    def _can_skip(self, qualname: str) -> bool:
167
168
169
        return any(qualname.startswith(p) for p in self.skip_prefixes) or any(
            substr in qualname for substr in self.skip_substrs
        )
170
171

    def _can_ignore_unexpected(self, qualname: str) -> bool:
172
173
174
        iup = (qualname.startswith(p) for p in self.ignore_unexpected_prefixes)
        ius = (qualname.endswith(s) for s in self.ignore_unexpected_suffixes)
        return any(iup) or any(ius)
175
176
177
178
179

    def _load_param(
        self,
        base_prefix: str,
        param: nn.Parameter,
180
        weights: Iterable[tuple[str, torch.Tensor]],
181
    ) -> Iterable[str]:
182
183
184
185
        for weight_name, weight_data in weights:
            weight_qualname = self._get_qualname(base_prefix, weight_name)

            if self._can_skip(weight_qualname):
186
187
                logger.debug("Skipping weight %s", weight_qualname)

188
189
190
                continue

            if weight_name != "":
191
192
                if self._can_ignore_unexpected(weight_qualname):
                    logger.debug("Ignoring weight %s", weight_qualname)
193

194
195
196
197
                    continue

                raise ValueError(
                    f"Attempted to load nested weight '{weight_qualname}' "
198
199
                    f"into a single parameter '{base_prefix}'"
                )
200

201
            weight_loader = getattr(param, "weight_loader", default_weight_loader)
202
203
            weight_loader(param, weight_data)

204
            logger.debug("Loaded weight %s with shape %s", weight_qualname, param.shape)
205

206
207
            yield weight_qualname

208
209
210
    def _add_loadable_non_param_tensors(
        self, module: nn.Module, child_params: dict[str, torch.Tensor]
    ):
211
212
213
214
        """
        Add tensor names that are not in the model params that may be in the
        safetensors, e.g., batch normalization stats.
        """
215
216
217
        if isinstance(
            module,
            (
218
219
220
221
222
223
224
                nn.BatchNorm1d,
                nn.BatchNorm2d,
                nn.BatchNorm3d,
                nn.LazyBatchNorm1d,
                nn.LazyBatchNorm2d,
                nn.LazyBatchNorm3d,
                nn.SyncBatchNorm,
225
226
            ),
        ):
227
            module_state_dict = module.state_dict()
228
            for stat_name in ("running_mean", "running_var", "num_batches_tracked"):
229
230
                child_params[stat_name] = module_state_dict[stat_name]

231
232
233
234
    def _load_module(
        self,
        base_prefix: str,
        module: nn.Module,
235
        weights: Iterable[tuple[str, torch.Tensor]],
236
    ) -> Iterable[str]:
237
238
239
240
241
242
243
244
        if isinstance(module, PPMissingLayer):
            return

        # Avoid infinite recursion since this function is typically
        # called inside load_weights of the module itself
        if module != self.module:
            module_load_weights = getattr(module, "load_weights", None)
            if callable(module_load_weights):
245
                loaded_params = module_load_weights(weights)
246
247
                if loaded_params is None:
                    logger.warning(
248
249
                        "Unable to collect loaded parameters for module %s", module
                    )
250
251
252
253
254
                else:
                    yield from map(
                        lambda x: self._get_qualname(base_prefix, x),
                        loaded_params,
                    )
255
256
257
258

        child_modules = dict(module.named_children())
        child_params = dict(module.named_parameters(recurse=False))

259
260
261
262
        # Add missing tensors the weight loader needs to be able to load
        # that aren't registered as params, e.g., batchnorm statistics.
        self._add_loadable_non_param_tensors(module, child_params)

263
264
265
266
        for child_prefix, child_weights in self._groupby_prefix(weights):
            prefix = self._get_qualname(base_prefix, child_prefix)

            if child_prefix in child_modules:
267
268
269
270
271
                if self._can_skip(prefix + "."):
                    logger.debug("Skipping module %s", prefix)

                    continue

272
273
274
                yield from self._load_module(
                    prefix, child_modules[child_prefix], child_weights
                )
275
            elif child_prefix in child_params:
276
277
278
279
280
                if self._can_skip(prefix):
                    logger.debug("Skipping param %s", prefix)

                    continue

281
282
283
                yield from self._load_param(
                    prefix, child_params[child_prefix], child_weights
                )
284
            else:
285
286
287
288
289
290
291
292
293
294
295
296
297
298
                can_skip_module = self._can_skip(prefix + ".")
                can_skip_param = self._can_skip(prefix)
                if can_skip_module or can_skip_param:
                    logger.debug("Skipping missing %s", prefix)

                    continue

                can_ignore_module = self._can_ignore_unexpected(prefix + ".")
                can_ignore_param = self._can_ignore_unexpected(prefix)
                if can_ignore_module or can_ignore_param:
                    logger.debug("Ignoring missing %s", prefix)

                    continue

299
300
301
302
                msg = (
                    f"There is no module or parameter named '{prefix}' "
                    f"in {type(self.module).__name__}"
                )
303
                raise ValueError(msg)
304
305
306

    def load_weights(
        self,
307
        weights: Iterable[tuple[str, torch.Tensor]],
308
309
        *,
        mapper: Optional[WeightsMapper] = None,
310
    ) -> set[str]:
311
312
        if mapper is not None:
            weights = mapper.apply(weights)
313
        # filter out weights with first-prefix/substr to skip in name
314
315
316
        weights = (
            (name, weight) for name, weight in weights if not self._can_skip(name)
        )
317

318
        autoloaded_weights = set(self._load_module("", self.module, weights))
319
        return autoloaded_weights
320
321


322
def init_vllm_registered_model(
323
    vllm_config: VllmConfig,
324
    *,
325
    prefix: str = "",
326
327
    hf_config: Optional[PretrainedConfig] = None,
    architectures: Optional[list[str]] = None,
328
329
330
331
332
) -> nn.Module:
    """
    Helper function to initialize an inner model registered to vLLM,
    based on the arguments passed to the outer vLLM model.
    """
333
    from vllm.model_executor.model_loader.utils import initialize_model
334

335
336
337
338
    if hf_config is None and architectures is not None:
        # So that the architectures field is overridden
        hf_config = vllm_config.model_config.hf_config

339
    if hf_config is not None:
340
        vllm_config = vllm_config.with_hf_config(hf_config, architectures=architectures)
341

342
    return initialize_model(vllm_config=vllm_config, prefix=prefix)
343
344


345
@overload
346
def flatten_bn(x: torch.Tensor) -> torch.Tensor: ...
347
348
349


@overload
350
def flatten_bn(x: list[torch.Tensor]) -> list[torch.Tensor]: ...
351
352
353
354


@overload
def flatten_bn(
355
    x: Union[list[torch.Tensor], torch.Tensor],
356
357
    *,
    concat: Literal[True],
358
) -> torch.Tensor: ...
359
360


361
362
@overload
def flatten_bn(
363
    x: Union[list[torch.Tensor], torch.Tensor],
364
365
    *,
    concat: bool = False,
366
) -> Union[list[torch.Tensor], torch.Tensor]: ...
367
368


369
def flatten_bn(
370
    x: Union[list[torch.Tensor], torch.Tensor],
371
372
    *,
    concat: bool = False,
373
) -> Union[list[torch.Tensor], torch.Tensor]:
374
375
376
377
378
379
380
381
382
383
384
385
386
387
    """
    Flatten the ``B`` and ``N`` dimensions of batched multimodal inputs.

    The input tensor should have shape ``(B, N, ...)```.
    """
    if isinstance(x, torch.Tensor):
        return x.flatten(0, 1)

    if concat:
        return torch.cat(x)

    return [x_n for x_b in x for x_n in x_b]


388
389
def _flatten_embeddings(embeddings: NestedTensors) -> torch.Tensor:
    """
390
391
    Recursively flattens and concatenates NestedTensors on all but the last
    dimension.
392
393
394
    """

    if isinstance(embeddings, torch.Tensor):
395
396
        # Flatten all but the last dimension.
        return embeddings.flatten(0, -2)
397
398
399
400
401
402
403
404
405
406
407
408
409

    return torch.cat(tuple(_flatten_embeddings(t) for t in embeddings))


def _embedding_count_expression(embeddings: NestedTensors) -> str:
    """
    Constructs a debugging representation of the number of embeddings in the
    NestedTensors.
    """

    if isinstance(embeddings, torch.Tensor):
        return " x ".join([str(dim) for dim in embeddings.shape[:-1]])

410
    return " + ".join(_embedding_count_expression(inner) for inner in embeddings)
411
412


Cyrus Leung's avatar
Cyrus Leung committed
413
414
415
def _merge_multimodal_embeddings(
    inputs_embeds: torch.Tensor,
    multimodal_embeddings: NestedTensors,
416
    is_multimodal: torch.Tensor,
Cyrus Leung's avatar
Cyrus Leung committed
417
) -> torch.Tensor:
418
    """
419
420
    Merge ``multimodal_embeddings`` into ``inputs_embeds`` by overwriting the
    positions in ``inputs_embeds`` corresponding to placeholder tokens in
421
    ``input_ids``.
422
423

    Note:
424
        This updates ``inputs_embeds`` in place.
425
    """
426
427
428
429
430
431
    if len(multimodal_embeddings) == 0:
        return inputs_embeds

    mm_embeds_flat = _flatten_embeddings(multimodal_embeddings)
    input_dtype = inputs_embeds.dtype

432
    try:
433
434
435
436
437
        # For debugging
        # inputs_embeds[is_multimodal] = mm_embeds_flat.to(dtype=input_dtype)

        # NOTE: This can avoid D2H sync (#22105), but fails to
        # raise an error if is_multimodal.sum() < len(mm_embeds_flat)
438
439
440
        inputs_embeds.masked_scatter_(
            is_multimodal.unsqueeze(-1), mm_embeds_flat.to(dtype=input_dtype)
        )
441
    except RuntimeError as e:
442
        num_actual_tokens = len(mm_embeds_flat)
443
444
        num_expected_tokens = is_multimodal.sum().item()

445
        if num_actual_tokens != num_expected_tokens:
446
            expr = _embedding_count_expression(multimodal_embeddings)
447

448
            raise ValueError(
449
                f"Attempted to assign {expr} = {num_actual_tokens} "
450
451
                f"multimodal tokens to {num_expected_tokens} placeholders"
            ) from e
Cyrus Leung's avatar
Cyrus Leung committed
452

453
        raise ValueError("Error during masked scatter operation") from e
Cyrus Leung's avatar
Cyrus Leung committed
454

455
    return inputs_embeds
Cyrus Leung's avatar
Cyrus Leung committed
456
457


458
459
460
461
462
@deprecated(
    "`merge_multimodal_embeddings` has been replaced with "
    "`SupportsMultiModal.get_input_embeddings` and will be "
    "removed in v0.12."
)
Cyrus Leung's avatar
Cyrus Leung committed
463
464
465
466
def merge_multimodal_embeddings(
    input_ids: torch.Tensor,
    inputs_embeds: torch.Tensor,
    multimodal_embeddings: NestedTensors,
467
    placeholder_token_id: Union[int, list[int]],
Cyrus Leung's avatar
Cyrus Leung committed
468
469
470
471
472
) -> torch.Tensor:
    """
    Merge ``multimodal_embeddings`` into ``inputs_embeds`` by overwriting the
    positions in ``inputs_embeds`` corresponding to placeholder tokens in
    ``input_ids``.
473
474
475
476
477

    ``placeholder_token_id`` can be a list of token ids (e.g, token ids
    of img_start, img_break, and img_end tokens) when needed: This means
    the order of these tokens in the ``input_ids`` MUST MATCH the order of
    their embeddings in ``multimodal_embeddings`` since we need to
478
479
480
481
482
483
484
485
    slice-merge instead of individually scattering.

    For example, if input_ids is "TTTTTSIIIBIIIBIIIETTT", where
    - T is text token
    - S is image start token
    - I is image embedding token
    - B is image break token
    - E is image end token.
486
487
488

    Then the image embeddings (that correspond to I's) from vision encoder
    must be padded with embeddings of S, B, and E in the same order of
489
    input_ids for a correct embedding merge.
Cyrus Leung's avatar
Cyrus Leung committed
490
491
492
493

    Note:
        This updates ``inputs_embeds`` in place.
    """
494
    if isinstance(placeholder_token_id, list):
495
496
        is_multimodal = isin_list(input_ids, placeholder_token_id)
    else:
497
        is_multimodal = input_ids == placeholder_token_id
498

Cyrus Leung's avatar
Cyrus Leung committed
499
500
    return _merge_multimodal_embeddings(
        inputs_embeds,
501
502
        multimodal_embeddings=multimodal_embeddings,
        is_multimodal=is_multimodal,
Cyrus Leung's avatar
Cyrus Leung committed
503
504
505
    )


506
507
508
509
510
511
512
513
514
515
516
517
def isin_list(
    elements: torch.Tensor,
    test_elements_list: list[int],
) -> torch.Tensor:
    test_elements = torch.tensor(
        test_elements_list,
        pin_memory=is_pin_memory_available(),
    ).to(device=elements.device, non_blocking=True)

    return torch.isin(elements, test_elements)


518
class LayerFn(Protocol):
519
    def __call__(self, prefix: str) -> torch.nn.Module: ...
520
521


522
523
524
525
526
527
528
class PPMissingLayer(torch.nn.Identity):
    """
    A placeholder layer for missing layers in a pipeline parallel model.
    """

    def __init__(self, *args, **kwargs):
        super().__init__()
529
530

    def forward(self, *args, **kwargs):
531
532
        """Return the first arg from args or the first value from kwargs."""
        return args[0] if args else next(iter(kwargs.values()))
533
534


535
536
537
538
539
540
541
542
543
544
545
_CPU_OFFLOAD_BYTES = 0
_CPU_OFFLOAD_MAX_BYTES = 0


def set_cpu_offload_max_bytes(max_bytes: int) -> None:
    global _CPU_OFFLOAD_MAX_BYTES, _CPU_OFFLOAD_BYTES
    _CPU_OFFLOAD_BYTES = 0
    _CPU_OFFLOAD_MAX_BYTES = max_bytes


def maybe_offload_to_cpu(module: torch.nn.Module) -> torch.nn.Module:
546
547
548
549
    if (params := next(module.parameters(), None)) is None:
        return module

    device = params.device
550
551
552
553
554
555
556
557
558

    if device == torch.device("cpu"):
        return module

    global _CPU_OFFLOAD_MAX_BYTES, _CPU_OFFLOAD_BYTES
    if _CPU_OFFLOAD_BYTES >= _CPU_OFFLOAD_MAX_BYTES:
        return module

    pin_memory = is_pin_memory_available()
559
560
561
    uva_available = is_uva_available()

    if envs.VLLM_USE_V1:
562
        assert uva_available, "V1 CPU offloading requires uva (pin memory) support"
563
564
565
        uva_offloading = True
    else:
        uva_offloading = False
566
567
568

    # offload parameters to CPU
    # use pin_memory if possible, which helps cudagraph capture speed
569
    offloaded_parameters = False
570
571
572
573
574
575
576
    for p in module.parameters():
        if _CPU_OFFLOAD_BYTES >= _CPU_OFFLOAD_MAX_BYTES:
            # we use per-parameter offloading
            # one module might have some parameters offloaded and some not
            break

        # `torch.empty_like` does not support `pin_memory` argument
577
578
579
580
581
582
583
584
        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,
        )
585
        cpu_data.copy_(p.data)
586
587
588
589
590
591
        if not uva_offloading:
            p.data = cpu_data
        else:
            # keep the cpu data alive
            p._vllm_offloaded_cpu_data = cpu_data
            p.data = get_cuda_view_from_cpu_tensor(cpu_data)
592
        _CPU_OFFLOAD_BYTES += p.data.numel() * p.data.element_size()
593
594
        offloaded_parameters = True

595
    if offloaded_parameters and not uva_offloading:
596
597
598
599
600
601
602
603
604
605
        original_forward = module.forward

        def forward(*args, **kwargs):
            module.forward = original_forward
            device_state = {
                # here we blindly call `to(device)`
                # if the parameter is already on the device, it will be a no-op
                k: v.to(device, non_blocking=True)
                for k, v in module.state_dict().items()
            }
606
            output = functional_call(module, device_state, args=args, kwargs=kwargs)
607
608
            module.forward = forward
            return output
609
610
611
612
613
614

        module.forward = forward

    return module


615
def make_layers(
616
617
618
    num_hidden_layers: int,
    layer_fn: LayerFn,
    prefix: str,
619
) -> tuple[int, int, torch.nn.ModuleList]:
620
621
622
623
624
    """Make a list of layers with the given layer function, taking
    pipeline parallelism into account.
    """
    from vllm.distributed.parallel_state import get_pp_group
    from vllm.distributed.utils import get_pp_indices
625
626
627
628

    start_layer, end_layer = get_pp_indices(
        num_hidden_layers, get_pp_group().rank_in_group, get_pp_group().world_size
    )
629
    modules = torch.nn.ModuleList(
630
631
        [PPMissingLayer() for _ in range(start_layer)]
        + [
632
633
            maybe_offload_to_cpu(layer_fn(prefix=f"{prefix}.{idx}"))
            for idx in range(start_layer, end_layer)
634
635
636
        ]
        + [PPMissingLayer() for _ in range(end_layer, num_hidden_layers)]
    )
637
638
639
640
    return start_layer, end_layer, modules


# NOTE: don't use lru_cache here because it can prevent garbage collection
641
_model_to_pp_missing_layer_names: dict[int, list[str]] = {}
642
643


644
def get_pp_missing_layer_names(model: torch.nn.Module) -> list[str]:
645
646
647
648
649
650
651
652
    """Get the names of the missing layers in a pipeline parallel model."""
    model_id = id(model)
    if model_id in _model_to_pp_missing_layer_names:
        return _model_to_pp_missing_layer_names[model_id]

    missing_layer_names = []
    for name, module in model.named_modules():
        if isinstance(module, PPMissingLayer):
653
654
655
            # NOTE: the trailing dot is used to match the prefix of the layer.
            # without the dot, we could match a layer that is not missing,
            # e.g., 'encoder.layer.1' would match 'encoder.layer.11'
656
            missing_layer_names.append(name + ".")
657
658
659
660
661
662
663
    _model_to_pp_missing_layer_names[model_id] = missing_layer_names

    return missing_layer_names


def is_pp_missing_parameter(name: str, model: torch.nn.Module) -> bool:
    """Check if a parameter is missing in a pipeline parallel model."""
664
665
666
667
668
    if isinstance(model, PPMissingLayer):
        return True

    return any(
        name.startswith(missing_layer_name)
669
670
        for missing_layer_name in get_pp_missing_layer_names(model)
    )
671
672


673
def make_empty_intermediate_tensors_factory(keys: list[str], hidden_size: int):
674
    def make_empty_intermediate_tensors(
675
676
677
678
        batch_size: int,
        dtype: torch.dtype,
        device: torch.device,
    ) -> IntermediateTensors:
679
680
681
682
683
684
        return IntermediateTensors(
            {
                key: torch.zeros((batch_size, hidden_size), dtype=dtype, device=device)
                for key in keys
            }
        )
685
686

    return make_empty_intermediate_tensors
687
688


689
690
691
692
693
694
695
696
697
698
699
def maybe_prefix(prefix: str, name: str) -> str:
    """Add a prefix to a name if the prefix is non-empty.

    Args:
        prefix: The prefix to add. If empty, no prefix will be added.
        name: The name to potentially prefix.

    Returns:
        The string "prefix.name" if prefix was non-empty, otherwise just "name".
    """
    return name if not prefix else f"{prefix}.{name}"
700
701


XuruiYang's avatar
XuruiYang committed
702
def extract_layer_index(layer_name: str, num_attn_module: int = 1) -> int:
703
704
705
706
707
708
    """
    Extract the layer index from the module name.
    Examples:
    - "encoder.layers.0" -> 0
    - "encoder.layers.1.self_attn" -> 1
    - "2.self_attn" -> 2
XuruiYang's avatar
XuruiYang committed
709
    - "model.encoder.layers.0.sub.1" -> ValueError if num_attn_module == 1
710
711
    """
    subnames = layer_name.split(".")
712
    int_vals: list[int] = []
713
714
715
716
717
    for subname in subnames:
        try:
            int_vals.append(int(subname))
        except ValueError:
            continue
XuruiYang's avatar
XuruiYang committed
718
    if num_attn_module == 1 or "attn" not in layer_name:
719
720
721
        assert len(int_vals) == 1, (
            f"layer name {layer_name} should only contain one integer"
        )
XuruiYang's avatar
XuruiYang committed
722
723
724

        return int_vals[0]
    else:
725
726
727
728
729
730
731
732
        assert len(int_vals) <= 2, (
            f"layer name {layer_name} should contain most two integers"
        )
        layer_index = (
            int_vals[0] * num_attn_module + int_vals[1]
            if len(int_vals) == 2
            else int_vals[0]
        )
XuruiYang's avatar
XuruiYang committed
733
        return layer_index
734
735
736
737
738
739
740
741
742


def cast_overflow_tensors(
    tensors: torch.Tensor,
    offset: float = 1000,
) -> torch.Tensor:
    if tensors.isinf().any() or tensors.isnan().any():
        clamp_value = torch.finfo(tensors.dtype).max - offset
        tensors = torch.clamp(tensors, min=-clamp_value, max=clamp_value)
743
    return tensors
744
745


746
747
748
def fast_topk(
    values: torch.Tensor, topk: int, dim: int
) -> tuple[torch.Tensor, torch.Tensor]:
749
750
    """
    Optimized topk implementation that uses torch.max for k=1 case.
751

752
753
    This function provides better performance for the common case of k=1
    by using torch.max instead of the more general torch.topk.
754

755
756
757
758
    Args:
        values: Input tensor to find top-k values from
        topk: Number of top values to return (k). Must be > 0.
        dim: Dimension along which to compute topk
759

760
761
762
763
    Returns:
        Tuple of (values, indices) where values are the top-k values
        and indices are their corresponding indices in the input tensor
    """
764
765
766
767
768
769
    if topk == 1:
        # Use max along the specified dimension to get both value and index
        return torch.max(values, dim=dim, keepdim=True)
    else:
        # Use topk for efficiency with larger k values
        return torch.topk(values, topk, dim=dim)
770
771
772
773
774
775
776


def get_model_hidden_size(hf_config: PretrainedConfig) -> int:
    if hasattr(hf_config, "hidden_size"):
        return hf_config.hidden_size
    text_config = hf_config.get_text_config()
    return text_config.hidden_size
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817


# Chunk x along the num_tokens axis for sequence parallelism
# NOTE: This is wrapped in a torch custom op to work around the following issue:
# The output tensor can have a sequence length 0 at small input sequence lengths
# even though we explicitly pad to avoid this.
def sequence_parallel_chunk(x: torch.Tensor) -> torch.Tensor:
    return torch.ops.vllm.sequence_parallel_chunk_impl(x)


def sequence_parallel_chunk_impl(x: torch.Tensor) -> torch.Tensor:
    tp_size = get_tensor_model_parallel_world_size()
    tp_rank = get_tensor_model_parallel_rank()

    # all_gather needs the sequence length to be divisible by tp_size
    seq_len = x.size(0)
    remainder = seq_len % tp_size
    if remainder != 0:
        pad_len = tp_size - remainder
        y = nn.functional.pad(x, (0, 0, 0, pad_len))
    else:
        y = x

    chunk = y.shape[0] // tp_size
    start = tp_rank * chunk
    return torch.narrow(y, 0, start, chunk)


def sequence_parallel_chunk_impl_fake(x: torch.Tensor) -> torch.Tensor:
    tp_size = get_tensor_model_parallel_world_size()
    seq_len = cdiv(x.size(0), tp_size)
    shape = list(x.shape)
    shape[0] = seq_len
    out = torch.empty(shape, dtype=x.dtype, device=x.device)
    return out


direct_register_custom_op(
    op_name="sequence_parallel_chunk_impl",
    op_func=sequence_parallel_chunk_impl,
    fake_impl=sequence_parallel_chunk_impl_fake,
818
    tags=(torch.Tag.needs_fixed_stride_order,),
819
)