| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import torch |
| 5 | import torch.nn as nn |
| 6 | import torch.nn.functional as F |
| 7 | from transformers import PreTrainedConfig |
| 8 | |
| 9 | from vllm.config.lora import LoRAConfig |
| 10 | from vllm.model_executor.custom_op import maybe_get_oot_by_class |
| 11 | from vllm.model_executor.layers.vocab_parallel_embedding import ( |
| 12 | ParallelLMHead, |
| 13 | VocabParallelEmbedding, |
| 14 | ) |
| 15 | from vllm.platforms import current_platform |
| 16 | |
| 17 | from .base import BaseLayerWithLoRA |
| 18 | |
| 19 | |
| 20 | class VocabParallelEmbeddingWithLoRA(BaseLayerWithLoRA): |
| 21 | def __init__(self, base_layer: VocabParallelEmbedding) -> None: |
| 22 | super().__init__() |
| 23 | self.base_layer = base_layer |
| 24 | |
| 25 | def create_lora_weights( |
| 26 | self, |
| 27 | max_loras: int, |
| 28 | lora_config: LoRAConfig, |
| 29 | model_config: PreTrainedConfig | None = None, |
| 30 | ) -> None: |
| 31 | self.lora_a_stacked = torch.zeros( |
| 32 | ( |
| 33 | max_loras, |
| 34 | self.base_layer.org_vocab_size, |
| 35 | lora_config.max_lora_rank, |
| 36 | ), |
| 37 | dtype=lora_config.lora_dtype, |
| 38 | device=self.base_layer.weight.device, |
| 39 | ) |
| 40 | self.lora_b_stacked = torch.zeros( |
| 41 | ( |
| 42 | max_loras, |
| 43 | 1, |
| 44 | self.base_layer.embedding_dim, |
| 45 | lora_config.max_lora_rank, |
| 46 | ), |
| 47 | dtype=lora_config.lora_dtype, |
| 48 | device=self.base_layer.weight.device, |
| 49 | ) |
| 50 | self.lora_a_stacked_2d = self.lora_a_stacked.view( |
| 51 | self.lora_a_stacked.shape[0] * self.lora_a_stacked.shape[1], |
| 52 | self.lora_a_stacked.shape[2], |
| 53 | ) |
| 54 | |
| 55 | def reset_lora(self, index: int): |
| 56 | self.lora_a_stacked[index] = 0 |
| 57 | self.lora_b_stacked[index] = 0 |
| 58 | |
| 59 | def set_lora( |
| 60 | self, |
| 61 | index: int, |
| 62 | lora_a: torch.Tensor | list[torch.Tensor], |
| 63 | lora_b: torch.Tensor | list[torch.Tensor], |
| 64 | ): |
| 65 | assert isinstance(lora_a, torch.Tensor) |
| 66 | assert isinstance(lora_b, torch.Tensor) |
| 67 | self.reset_lora(index) |
| 68 | # NOTE self.lora_a_stacked is row-major, and lora_a is col-major, |
| 69 | # so we need transpose here |
| 70 | |
| 71 | self.lora_a_stacked[index, : lora_a.shape[1], : lora_a.shape[0]].copy_( |
| 72 | lora_a.T, non_blocking=True |
| 73 | ) |
| 74 | self.lora_b_stacked[index, 0, : lora_b.shape[0], : lora_b.shape[1]].copy_( |
| 75 | lora_b, non_blocking=True |
| 76 | ) |
| 77 | |
| 78 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 79 | # NB: Don't use torch.narrow here. torch.narrow triggers some |
| 80 | # Dynamic Shape specialization in torch.compile |
| 81 | num_tokens = x.shape[0] |
| 82 | embeddings_indices = self.punica_wrapper._embeddings_indices[:num_tokens] |
| 83 | |
| 84 | full_lora_a_embeddings = F.embedding( |
| 85 | x + embeddings_indices, |
| 86 | self.lora_a_stacked_2d, |
| 87 | ) |
| 88 | full_output = self.base_layer.forward(x) |
| 89 | |
| 90 | full_output_org = full_output |
| 91 | if full_output.ndim == 3: |
| 92 | full_output = full_output.view( |
| 93 | full_output.shape[0] * full_output.shape[1], -1 |
| 94 | ) |
| 95 | if full_lora_a_embeddings.ndim == 3: |
| 96 | full_lora_a_embeddings = full_lora_a_embeddings.view( |
| 97 | full_lora_a_embeddings.shape[0] * full_lora_a_embeddings.shape[1], |
| 98 | -1, |
| 99 | ) |
| 100 | |
| 101 | lora_output: torch.Tensor | None = self.punica_wrapper.add_lora_embedding( |
| 102 | full_output, full_lora_a_embeddings, self.lora_b_stacked, add_input=True |
| 103 | ) |
| 104 | |
| 105 | if not current_platform.can_update_inplace(): |
| 106 | full_output = lora_output |
| 107 | |
| 108 | return full_output.view_as(full_output_org) |
| 109 | |
| 110 | @classmethod |
| 111 | def can_replace_layer( |
| 112 | cls, |
| 113 | source_layer: nn.Module, |
| 114 | lora_config: LoRAConfig, |
| 115 | packed_modules_list: list, |
| 116 | model_config: PreTrainedConfig | None = None, |
| 117 | ) -> bool: |
| 118 | # Accept both the in-tree and any OOT class, because layers built on top of |
| 119 | # `VocabParallelEmbedding` (e.g. by the Transformers modelling backend) use the |
| 120 | # former while `PluggableLayer.__new__` instantiates the latter |
| 121 | embedding_classes = ( |
| 122 | VocabParallelEmbedding, |
| 123 | maybe_get_oot_by_class(VocabParallelEmbedding), |
| 124 | ) |
| 125 | lm_head_classes = (ParallelLMHead, maybe_get_oot_by_class(ParallelLMHead)) |
| 126 | return isinstance(source_layer, embedding_classes) and not isinstance( |
| 127 | source_layer, lm_head_classes |
| 128 | ) |
| 129 | |
| 130 | @property |
| 131 | def weight(self): |
| 132 | return self.base_layer.weight |