| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | |
| 5 | import torch |
| 6 | import torch.nn as nn |
| 7 | from transformers import PreTrainedConfig |
| 8 | |
| 9 | from vllm.config.lora import LoRAConfig |
| 10 | from vllm.distributed import ( |
| 11 | split_tensor_along_last_dim, |
| 12 | tensor_model_parallel_all_reduce, |
| 13 | ) |
| 14 | from vllm.model_executor.custom_op import maybe_get_oot_by_class |
| 15 | from vllm.model_executor.layers.linear import RowParallelLinear |
| 16 | from vllm.platforms import current_platform |
| 17 | |
| 18 | from .base_linear import BaseLinearLayerWithLoRA |
| 19 | from .utils import _fully_sharded_can_replace, _not_fully_sharded_can_replace |
| 20 | |
| 21 | |
| 22 | class RowParallelLinearWithLoRA(BaseLinearLayerWithLoRA): |
| 23 | def __init__(self, base_layer: RowParallelLinear) -> None: |
| 24 | super().__init__(base_layer) |
| 25 | |
| 26 | # reset input_size |
| 27 | self.input_size = self.base_layer.input_size_per_partition |
| 28 | self.output_size = self.base_layer.output_size |
| 29 | # There is only one LoRA layer. |
| 30 | self.n_slices = 1 |
| 31 | |
| 32 | def slice_lora_a(self, lora_a: torch.Tensor) -> torch.Tensor: |
| 33 | shard_size = self.input_size |
| 34 | start_idx = self.tp_rank * shard_size |
| 35 | end_idx = (self.tp_rank + 1) * shard_size |
| 36 | lora_a = lora_a[:, start_idx:end_idx] |
| 37 | return lora_a |
| 38 | |
| 39 | def slice_lora_b(self, lora_b: torch.Tensor) -> torch.Tensor: |
| 40 | return lora_b |
| 41 | |
| 42 | def forward( |
| 43 | self, input_: torch.Tensor |
| 44 | ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor | None]: |
| 45 | """Forward of RowParallelLinear. |
| 46 | |
| 47 | Args: |
| 48 | input_: tensor whose last dimension is `input_size`. If |
| 49 | `input_is_parallel` is set, then the last dimension |
| 50 | is `input_size // tp_size`. |
| 51 | |
| 52 | Returns: |
| 53 | - output |
| 54 | - bias |
| 55 | |
| 56 | """ |
| 57 | # set up backprop all-reduce. |
| 58 | if self.base_layer.input_is_parallel: |
| 59 | input_parallel = input_ |
| 60 | else: |
| 61 | # TODO: simplify code below |
| 62 | split_input = split_tensor_along_last_dim( |
| 63 | input_, num_partitions=self.tp_size |
| 64 | ) |
| 65 | input_parallel = split_input[self.tp_rank].contiguous() |
| 66 | |
| 67 | # Matrix multiply. |
| 68 | bias_ = ( |
| 69 | None |
| 70 | if (self.tp_rank > 0 or self.base_layer.skip_bias_add) |
| 71 | else self.base_layer.bias |
| 72 | ) |
| 73 | output_parallel = self.apply(input_parallel, bias_) |
| 74 | if self.base_layer.reduce_results and self.tp_size > 1: |
| 75 | output = tensor_model_parallel_all_reduce(output_parallel) |
| 76 | else: |
| 77 | output = output_parallel |
| 78 | |
| 79 | output_bias = self.base_layer.bias if self.base_layer.skip_bias_add else None |
| 80 | if not self.base_layer.return_bias: |
| 81 | return output |
| 82 | |
| 83 | return output, output_bias |
| 84 | |
| 85 | @classmethod |
| 86 | @_not_fully_sharded_can_replace |
| 87 | def can_replace_layer( |
| 88 | cls, |
| 89 | source_layer: nn.Module, |
| 90 | lora_config: LoRAConfig, |
| 91 | packed_modules_list: list, |
| 92 | model_config: PreTrainedConfig | None = None, |
| 93 | ) -> bool: |
| 94 | return type(source_layer) is maybe_get_oot_by_class(RowParallelLinear) |
| 95 | |
| 96 | |
| 97 | # The following layer is based on the tensor parallelism strategy given in |
| 98 | # Y. Sheng et al., S-LoRA: Serving Thousands of Concurrent LoRA Adapters. 2023, |
| 99 | # https://arxiv.org/abs/2311.03285. |
| 100 | |
| 101 | |
| 102 | class RowParallelLinearWithShardedLoRA(RowParallelLinearWithLoRA): |
| 103 | """Differs from RowParallelLinearWithLoRA by slicing the |
| 104 | LoRA B's also. |
| 105 | |
| 106 | Based on S-LoRA, slicing happens along the output dim. |
| 107 | This yields a combined partial sum from the row parallel base |
| 108 | layer and column partitioned output from the LoRA. |
| 109 | """ |
| 110 | |
| 111 | def slice_lora_b(self, lora_b: torch.Tensor) -> torch.Tensor: |
| 112 | shard_size = self.lora_b_stacked[0].shape[2] |
| 113 | start_idx = self.tp_rank * shard_size |
| 114 | end_idx = (self.tp_rank + 1) * shard_size |
| 115 | lora_b = lora_b[start_idx:end_idx, :] |
| 116 | return lora_b |
| 117 | |
| 118 | def apply(self, x: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor: |
| 119 | output = self._get_quant_method().apply(self.base_layer, x, bias) |
| 120 | |
| 121 | x = x.view(-1, x.shape[-1]) |
| 122 | output, out_orig_shape = output.view(-1, output.shape[-1]), output.shape |
| 123 | buffer = torch.zeros( |
| 124 | (self.n_slices, x.shape[0], self.lora_a_stacked[0].shape[2]), |
| 125 | dtype=torch.float32, |
| 126 | device=x.device, |
| 127 | ) |
| 128 | |
| 129 | shrunk_buffer: torch.Tensor | None = self.punica_wrapper.add_shrink( |
| 130 | buffer, x, self.lora_a_stacked, 1.0 |
| 131 | ) |
| 132 | if not current_platform.can_update_inplace(): |
| 133 | buffer = shrunk_buffer |
| 134 | if self.tp_size > 1: |
| 135 | buffer = tensor_model_parallel_all_reduce(buffer) |
| 136 | |
| 137 | # following S-LoRA, allows the fusing of all_gather and all_reduce |
| 138 | # by adding the column partitioned lora output to a slice of output |
| 139 | # tensor, which is a partial sum due to row parallel. All that |
| 140 | # remains is a standard all_reduce. User should be aware though that |
| 141 | # the output is not the same as a normal row_parallel, it should be |
| 142 | # reduced before being used |
| 143 | # NOTE offset are based on the rank. |
| 144 | shard_size = self.lora_b_stacked[0].shape[2] |
| 145 | offset_start = self.tp_rank * shard_size |
| 146 | lora_output: torch.Tensor | None = self.punica_wrapper.add_expand( |
| 147 | output, |
| 148 | buffer, |
| 149 | self.lora_b_stacked, |
| 150 | self.output_slices, |
| 151 | offset_start=offset_start, |
| 152 | add_input=True, |
| 153 | ) |
| 154 | |
| 155 | if not current_platform.can_update_inplace(): |
| 156 | output = lora_output |
| 157 | |
| 158 | output = output.view(*out_orig_shape) |
| 159 | return output |
| 160 | |
| 161 | @classmethod |
| 162 | @_fully_sharded_can_replace |
| 163 | def can_replace_layer( |
| 164 | cls, |
| 165 | source_layer: nn.Module, |
| 166 | lora_config: LoRAConfig, |
| 167 | packed_modules_list: list, |
| 168 | model_config: PreTrainedConfig | None = None, |
| 169 | ) -> bool: |
| 170 | # specifying kwargs so they can be easily accessed in decorator |
| 171 | return super().can_replace_layer( |
| 172 | source_layer=source_layer, |
| 173 | lora_config=lora_config, |
| 174 | packed_modules_list=packed_modules_list, |
| 175 | model_config=model_config, |
| 176 | decorate=False, |
| 177 | ) |