utils.py 27.3 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, Callable, 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

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

logger = init_logger(__name__)
27

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

31

32
33
34
@dataclass
class WeightsMapper:
    """Maps the name of each weight if they match the following patterns."""
35

36
37
38
    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)
39

40
41
42
43
44
    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
45

46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
                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
63

64
    def apply(
65
66
        self, weights: Iterable[tuple[str, torch.Tensor]]
    ) -> Iterable[tuple[str, torch.Tensor]]:
67
68
        return ((out_name, data) for name, data in weights
                if (out_name := self._map_name(name)) is not None)
69

70
71
72
73
74
75
76
77
78
79
80
81
82
    def apply_list(self, values: list[str]) -> list[str]:
        return [
            out_name for name in values
            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
        }

83
84

class AutoWeightsLoader:
85
    """
86
    Helper class to load weights into a [`torch.nn.Module`][]. It is able
87
88
89
90
91
92
93
94
    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.
95
96
97

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

100
101
102
103
104
105
106
107
    # 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",
    ]

108
109
110
111
    def __init__(
        self,
        module: nn.Module,
        *,
112
        skip_prefixes: Optional[list[str]] = None,
113
        skip_substrs: Optional[list[str]] = None,
114
        ignore_unexpected_prefixes: Optional[list[str]] = None,
115
116
117
118
119
    ) -> None:
        super().__init__()

        self.module = module
        self.skip_prefixes = skip_prefixes or []
120
        self.skip_substrs = skip_substrs or []
121
        self.ignore_unexpected_prefixes = ignore_unexpected_prefixes or []
122
123
        # update default skip_substrs
        self.skip_substrs += self.ROTARY_EMBEDS_UNUSED_WEIGHTS
124
125
126

    def _groupby_prefix(
        self,
127
128
        weights: Iterable[tuple[str, torch.Tensor]],
    ) -> Iterable[tuple[str, Iterable[tuple[str, torch.Tensor]]]]:
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
        weights_by_parts = ((weight_name.split(".", 1), weight_data)
                            for weight_name, weight_data in weights)

        for prefix, group in itertools.groupby(weights_by_parts,
                                               key=lambda x: x[0][0]):
            yield (
                prefix,
                # Because maxsplit=1 in weight_name.split(...),
                # the length of `parts` must either be 1 or 2
                (("" if len(parts) == 1 else parts[1], weights_data)
                 for parts, weights_data in group),
            )

    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:
151
152
        return (any(qualname.startswith(p) for p in self.skip_prefixes)
                or any(substr in qualname for substr in self.skip_substrs))
153
154
155
156
157
158
159
160
161

    def _can_ignore_unexpected(self, qualname: str) -> bool:
        return any(
            qualname.startswith(p) for p in self.ignore_unexpected_prefixes)

    def _load_param(
        self,
        base_prefix: str,
        param: nn.Parameter,
162
        weights: Iterable[tuple[str, torch.Tensor]],
163
    ) -> Iterable[str]:
164
165
166
167
        for weight_name, weight_data in weights:
            weight_qualname = self._get_qualname(base_prefix, weight_name)

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

170
171
172
                continue

            if weight_name != "":
173
174
                if self._can_ignore_unexpected(weight_qualname):
                    logger.debug("Ignoring weight %s", weight_qualname)
175

176
177
178
179
180
                    continue

                raise ValueError(
                    f"Attempted to load nested weight '{weight_qualname}' "
                    f"into a single parameter '{base_prefix}'")
181
182
183
184
185

            weight_loader = getattr(param, "weight_loader",
                                    default_weight_loader)
            weight_loader(param, weight_data)

186
187
188
            logger.debug("Loaded weight %s with shape %s", weight_qualname,
                         param.shape)

189
190
            yield weight_qualname

191
    def _add_loadable_non_param_tensors(self, module: nn.Module,
192
                                        child_params: dict[str, torch.Tensor]):
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
        """
        Add tensor names that are not in the model params that may be in the
        safetensors, e.g., batch normalization stats.
        """
        if isinstance(module, (
                nn.BatchNorm1d,
                nn.BatchNorm2d,
                nn.BatchNorm3d,
                nn.LazyBatchNorm1d,
                nn.LazyBatchNorm2d,
                nn.LazyBatchNorm3d,
                nn.SyncBatchNorm,
        )):
            module_state_dict = module.state_dict()
            for stat_name in ("running_mean", "running_var",
                              "num_batches_tracked"):
                child_params[stat_name] = module_state_dict[stat_name]

211
212
213
214
    def _load_module(
        self,
        base_prefix: str,
        module: nn.Module,
215
        weights: Iterable[tuple[str, torch.Tensor]],
216
    ) -> Iterable[str]:
217
218
219
220
221
222
223
224
        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):
225
                loaded_params = module_load_weights(weights)
226
227
228
229
230
231
232
233
234
                if loaded_params is None:
                    logger.warning(
                        "Unable to collect loaded parameters "
                        "for module %s", module)
                else:
                    yield from map(
                        lambda x: self._get_qualname(base_prefix, x),
                        loaded_params,
                    )
235
236
237
238

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

239
240
241
242
        # 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)

243
244
245
246
        for child_prefix, child_weights in self._groupby_prefix(weights):
            prefix = self._get_qualname(base_prefix, child_prefix)

            if child_prefix in child_modules:
247
248
249
250
251
                if self._can_skip(prefix + "."):
                    logger.debug("Skipping module %s", prefix)

                    continue

252
253
254
                yield from self._load_module(prefix,
                                             child_modules[child_prefix],
                                             child_weights)
255
            elif child_prefix in child_params:
256
257
258
259
260
                if self._can_skip(prefix):
                    logger.debug("Skipping param %s", prefix)

                    continue

261
262
                yield from self._load_param(prefix, child_params[child_prefix],
                                            child_weights)
263
            else:
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
                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

                msg = (f"There is no module or parameter named '{prefix}' "
                       f"in {type(self.module).__name__}")
                raise ValueError(msg)
281
282
283

    def load_weights(
        self,
284
        weights: Iterable[tuple[str, torch.Tensor]],
285
286
        *,
        mapper: Optional[WeightsMapper] = None,
287
    ) -> set[str]:
288
289
        if mapper is not None:
            weights = mapper.apply(weights)
290
291
292
        # filter out weights with first-prefix/substr to skip in name
        weights = ((name, weight) for name, weight in weights
                   if not self._can_skip(name))
293

294
        autoloaded_weights = set(self._load_module("", self.module, weights))
295
        return autoloaded_weights
296
297


298
def init_vllm_registered_model(
299
    vllm_config: VllmConfig,
300
    *,
301
    prefix: str = "",
302
303
    hf_config: Optional[PretrainedConfig] = None,
    architectures: Optional[list[str]] = None,
304
305
306
307
308
) -> nn.Module:
    """
    Helper function to initialize an inner model registered to vLLM,
    based on the arguments passed to the outer vLLM model.
    """
309
    from vllm.model_executor.model_loader.utils import initialize_model
310

311
312
313
314
    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

315
    if hf_config is not None:
316
317
        vllm_config = vllm_config.with_hf_config(hf_config,
                                                 architectures=architectures)
318

319
    return initialize_model(vllm_config=vllm_config, prefix=prefix)
320
321


322
323
324
325
326
327
@overload
def flatten_bn(x: torch.Tensor) -> torch.Tensor:
    ...


@overload
328
def flatten_bn(x: list[torch.Tensor]) -> list[torch.Tensor]:
329
330
331
332
333
    ...


@overload
def flatten_bn(
334
    x: Union[list[torch.Tensor], torch.Tensor],
335
336
337
338
339
340
    *,
    concat: Literal[True],
) -> torch.Tensor:
    ...


341
342
@overload
def flatten_bn(
343
    x: Union[list[torch.Tensor], torch.Tensor],
344
345
    *,
    concat: bool = False,
346
) -> Union[list[torch.Tensor], torch.Tensor]:
347
348
349
    ...


350
def flatten_bn(
351
    x: Union[list[torch.Tensor], torch.Tensor],
352
353
    *,
    concat: bool = False,
354
) -> Union[list[torch.Tensor], torch.Tensor]:
355
356
357
358
359
360
361
362
363
364
365
366
367
368
    """
    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]


