diff --git a/vllm/lora/layers/__init__.py b/vllm/lora/layers/__init__.py index 1f3fdea2cd..e0c5f202bd 100644 --- a/vllm/lora/layers/__init__.py +++ b/vllm/lora/layers/__init__.py @@ -20,7 +20,7 @@ from vllm.lora.layers.row_parallel_linear import ( RowParallelLinearWithShardedLoRA, ) from vllm.lora.layers.utils import LoRAMapping, LoRAMappingType -from vllm.lora.layers.vocal_parallel_embedding import VocabParallelEmbeddingWithLoRA +from vllm.lora.layers.vocab_parallel_embedding import VocabParallelEmbeddingWithLoRA __all__ = [ "BaseLayerWithLoRA", diff --git a/vllm/lora/layers/logits_processor.py b/vllm/lora/layers/logits_processor.py index 7db5462d8e..0199437ba2 100644 --- a/vllm/lora/layers/logits_processor.py +++ b/vllm/lora/layers/logits_processor.py @@ -20,7 +20,7 @@ from .base import BaseLayerWithLoRA class LogitsProcessorWithLoRA(BaseLayerWithLoRA): """ LoRA wrapper for LogitsProcessor, with extra logic to handle the - application of the LoRA adapter and added LoRA vocabulary. + application of the LoRA adapter. Args: base_layer: LogitsProcessor layer @@ -183,19 +183,19 @@ class LogitsProcessorWithLoRA(BaseLayerWithLoRA): return None if self.sharded_to_full_mapping_gpu is not None: - # Reindex full logits tensor to ensure 1:1 mapping between - # index and token_id + # Reindex the gathered logits so that index == token_id. Each TP + # shard is padded to `num_embeddings_per_partition`, so the gather + # interleaves the shards' padding with real vocab entries. # Example for: - # org_vocab_size = 4 - # added_vocab_size = 2 + # org_vocab_size = 6 # pad_to_size = 8 # tp_size = 2 - # indices: [0, 1, 2, 3, 4, 5, 6, 7] - # token_id: [0, 1, 4, -1, 2, 3, 5, -1] + # indices: [0, 1, 2, 3, 4, 5, 6, 7] + # token_id: [0, 1, 2, -1, 3, 4, 5, -1] # Therefore, the mapping is expected to be: - # [0, 1, 4, 6, 2, 3, 5, 7] so that when we reindex, + # [0, 1, 2, 4, 5, 6, 3, 7] so that when we reindex, # we get: # indices: [0, 1, 2, 3, 4, 5, 6, 7] # token_id: [0, 1, 2, 3, 4, 5, -1, -1] diff --git a/vllm/lora/layers/vocal_parallel_embedding.py b/vllm/lora/layers/vocab_parallel_embedding.py similarity index 79% rename from vllm/lora/layers/vocal_parallel_embedding.py rename to vllm/lora/layers/vocab_parallel_embedding.py index 6aade587cd..c6f294efcd 100644 --- a/vllm/lora/layers/vocal_parallel_embedding.py +++ b/vllm/lora/layers/vocab_parallel_embedding.py @@ -21,8 +21,6 @@ class VocabParallelEmbeddingWithLoRA(BaseLayerWithLoRA): def __init__(self, base_layer: VocabParallelEmbedding) -> None: super().__init__() self.base_layer = base_layer - self.embeddings_slice: tuple[int, int] | None - self.embeddings_weights: torch.Tensor | None def create_lora_weights( self, @@ -30,25 +28,6 @@ class VocabParallelEmbeddingWithLoRA(BaseLayerWithLoRA): lora_config: LoRAConfig, model_config: PretrainedConfig | None = None, ) -> None: - if self.base_layer.num_added_embeddings_per_partition > 0: - # We can start adding lora weights - self.embeddings_weights = self.base_layer.weight.data[ - self.base_layer.num_org_embeddings_per_partition : self.base_layer.num_org_embeddings_per_partition # noqa: E501 - + self.base_layer.num_added_embeddings_per_partition - ] - self.embeddings_slice = ( - self.base_layer.shard_indices.added_vocab_start_index - - self.base_layer.org_vocab_size, - self.base_layer.shard_indices.added_vocab_end_index - - self.base_layer.org_vocab_size, - ) - self.base_layer.weight.data[ - self.base_layer.num_org_embeddings_per_partition : - ].fill_(0) - else: - self.embeddings_slice = None - self.embeddings_weights = None - self.lora_a_stacked = torch.zeros( ( max_loras, @@ -100,10 +79,10 @@ class VocabParallelEmbeddingWithLoRA(BaseLayerWithLoRA): # NB: Don't use torch.narrow here. torch.narrow triggers some # Dynamic Shape specialization in torch.compile num_tokens = x.shape[0] - indices_1 = self.punica_wrapper._embeddings_indices[1][:num_tokens] + embeddings_indices = self.punica_wrapper._embeddings_indices[:num_tokens] full_lora_a_embeddings = F.embedding( - x + indices_1, + x + embeddings_indices, self.lora_a_stacked_2d, ) full_output = self.base_layer.forward(x) diff --git a/vllm/lora/punica_wrapper/punica_base.py b/vllm/lora/punica_wrapper/punica_base.py index 0867129919..f5c7ca9cb7 100644 --- a/vllm/lora/punica_wrapper/punica_base.py +++ b/vllm/lora/punica_wrapper/punica_base.py @@ -145,7 +145,7 @@ class PunicaWrapperBase(PunicaWrapperABC): max_num_batched_tokens, dtype=torch.long, device=device ) self._embeddings_indices = torch.empty( - 2, max_num_batched_tokens, dtype=torch.long, device=device + max_num_batched_tokens, dtype=torch.long, device=device ) # 4 is the number of indices tensors. @@ -172,10 +172,6 @@ class PunicaWrapperBase(PunicaWrapperABC): max_loras: int, vocab_size: int, ): - # NOTE We have remove lora extra vocab support for now. So we set - # extra_vocab_size always to 0, and extra_vocab_size will be removed. - - extra_vocab_size = 0 ( base_indices, sampler_indices, @@ -187,7 +183,6 @@ class PunicaWrapperBase(PunicaWrapperABC): lora_index_to_id, max_loras, vocab_size, - extra_vocab_size, self.device, ) self._token_lora_indices[: base_indices.shape[0]].copy_(base_indices) @@ -195,9 +190,9 @@ class PunicaWrapperBase(PunicaWrapperABC): self._sampler_indices_padded[: sampler_indices_padded.shape[0]].copy_( sampler_indices_padded ) - self._embeddings_indices[ - : embeddings_indices.shape[0], : embeddings_indices.shape[1] - ].copy_(embeddings_indices) + self._embeddings_indices[: embeddings_indices.shape[0]].copy_( + embeddings_indices + ) self.indices_len[:] = indices_len @@ -272,15 +267,6 @@ class PunicaWrapperBase(PunicaWrapperABC): indices_padded_len = self.indices_len[2] return self._sampler_indices_padded[:indices_padded_len] - @property - def embeddings_indices(self) -> torch.Tensor: - """ - This property provides access to the indices used for lora embeddings, - specifically for VocabParallelEmbeddingWithLoRA. - """ - embeddings_indices_len = self.indices_len[3] - return self._embeddings_indices[:, :embeddings_indices_len] - def update_metadata( self, mapping: "LoRAMapping", diff --git a/vllm/lora/punica_wrapper/punica_xpu.py b/vllm/lora/punica_wrapper/punica_xpu.py index 96b411abd1..7f040732ee 100755 --- a/vllm/lora/punica_wrapper/punica_xpu.py +++ b/vllm/lora/punica_wrapper/punica_xpu.py @@ -44,7 +44,7 @@ class PunicaWrapperXPU(PunicaWrapperBase): ): PunicaWrapperBase.__init__(self, max_num_batched_tokens, max_batches, device) torch._dynamo.mark_dynamic(self._token_lora_indices, 0) - torch._dynamo.mark_dynamic(self._embeddings_indices, 1) + torch._dynamo.mark_dynamic(self._embeddings_indices, 0) torch._dynamo.mark_dynamic(self._sampler_indices_padded, 0) self.lora_config = kwargs["lora_config"] diff --git a/vllm/lora/punica_wrapper/utils.py b/vllm/lora/punica_wrapper/utils.py index 257e88c10c..c29e359f19 100644 --- a/vllm/lora/punica_wrapper/utils.py +++ b/vllm/lora/punica_wrapper/utils.py @@ -56,7 +56,6 @@ def convert_mapping( lora_index_to_id: list[int | None], max_loras: int, vocab_size: int, - extra_vocab_size: int, device: torch.device, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, list[int]]: """Converts LoRAMapping to index tensors. @@ -66,7 +65,6 @@ def convert_mapping( lora_index_to_id: List mapping LoRA ids to LoRA indices. max_loras: Maximum number of LoRAs. vocab_size: Model vocab size. - extra_vocab_size: Extra vocab size each LoRA can have. Returns: A tuple of tensors: @@ -80,10 +78,8 @@ def convert_mapping( requests to LoRA indices for sampler with padding. Same as sampler_indices, but -1 is replaced with max_loras. - embeddings_indices: Tensor of shape [2, batch_size] mapping - requests to embedding indices. First row is for embeddings - added by the LoRAs, second row is for the LoRA.lora_a - embeddings. + embeddings_indices: Tensor of shape [batch_size] mapping + requests to the row offset of their LoRA.lora_a embeddings. indices_len: List of lengths of the above tensors. It contains (base_indices, sampler_indices, sampler_indices_padded, embeddings_indices). @@ -124,15 +120,7 @@ def convert_mapping( prompt_mapping_tensor = async_tensor_h2d( prompt_mapping, dtype=torch.long, device=device ) - embeddings_indices = torch.stack( - [ - indices[2] * extra_vocab_size, - indices[2] * (vocab_size + extra_vocab_size), - ] - ) - embeddings_indices = torch.where( - embeddings_indices == -1, max_loras - 1, embeddings_indices - ) + embeddings_indices = indices[2] * vocab_size base_indices = indices[1] sampler_indices = prompt_mapping_tensor sampler_indices_padded = sampler_indices.clone()