| 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 | from transformers import PreTrainedConfig |
| 7 | |
| 8 | from vllm.config.lora import LoRAConfig |
| 9 | from vllm.distributed import ( |
| 10 | get_tensor_model_parallel_rank, |
| 11 | get_tensor_model_parallel_world_size, |
| 12 | ) |
| 13 | from vllm.model_executor.layers.logits_processor import LogitsProcessor |
| 14 | from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding |
| 15 | from vllm.platforms import current_platform |
| 16 | |
| 17 | from .base import BaseLayerWithLoRA |
| 18 | |
| 19 | |
| 20 | class LogitsProcessorWithLoRA(BaseLayerWithLoRA): |
| 21 | """LoRA wrapper for LogitsProcessor, with extra logic to handle the |
| 22 | application of the LoRA adapter. |
| 23 | |
| 24 | Args: |
| 25 | base_layer: LogitsProcessor layer |
| 26 | hidden_size: hidden size of the model |
| 27 | dtype: data type of the model |
| 28 | device: device of the model |
| 29 | sharded_to_full_mapping: index mapping from sharded vocab to full vocab |
| 30 | received from base_layer.get_sharded_to_full_mapping(). If None, |
| 31 | no reindexing will be done. |
| 32 | |
| 33 | """ |
| 34 | |
| 35 | def __init__( |
| 36 | self, |
| 37 | base_layer: LogitsProcessor, |
| 38 | hidden_size: int, |
| 39 | dtype: torch.dtype, |
| 40 | device: torch.device, |
| 41 | sharded_to_full_mapping: list[int] | None, |
| 42 | ) -> None: |
| 43 | super().__init__() |
| 44 | self.base_layer = base_layer |
| 45 | self.hidden_size = hidden_size |
| 46 | self.dtype = dtype |
| 47 | self.device = device |
| 48 | self.tp_size = get_tensor_model_parallel_world_size() |
| 49 | self.tp_rank = get_tensor_model_parallel_rank() |
| 50 | self.sharded_to_full_mapping = sharded_to_full_mapping |
| 51 | |
| 52 | @property |
| 53 | def logits_as_input(self): |
| 54 | return self.base_layer.logits_as_input |
| 55 | |
| 56 | @property |
| 57 | def vocab_size(self): |
| 58 | return self.base_layer.vocab_size |
| 59 | |
| 60 | @property |
| 61 | def scale(self): |
| 62 | return self.base_layer.scale |
| 63 | |
| 64 | @property |
| 65 | def soft_cap(self): |
| 66 | return self.base_layer.soft_cap |
| 67 | |
| 68 | @property |
| 69 | def use_all_gather(self): |
| 70 | return self.base_layer.use_all_gather |
| 71 | |
| 72 | @property |
| 73 | def org_vocab_size(self): |
| 74 | return self.base_layer.org_vocab_size |
| 75 | |
| 76 | @property |
| 77 | def include_gpu_probs_tensor(self): |
| 78 | return self.base_layer.include_gpu_probs_tensor |
| 79 | |
| 80 | @property |
| 81 | def should_modify_greedy_probs_inplace(self): |
| 82 | return self.base_layer.should_modify_greedy_probs_inplace |
| 83 | |
| 84 | def create_lora_weights( |
| 85 | self, |
| 86 | max_loras: int, |
| 87 | lora_config: LoRAConfig, |
| 88 | model_config: PreTrainedConfig | None = None, |
| 89 | ) -> None: |
| 90 | # TODO: Verify if this condition can be further relaxed |
| 91 | if self.base_layer.vocab_size > 258048: |
| 92 | raise ValueError("When using LoRA, vocab size must be <= 258048") |
| 93 | self.lora_a_stacked = torch.zeros( |
| 94 | ( |
| 95 | max_loras, |
| 96 | 1, |
| 97 | lora_config.max_lora_rank, |
| 98 | self.hidden_size, |
| 99 | ), |
| 100 | dtype=lora_config.lora_dtype, |
| 101 | device=self.device, |
| 102 | ) |
| 103 | self.lora_b_stacked = torch.zeros( |
| 104 | ( |
| 105 | max_loras, |
| 106 | 1, |
| 107 | self.base_layer.vocab_size, |
| 108 | lora_config.max_lora_rank, |
| 109 | ), |
| 110 | dtype=lora_config.lora_dtype, |
| 111 | device=self.device, |
| 112 | ) |
| 113 | |
| 114 | if self.sharded_to_full_mapping is not None: |
| 115 | self.sharded_to_full_mapping_gpu = torch.tensor( |
| 116 | self.sharded_to_full_mapping, device=self.device, dtype=torch.long |
| 117 | ) |
| 118 | else: |
| 119 | self.sharded_to_full_mapping_gpu = None |
| 120 | |
| 121 | def reset_sharded_to_full_mapping(self) -> None: |
| 122 | """Restore the TP logits mapping after its GPU memory is reused.""" |
| 123 | mapping_gpu = self.sharded_to_full_mapping_gpu |
| 124 | if mapping_gpu is not None: |
| 125 | mapping_gpu.copy_( |
| 126 | torch.tensor( |
| 127 | self.sharded_to_full_mapping, |
| 128 | device=mapping_gpu.device, |
| 129 | dtype=mapping_gpu.dtype, |
| 130 | ) |
| 131 | ) |
| 132 | |
| 133 | def reset_lora(self, index: int): |
| 134 | self.lora_a_stacked[index] = 0 |
| 135 | self.lora_b_stacked[index] = 0 |
| 136 | |
| 137 | def set_lora( |
| 138 | self, |
| 139 | index: int, |
| 140 | lora_a: torch.Tensor | list[torch.Tensor], |
| 141 | lora_b: torch.Tensor | list[torch.Tensor], |
| 142 | ): |
| 143 | assert isinstance(lora_a, torch.Tensor) |
| 144 | assert isinstance(lora_b, torch.Tensor) |
| 145 | self.reset_lora(index) |
| 146 | self.lora_a_stacked[index, 0, : lora_a.shape[0], : lora_a.shape[1]].copy_( |
| 147 | lora_a, non_blocking=True |
| 148 | ) |
| 149 | self.lora_b_stacked[index, 0, : lora_b.shape[0], : lora_b.shape[1]].copy_( |
| 150 | lora_b, non_blocking=True |
| 151 | ) |
| 152 | |
| 153 | def _get_logits( |
| 154 | self, |
| 155 | hidden_states: torch.Tensor, |
| 156 | lm_head: VocabParallelEmbedding, |
| 157 | embedding_bias: torch.Tensor | None = None, |
| 158 | skip_gather: bool = False, |
| 159 | ) -> torch.Tensor | None: |
| 160 | # The LoRA delta is accumulated into the full gathered logits, so the |
| 161 | # TP gather cannot be skipped here. |
| 162 | if skip_gather: |
| 163 | raise NotImplementedError( |
| 164 | "Skipping the logits TP gather is not supported with an lm_head LoRA." |
| 165 | ) |
| 166 | # Get the logits for the next tokens. |
| 167 | if hasattr(lm_head, "base_layer"): |
| 168 | actual_lm_head = lm_head.base_layer |
| 169 | else: |
| 170 | actual_lm_head = lm_head |
| 171 | # Run the base projection through the LogitsProcessor so head_dtype |
| 172 | # (e.g. an fp32 lm_head) is honored on the LoRA path too. The LoRA |
| 173 | # delta is accumulated into these logits by add_lora_logits below, |
| 174 | # whose internal buffer is already fp32. |
| 175 | logits = self.base_layer._apply_head( |
| 176 | actual_lm_head, hidden_states, embedding_bias |
| 177 | ) |
| 178 | |
| 179 | # Gather logits for TP |
| 180 | logits = self.base_layer._gather_logits(logits) |
| 181 | |
| 182 | if logits is None: |
| 183 | return None |
| 184 | |
| 185 | if self.sharded_to_full_mapping_gpu is not None: |
| 186 | # Reindex the gathered logits so that index == token_id. Each TP |
| 187 | # shard is padded to `num_embeddings_per_partition`, so the gather |
| 188 | # interleaves the shards' padding with real vocab entries. |
| 189 | # Example for: |
| 190 | # org_vocab_size = 6 |
| 191 | # pad_to_size = 8 |
| 192 | # tp_size = 2 |
| 193 | |
| 194 | # indices: [0, 1, 2, 3, 4, 5, 6, 7] |
| 195 | # token_id: [0, 1, 2, -1, 3, 4, 5, -1] |
| 196 | |
| 197 | # Therefore, the mapping is expected to be: |
| 198 | # [0, 1, 2, 4, 5, 6, 3, 7] so that when we reindex, |
| 199 | # we get: |
| 200 | # indices: [0, 1, 2, 3, 4, 5, 6, 7] |
| 201 | # token_id: [0, 1, 2, 3, 4, 5, -1, -1] |
| 202 | logits = logits[:, self.sharded_to_full_mapping_gpu] |
| 203 | |
| 204 | lora_output: torch.Tensor | None = self.punica_wrapper.add_lora_logits( |
| 205 | logits, hidden_states, self.lora_a_stacked, self.lora_b_stacked, 1.0 |
| 206 | ) |
| 207 | |
| 208 | if not current_platform.can_update_inplace(): |
| 209 | logits = lora_output |
| 210 | |
| 211 | # Remove paddings in vocab (if any). |
| 212 | logits = logits[:, : self.base_layer.vocab_size] |
| 213 | return logits |
| 214 | |
| 215 | def forward(self, *args, **kwargs): |
| 216 | return type(self.base_layer).forward(self, *args, **kwargs) |
| 217 | |
| 218 | @classmethod |
| 219 | def can_replace_layer( |
| 220 | cls, |
| 221 | source_layer: nn.Module, |
| 222 | lora_config: LoRAConfig, |
| 223 | packed_modules_list: list, |
| 224 | model_config: PreTrainedConfig | None = None, |
| 225 | ) -> bool: |
| 226 | # Special handling for the LogitsProcessor. |
| 227 | return False |