369
370
def _flatten_embeddings(embeddings: NestedTensors) -> torch.Tensor:
    """
371
372
    Recursively flattens and concatenates NestedTensors on all but the last
    dimension.
373
374
375
    """

    if isinstance(embeddings, torch.Tensor):
376
377
        # Flatten all but the last dimension.
        return embeddings.flatten(0, -2)
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394

    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]])

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


Cyrus Leung's avatar
Cyrus Leung committed
395
396
397
398
399
def _merge_multimodal_embeddings(
    inputs_embeds: torch.Tensor,
    is_multimodal: torch.Tensor,
    multimodal_embeddings: NestedTensors,
) -> torch.Tensor:
400
    """
401
402
    Merge ``multimodal_embeddings`` into ``inputs_embeds`` by overwriting the
    positions in ``inputs_embeds`` corresponding to placeholder tokens in
403
    ``input_ids``.
404
405

    Note:
406
        This updates ``inputs_embeds`` in place.
407
    """
408
    flattened = _flatten_embeddings(multimodal_embeddings)
409
410
    try:
        # This is equivalent to: inputs_embeds[is_multimodal] = flattened.
411
412
        inputs_embeds.masked_scatter_(is_multimodal.unsqueeze(-1),
                                      flattened.to(dtype=inputs_embeds.dtype))
413
414
415
416
417
418
419
420
421
422
423
424
    except RuntimeError as e:
        num_expected_tokens = is_multimodal.sum().item()
        assert isinstance(num_expected_tokens, int)

        if flattened.shape[0] != num_expected_tokens:
            expr = _embedding_count_expression(multimodal_embeddings)
            raise ValueError(
                f"Attempted to assign {expr} = {flattened.shape[0]} "
                f"multimodal tokens to {num_expected_tokens} placeholders"
            ) from e
        else:
            raise ValueError("Error during masked scatter operation") from e
425
426

    return inputs_embeds
427
428


Cyrus Leung's avatar
Cyrus Leung committed
429
430
431
432
def embed_multimodal(
    input_ids: torch.Tensor,
    multimodal_token_id: int,
    get_text_embeds: Callable[[torch.Tensor], torch.Tensor],
433
    multimodal_embeds: NestedTensors,
Cyrus Leung's avatar
Cyrus Leung committed
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
) -> torch.Tensor:
    """
    Embed token IDs and multimodal inputs and combine their embeddings.

    ``multimodal_token_id`` is used to determine whether a token ID should
    be embedded using ``get_text_embeds`` or ``get_multimodal_embeds``.

    Compared to ``merge_multimodal_embeddings`, this avoids running
    ``get_text_embeds`` on ``input_ids[input_ids == multimodal_token_id]``
    which causes issues when the placeholder token ID exceeds the
    vocabulary size of the language model.
    """
    is_multimodal = input_ids == multimodal_token_id
    is_text = ~is_multimodal

    text_embeds = get_text_embeds(input_ids[is_text])
    merged_embeds = torch.empty(
        (input_ids.shape[0], text_embeds.shape[1]),
        dtype=text_embeds.dtype,
        device=text_embeds.device,
    )

    merged_embeds[is_text] = text_embeds

    return _merge_multimodal_embeddings(
        merged_embeds,
        is_multimodal,
        multimodal_embeds,
    )


def merge_multimodal_embeddings(
    input_ids: torch.Tensor,
    inputs_embeds: torch.Tensor,
    multimodal_embeddings: NestedTensors,
469
    placeholder_token_id: Union[int, list[int]],
Cyrus Leung's avatar
Cyrus Leung committed
470
471
472
473
474
) -> torch.Tensor:
    """
    Merge ``multimodal_embeddings`` into ``inputs_embeds`` by overwriting the
    positions in ``inputs_embeds`` corresponding to placeholder tokens in
    ``input_ids``.
475
476
477
478
479

    ``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
480
481
482
483
484
485
486
487
    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.
488
489
490

    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
491
    input_ids for a correct embedding merge.
Cyrus Leung's avatar
Cyrus Leung committed
492
493
494
495

    Note:
        This updates ``inputs_embeds`` in place.
    """
496
    if isinstance(placeholder_token_id, list):
497
498
499
500
        placeholder_token_id = torch.tensor(
            placeholder_token_id,
            pin_memory=is_pin_memory_available()).to(device=input_ids.device,
                                                     non_blocking=True)
501
502
503
504
505
506
        return _merge_multimodal_embeddings(
            inputs_embeds,
            torch.isin(input_ids, placeholder_token_id),
            multimodal_embeddings,
        )

Cyrus Leung's avatar
Cyrus Leung committed
507
508
509
510
511
512
513
    return _merge_multimodal_embeddings(
        inputs_embeds,
        (input_ids == placeholder_token_id),
        multimodal_embeddings,
    )


514
515
class LayerFn(Protocol):

516
    def __call__(self, prefix: str) -> torch.nn.Module:
517
518
519
        ...


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

    def __init__(self, *args, **kwargs):
        super().__init__()
527
528

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


533
534
535
536
537
538
539
540
541
542
543
_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:
544
545
546
547
    if (params := next(module.parameters(), None)) is None:
        return module

    device = params.device
548
549
550
551
552
553
554
555
556

    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()
557
558
559
560
561
562
563
564
    uva_available = is_uva_available()

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

    # offload parameters to CPU
    # use pin_memory if possible, which helps cudagraph capture speed
568
    offloaded_parameters = False
569
570
571
572
573
574
575
    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
576
577
578
579
580
581
        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)
582
        cpu_data.copy_(p.data)
583
584
585
586
587
588
        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)
589
        _CPU_OFFLOAD_BYTES += p.data.numel() * p.data.element_size()
590
591
        offloaded_parameters = True

592
    if offloaded_parameters and not uva_offloading:
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
        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()
            }
            output = functional_call(module,
                                     device_state,
                                     args=args,
                                     kwargs=kwargs)
            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
625
626
627
628
    """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
    start_layer, end_layer = get_pp_indices(num_hidden_layers,
                                            get_pp_group().rank_in_group,
                                            get_pp_group().world_size)
    modules = torch.nn.ModuleList(
629
        [PPMissingLayer() for _ in range(start_layer)] + [
630
631
            maybe_offload_to_cpu(layer_fn(prefix=f"{prefix}.{idx}"))
            for idx in range(start_layer, end_layer)
632
        ] + [PPMissingLayer() for _ in range(end_layer, num_hidden_layers)])
633
634
635
636
    return start_layer, end_layer, modules


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


640
def get_pp_missing_layer_names(model: torch.nn.Module) -> list[str]:
641
642
643
644
645
646
647
648
    """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):
649
650
651
652
            # 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'
            missing_layer_names.append(name + '.')
653
654
655
656
657
658
659
    _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."""
660
661
662
663
664
665
    if isinstance(model, PPMissingLayer):
        return True

    return any(
        name.startswith(missing_layer_name)
        for missing_layer_name in get_pp_missing_layer_names(model))
666
667


668
def make_empty_intermediate_tensors_factory(keys: list[str], hidden_size: int):
669
670

    def make_empty_intermediate_tensors(
671
672
673
674
        batch_size: int,
        dtype: torch.dtype,
        device: torch.device,
    ) -> IntermediateTensors:
675
        return IntermediateTensors({
676
677
            key:
            torch.zeros((batch_size, hidden_size), dtype=dtype, device=device)
678
679
680
681
            for key in keys
        })

    return make_empty_intermediate_tensors
682
683


684
685
686
687
688
689
690
691
692
693
694
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}"
695
696


XuruiYang's avatar
XuruiYang committed
697
def extract_layer_index(layer_name: str, num_attn_module: int = 1) -> int:
698
699
700
701
702
703
    """
    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
704
    - "model.encoder.layers.0.sub.1" -> ValueError if num_attn_module == 1
705
706
    """
    subnames = layer_name.split(".")
707
    int_vals: list[int] = []
708
709
710
711
712
    for subname in subnames:
        try:
            int_vals.append(int(subname))
        except ValueError:
            continue
XuruiYang's avatar
XuruiYang committed
713
714
715
716
717
718
719
720
721
722
723
    if num_attn_module == 1 or "attn" not in layer_name:
        assert len(int_vals) == 1, (f"layer name {layer_name} should"
                                    " only contain one integer")

        return int_vals[0]
    else:
        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]
        return layer_index
724
725
726
727
728
729
730
731
732


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)
733
    return tensors
734
735


736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
def fast_topk(values: torch.Tensor, topk: int,
              dim: int) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Optimized topk implementation that uses torch.max for k=1 case.
    
    This function provides better performance for the common case of k=1
    by using torch.max instead of the more general torch.topk.
    
    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
        
    Returns:
        Tuple of (values, indices) where values are the top-k values
        and indices are their corresponding indices in the input tensor
    """
753
754
755
756
757
758
    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)
759
760
761
762
763
764
765


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
766
767
768
769
770
771
772
773
774
775
776
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


# 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,
    tags=(torch.Tag.needs_fixed_stride_order, ),
)