A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[Core] Split LoRA layers (#24574)
Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>
Jee Jee Li committed
Sep 10, 2025 at 22:47 UTC
bb3eb80d92f66eed4016fa5ed9d9e66125a14482
14 files changed
+1668
-1557
tests/lora/test_layers.py
+5
-5
@@ -12,20 +12,20 @@ import torch
12
import torch.nn.functional as F
13
14
from vllm.config import LoRAConfig
15
-from vllm.lora.fully_sharded_layers import (
16
- ColumnParallelLinearWithShardedLoRA,
17
- MergedColumnParallelLinearWithShardedLoRA,
18
- MergedQKVParallelLinearWithShardedLoRA, QKVParallelLinearWithShardedLoRA,
19
- RowParallelLinearWithShardedLoRA)
15
# yapf conflicts with isort for this block
16
# yapf: disable
17
from vllm.lora.layers import (BaseLayerWithLoRA, ColumnParallelLinearWithLoRA,
18
+ ColumnParallelLinearWithShardedLoRA,
19
LogitsProcessorWithLoRA, LoRAMapping,
20
MergedColumnParallelLinearWithLoRA,
21
+ MergedColumnParallelLinearWithShardedLoRA,
22
MergedQKVParallelLinearWithLoRA,
23
+ MergedQKVParallelLinearWithShardedLoRA,
24
QKVParallelLinearWithLoRA,
25
+ QKVParallelLinearWithShardedLoRA,
26
ReplicatedLinearWithLoRA,
27
RowParallelLinearWithLoRA,
28
+ RowParallelLinearWithShardedLoRA,
29
VocabParallelEmbeddingWithLoRA)
30
# yapf: enable
31
from vllm.lora.models import LoRALayerWeights, PackedLoRALayerWeights
vllm/lora/fully_sharded_layers.py
deleted
-355
@@ -1,355 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-# pylint: disable=unused-argument
5
-from typing import TYPE_CHECKING, Optional, Union, cast
6
-
7
-import torch
8
-import torch.nn as nn
9
-from transformers import PretrainedConfig
10
-
11
-from vllm.config import LoRAConfig
12
-from vllm.distributed.communication_op import (
13
- tensor_model_parallel_all_gather, tensor_model_parallel_all_reduce)
14
-from vllm.distributed.parallel_state import get_tensor_model_parallel_rank
15
-from vllm.lora.layers import (ColumnParallelLinearWithLoRA,
16
- MergedColumnParallelLinearWithLoRA,
17
- MergedQKVParallelLinearWithLoRA,
18
- QKVParallelLinearWithLoRA,
19
- RowParallelLinearWithLoRA)
20
-from vllm.platforms import current_platform
21
-
22
-if TYPE_CHECKING:
23
- pass
24
-
25
-
26
-def _fully_sharded_can_replace(can_replace):
27
- """
28
- decorator which adds the condition of fully sharded loras
29
- intended to wrap can_replace_layer()
30
- """
31
-
32
- def dec(*args, **kwargs):
33
- return (can_replace(*args, **kwargs)
34
- and kwargs["lora_config"].fully_sharded_loras)
35
-
36
- return dec
37
-
38
-
39
-def _mcp_apply(x, bias, layer: ColumnParallelLinearWithLoRA):
40
- """
41
- For `ColumnParallelLinearWithLoRA` or classes that inherit from
42
- `ColumnParallelLinearWithLoRA`, they share the same `apply` logic.
43
- """
44
- assert (layer.n_slices == len(layer.lora_a_stacked) == len(
45
- layer.lora_b_stacked) == len(layer.output_slices))
46
- if layer.lora_bias_stacked is not None:
47
- assert layer.n_slices == len(layer.lora_bias_stacked)
48
-
49
- output = layer.base_layer.quant_method.apply(layer.base_layer, x, bias)
50
-
51
- x = x.view(-1, x.shape[-1])
52
- output, out_orig_shape = output.view(-1, output.shape[-1]), output.shape
53
-
54
- # Since communication is needed, the buffer is directly initialized as a
55
- # tensor rather than a tuple of tensor.
56
- buffers = torch.zeros(
57
- (layer.n_slices, x.shape[0], layer.lora_a_stacked[0].shape[2]),
58
- dtype=torch.float32,
59
- device=x.device,
60
- )
61
-
62
- shrunk_buffers: Optional[torch.Tensor] = layer.punica_wrapper.add_shrink(
63
- buffers, x, layer.lora_a_stacked, 1.0)
64
-
65
- if not current_platform.can_update_inplace():
66
- buffers = shrunk_buffers
67
-
68
- buffers = tensor_model_parallel_all_gather(buffers)
69
-
70
- lora_output: Optional[torch.Tensor] = layer.punica_wrapper.add_expand(
71
- output,
72
- buffers,
73
- layer.lora_b_stacked,
74
- layer.lora_bias_stacked,
75
- layer.output_slices,
76
- offset_start=0,
77
- add_input=True)
78
-
79
- if not current_platform.can_update_inplace():
80
- output = lora_output
81
-
82
- output = output.view(*out_orig_shape)
83
- # now have column partitioned and packed output
84
- return output
85
-
86
-
87
-# these layers are based on the tensor parallelism strategy given in
88
-# Y. Sheng et al., S-LoRA: Serving Thousands of Concurrent LoRA Adapters. 2023,
89
-# https://arxiv.org/abs/2311.03285.
90
-
91
-
92
-class ColumnParallelLinearWithShardedLoRA(ColumnParallelLinearWithLoRA):
93
- """
94
- Differs from ColumnParallelLinearWithLoRA by slicing LoRA A also.
95
-
96
- Based on S-LoRA, slicing happens along the rank dim.
97
- """
98
-
99
- # For all LoRA layers where the `base_layer` is `ColumnParallelLinear`,
100
- # their `lora_a` and `lora_b` have different sharding patterns. After
101
- # completing the `lora_a` GEMM , a gather operation is performed.
102
- # Therefore, the sharding of `lora_a` only needs to correspond with the
103
- # gather operation.
104
- def slice_lora_a(self, lora_a: torch.Tensor) -> torch.Tensor:
105
- tp_rank = get_tensor_model_parallel_rank()
106
- shard_size = self.lora_a_stacked[0].shape[2]
107
- start_idx = tp_rank * shard_size
108
- lora_a = lora_a[:, start_idx:start_idx + shard_size]
109
- return lora_a
110
-
111
- def apply(self,
112
- x: torch.Tensor,
113
- bias: Optional[torch.Tensor] = None) -> torch.Tensor:
114
- return _mcp_apply(x, bias, self)
115
-
116
- @classmethod
117
- @_fully_sharded_can_replace
118
- def can_replace_layer(
119
- cls,
120
- source_layer: nn.Module,
121
- lora_config: LoRAConfig,
122
- packed_modules_list: list,
123
- model_config: Optional[PretrainedConfig],
124
- ) -> bool:
125
- # specifying kwargs so they can be easily accessed in decorator
126
- return super().can_replace_layer(
127
- source_layer=source_layer,
128
- lora_config=lora_config,
129
- packed_modules_list=packed_modules_list,
130
- model_config=model_config,
131
- decorate=False,
132
- )
133
-
134
-
135
-class MergedColumnParallelLinearWithShardedLoRA(
136
- MergedColumnParallelLinearWithLoRA):
137
- """
138
- Differs from MergedColumnParallelLinearWithLoRA by slicing the
139
- LoRA A's also.
140
-
141
- Based on S-LoRA, slicing happens along the rank dim.
142
- """
143
-
144
- def slice_lora_a(
145
- self, lora_a: list[Union[torch.Tensor, None]]
146
- ) -> list[Union[torch.Tensor, None]]:
147
- #NOTE: lora_a contains 2 subloras, and each sublora could be None.
148
- output_shard_size = self.lora_a_stacked[0].shape[2]
149
- output_start_idx = self.tp_rank * output_shard_size
150
- lora_a = [
151
- lora_a[0][:, output_start_idx:output_start_idx +
152
- output_shard_size] if lora_a[0] is not None else None,
153
- lora_a[1][:, output_start_idx:output_start_idx +
154
- output_shard_size] if lora_a[1] is not None else None,
155
- ]
156
- return lora_a
157
-
158
- def apply(self,
159
- x: torch.Tensor,
160
- bias: Optional[torch.Tensor] = None) -> torch.Tensor:
161
- return _mcp_apply(x, bias, self)
162
-
163
- @classmethod
164
- @_fully_sharded_can_replace
165
- def can_replace_layer(
166
- cls,
167
- source_layer: nn.Module,
168
- lora_config: LoRAConfig,
169
- packed_modules_list: list,
170
- model_config: Optional[PretrainedConfig],
171
- ) -> bool:
172
- # specifying kwargs so they can be easily accessed in decorator
173
- return super().can_replace_layer(
174
- source_layer=source_layer,
175
- lora_config=lora_config,
176
- packed_modules_list=packed_modules_list,
177
- model_config=model_config,
178
- decorate=False,
179
- )
180
-
181
-
182
-class QKVParallelLinearWithShardedLoRA(QKVParallelLinearWithLoRA):
183
- """
184
- Differs from QKVParallelLinearWithLoRA by slicing the
185
- LoRA A's also.
186
-
187
- Based on S-LoRA, slicing happens along the rank dim.
188
- """
189
-
190
- def slice_lora_a(self, lora_a: torch.Tensor) -> torch.Tensor:
191
- tp_rank = get_tensor_model_parallel_rank()
192
- shard_size = self.lora_a_stacked[0].shape[2]
193
- start_idx = tp_rank * shard_size
194
- lora_a = lora_a[:, start_idx:start_idx + shard_size]
195
- return lora_a
196
-
197
- def apply(self,
198
- x: torch.Tensor,
199
- bias: Optional[torch.Tensor] = None) -> torch.Tensor:
200
- return _mcp_apply(x, bias, self)
201
-
202
- @classmethod
203
- @_fully_sharded_can_replace
204
- def can_replace_layer(cls, source_layer: nn.Module,
205
- lora_config: LoRAConfig, packed_modules_list: list,
206
- model_config: Optional[PretrainedConfig]) -> bool:
207
- # specifying kwargs so they can be easily accessed in decorator
208
- return super().can_replace_layer(
209
- source_layer=source_layer,
210
- lora_config=lora_config,
211
- packed_modules_list=packed_modules_list,
212
- model_config=model_config,
213
- decorate=False,
214
- )
215
-
216
-
217
-class MergedQKVParallelLinearWithShardedLoRA(MergedQKVParallelLinearWithLoRA):
218
- """
219
- Differs from MergedQKVParallelLinearWithLoRA by slicing the
220
- LoRA A's also.
221
-
222
- Based on S-LoRA, slicing happens along the rank dim.
223
- """
224
-
225
- def slice_lora_a(
226
- self, lora_a: list[Union[torch.Tensor, None]]
227
- ) -> list[Union[torch.Tensor, None]]:
228
- # NOTE: lora_a contains 3 subloras, and each sublora could be None.
229
- shard_size = [self.lora_a_stacked[i].shape[2] for i in range(3)]
230
- start_idx = [self.tp_rank * shard_size[i] for i in range(3)]
231
- lora_a = [
232
- lora_a[0][:, start_idx[0]:start_idx[0] +
233
- shard_size[0]] if lora_a[0] is not None else None,
234
- lora_a[1][:, start_idx[1]:start_idx[1] +
235
- shard_size[1]] if lora_a[1] is not None else None,
236
- lora_a[2][:, start_idx[2]:start_idx[2] +
237
- shard_size[2]] if lora_a[2] is not None else None,
238
- ]
239
- return lora_a
240
-
241
- def apply(self,
242
- x: torch.Tensor,
243
- bias: Optional[torch.Tensor] = None) -> torch.Tensor:
244
- return _mcp_apply(x, bias, self)
245
-
246
- @classmethod
247
- @_fully_sharded_can_replace
248
- def can_replace_layer(
249
- cls,
250
- source_layer: nn.Module,
251
- lora_config: LoRAConfig,
252
- packed_modules_list: list,
253
- model_config: Optional[PretrainedConfig],
254
- ) -> bool:
255
- # specifying kwargs so they can be easily accessed in decorator
256
- return super().can_replace_layer(
257
- source_layer=source_layer,
258
- lora_config=lora_config,
259
- packed_modules_list=packed_modules_list,
260
- model_config=model_config,
261
- decorate=False,
262
- )
263
-
264
-
265
-class RowParallelLinearWithShardedLoRA(RowParallelLinearWithLoRA):
266
- """
267
- Differs from RowParallelLinearWithLoRA by slicing the
268
- LoRA B's also.
269
-
270
- Based on S-LoRA, slicing happens along the output dim.
271
- This yields a combined partial sum from the row parallel base
272
- layer and column partitioned output from the LoRA.
273
- """
274
-
275
- def slice_lora_b(self, lora_b: torch.Tensor) -> torch.Tensor:
276
- shard_size = self.lora_b_stacked[0].shape[2]
277
- start_idx = self.tp_rank * shard_size
278
- end_idx = (self.tp_rank + 1) * shard_size
279
- lora_b = lora_b[:, start_idx:end_idx]
280
- return lora_b
281
-
282
- def slice_bias(self, bias: torch.Tensor) -> torch.Tensor:
283
- if bias is None:
284
- return bias
285
- self.lora_bias_stacked = cast(tuple[torch.Tensor, ...],
286
- self.lora_bias_stacked)
287
- shard_size = self.lora_bias_stacked[0].shape[2]
288
- start_idx = self.tp_rank * shard_size
289
- end_idx = (self.tp_rank + 1) * shard_size
290
- bias = bias[start_idx:end_idx]
291
- return bias
292
-
293
- def apply(self,
294
- x: torch.Tensor,
295
- bias: Optional[torch.Tensor] = None) -> torch.Tensor:
296
- output = self.base_layer.quant_method.apply(self.base_layer, x)
297
-
298
- x = x.view(-1, x.shape[-1])
299
- output, out_orig_shape = output.view(-1,
300
- output.shape[-1]), output.shape
301
- buffer = torch.zeros(
302
- (self.n_slices, x.shape[0], self.lora_a_stacked[0].shape[2]),
303
- dtype=torch.float32,
304
- device=x.device,
305
- )
306
-
307
- shrunk_buffer: Optional[torch.Tensor] = self.punica_wrapper.add_shrink(
308
- buffer, x, self.lora_a_stacked, 1.0)
309
- if not current_platform.can_update_inplace():
310
- buffer = shrunk_buffer
311
-
312
- buffer = tensor_model_parallel_all_reduce(buffer)
313
-
314
- # following S-LoRA, allows the fusing of all_gather and all_reduce
315
- # by adding the column partitioned lora output to a slice of output
316
- # tensor, which is a partial sum due to row parallel. All that
317
- # remains is a standard all_reduce. User should be aware though that
318
- # the output is not the same as a normal row_parallel, it should be
319
- # reduced before being used
320
- # NOTE offset are based on the rank.
321
- shard_size = self.lora_b_stacked[0].shape[2]
322
- offset_start = self.tp_rank * shard_size
323
- lora_output: Optional[torch.Tensor] = self.punica_wrapper.add_expand(
324
- output,
325
- buffer,
326
- self.lora_b_stacked,
327
- self.lora_bias_stacked,
328
- self.output_slices,
329
- offset_start=offset_start,
330
- add_input=True,
331
- )
332
-
333
- if not current_platform.can_update_inplace():
334
- output = lora_output
335
-
336
- output = output.view(*out_orig_shape)
337
- return output
338
-
339
- @classmethod
340
- @_fully_sharded_can_replace
341
- def can_replace_layer(
342
- cls,
343
- source_layer: nn.Module,
344
- lora_config: LoRAConfig,
345
- packed_modules_list: list,
346
- model_config: Optional[PretrainedConfig],
347
- ) -> bool:
348
- # specifying kwargs so they can be easily accessed in decorator
349
- return super().can_replace_layer(
350
- source_layer=source_layer,
351
- lora_config=lora_config,
352
- packed_modules_list=packed_modules_list,
353
- model_config=model_config,
354
- decorate=False,
355
- )
vllm/lora/layers.py
deleted
-1192
@@ -1,1192 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-# pylint: disable=unused-argument
5
-import math
6
-from dataclasses import dataclass
7
-from typing import TYPE_CHECKING, Optional, Union, cast
8
-
9
-import torch
10
-import torch.nn as nn
11
-import torch.nn.functional as F
12
-from transformers import PretrainedConfig
13
-
14
-from vllm.adapter_commons.layers import AdapterMapping
15
-from vllm.config import LoRAConfig
16
-from vllm.distributed import (get_tensor_model_parallel_rank,
17
- get_tensor_model_parallel_world_size,
18
- split_tensor_along_last_dim,
19
- tensor_model_parallel_all_gather,
20
- tensor_model_parallel_all_reduce)
21
-from vllm.distributed.utils import divide
22
-# yapf: disable
23
-from vllm.model_executor.layers.linear import (ColumnParallelLinear,
24
- LinearBase,
25
- MergedColumnParallelLinear,
26
- QKVParallelLinear,
27
- ReplicatedLinear,
28
- RowParallelLinear)
29
-# yapf: enable
30
-from vllm.model_executor.layers.logits_processor import LogitsProcessor
31
-from vllm.model_executor.layers.vocab_parallel_embedding import (
32
- VocabParallelEmbedding)
33
-from vllm.platforms import current_platform
34
-
35
-if TYPE_CHECKING:
36
- from vllm.lora.punica_wrapper import PunicaWrapperBase
37
-
38
-
39
-def _get_lora_device(base_layer: nn.Module) -> torch.device:
40
- # code borrowed from https://github.com/fmmoret/vllm/blob/fm-support-lora-on-quantized-models/vllm/lora/layers.py#L34
41
- """Returns the device for where to place the LoRA tensors."""
42
- # unquantizedLinear
43
- if hasattr(base_layer, "weight"):
44
- return base_layer.weight.device
45
- # Compressed Tensor
46
- elif hasattr(base_layer, "weight_packed"):
47
- return base_layer.weight_packed.device
48
- # GPTQ/AWQ
49
- elif hasattr(base_layer, "qweight"):
50
- return base_layer.qweight.device
51
- # HQQ marlin
52
- elif hasattr(base_layer, "W_q"):
53
- return base_layer.W_q.device
54
- else:
55
- raise ValueError(f"Unsupported base layer: {base_layer}")
56
-
57
-
58
-def _not_fully_sharded_can_replace(can_replace):
59
- """
60
- decorator which adds the condition of not using fully sharded loras
61
- intended to wrap can_replace_layer()
62
- """
63
-
64
- def dec(*args, **kwargs):
65
- decorate = kwargs.pop("decorate") if "decorate" in kwargs else True
66
- condition = (not kwargs["lora_config"].fully_sharded_loras
67
- if decorate else True)
68
- return can_replace(*args, **kwargs) and condition
69
-
70
- return dec
71
-
72
-
73
-@dataclass
74
-class LoRAMapping(AdapterMapping):
75
- is_prefill: bool = False
76
-
77
-
78
-class BaseLayerWithLoRA(nn.Module):
79
-
80
- def slice_lora_a(
81
- self, lora_a: Union[torch.Tensor, list[Union[torch.Tensor, None]]]
82
- ) -> Union[torch.Tensor, list[Union[torch.Tensor, None]]]:
83
- """Slice lora a if splitting for tensor parallelism."""
84
- ...
85
-
86
- def slice_lora_b(
87
- self, lora_b: Union[torch.Tensor, list[Union[torch.Tensor, None]]]
88
- ) -> Union[torch.Tensor, list[Union[torch.Tensor, None]]]:
89
- """Slice lora b if splitting with tensor parallelism."""
90
- ...
91
-
92
- def create_lora_weights(
93
- self,
94
- max_loras: int,
95
- lora_config: LoRAConfig,
96
- model_config: Optional[PretrainedConfig] = None,
97
- ) -> None:
98
- """Initializes lora matrices."""
99
- ...
100
-
101
- def reset_lora(self, index: int):
102
- """Resets the lora weights at index back to 0."""
103
- ...
104
-
105
- def set_lora(
106
- self,
107
- index: int,
108
- lora_a: torch.Tensor,
109
- lora_b: torch.Tensor,
110
- embeddings_tensor: Optional[torch.Tensor],
111
- bias: Optional[torch.Tensor] = None,
112
- ):
113
- """Overwrites lora tensors at index."""
114
- ...
115
-
116
- def set_mapping(
117
- self,
118
- punica_wrapper,
119
- ):
120
- self.punica_wrapper: PunicaWrapperBase = punica_wrapper
121
-
122
- @classmethod
123
- def can_replace_layer(
124
- cls,
125
- source_layer: nn.Module,
126
- lora_config: LoRAConfig,
127
- packed_modules_list: list,
128
- model_config: Optional[PretrainedConfig],
129
- ) -> bool:
130
- """Returns True if the layer can be replaced by this LoRA layer."""
131
- raise NotImplementedError
132
-
133
-
134
-class VocabParallelEmbeddingWithLoRA(BaseLayerWithLoRA):
135
-
136
- def __init__(self, base_layer: VocabParallelEmbedding) -> None:
137
- super().__init__()
138
- self.base_layer = base_layer
139
- self.embeddings_slice: Optional[tuple[int, int]]
140
- self.embeddings_weights: Optional[torch.Tensor]
141
-
142
- def create_lora_weights(
143
- self,
144
- max_loras: int,
145
- lora_config: LoRAConfig,
146
- model_config: Optional[PretrainedConfig] = None) -> None:
147
-
148
- if self.base_layer.num_added_embeddings_per_partition > 0:
149
- # We can start adding lora weights
150
- self.embeddings_weights = self.base_layer.weight.data[
151
- self.base_layer.num_org_embeddings_per_partition:self.
152
- base_layer.num_org_embeddings_per_partition +
153
- self.base_layer.num_added_embeddings_per_partition]
154
- self.embeddings_slice = (
155
- self.base_layer.shard_indices.added_vocab_start_index -
156
- self.base_layer.org_vocab_size,
157
- self.base_layer.shard_indices.added_vocab_end_index -
158
- self.base_layer.org_vocab_size)
159
- self.base_layer.weight.data[
160
- self.base_layer.num_org_embeddings_per_partition:].fill_(0)
161
- else:
162
- self.embeddings_slice = None
163
- self.embeddings_weights = None
164
-
165
- self.embeddings_tensors = torch.zeros(
166
- (
167
- max_loras,
168
- lora_config.lora_extra_vocab_size,
169
- self.base_layer.embedding_dim,
170
- ),
171
- dtype=self.base_layer.weight.dtype,
172
- device=self.base_layer.weight.device,
173
- )
174
- self.lora_a_stacked = torch.zeros(
175
- (
176
- max_loras,
177
- self.base_layer.org_vocab_size +
178
- lora_config.lora_extra_vocab_size,
179
- lora_config.max_lora_rank,
180
- ),
181
- dtype=lora_config.lora_dtype,
182
- device=self.base_layer.weight.device,
183
- )
184
- self.lora_b_stacked = torch.zeros(
185
- (
186
- max_loras,
187
- 1,
188
- self.base_layer.embedding_dim,
189
- lora_config.max_lora_rank,
190
- ),
191
- dtype=lora_config.lora_dtype,
192
- device=self.base_layer.weight.device,
193
- )
194
- self.lora_a_stacked_2d = self.lora_a_stacked.view(
195
- self.lora_a_stacked.shape[0] * self.lora_a_stacked.shape[1],
196
- self.lora_a_stacked.shape[2],
197
- )
198
-
199
- def reset_lora(self, index: int):
200
- self.lora_a_stacked[index] = 0
201
- self.lora_b_stacked[index] = 0
202
- self.embeddings_tensors[index] = 0
203
-
204
- def set_lora(
205
- self,
206
- index: int,
207
- lora_a: torch.Tensor,
208
- lora_b: torch.Tensor,
209
- embeddings_tensor: Optional[torch.Tensor],
210
- bias: Optional[torch.Tensor] = None,
211
- ):
212
- self.reset_lora(index)
213
- self.lora_a_stacked[index, :lora_a.shape[0], :lora_a.shape[1]].copy_(
214
- lora_a, non_blocking=True)
215
- self.lora_b_stacked[index,
216
- 0, :lora_b.shape[1], :lora_b.shape[0]].copy_(
217
- lora_b.T, non_blocking=True)
218
- if embeddings_tensor is not None:
219
- self.embeddings_tensors[
220
- index,
221
- :embeddings_tensor.shape[0],
222
- :embeddings_tensor.shape[1],
223
- ].copy_(embeddings_tensor, non_blocking=True)
224
- if self.embeddings_slice is not None:
225
- # TODO(yard1): Optimize this copy, we don't need to copy
226
- # everything, just the modified part
227
- embeddings = self.embeddings_tensors.view(
228
- self.embeddings_tensors.shape[0] *
229
- self.embeddings_tensors.shape[1],
230
- self.embeddings_tensors.shape[2],
231
- )[self.embeddings_slice[0]:self.embeddings_slice[1]]
232
- assert self.embeddings_weights is not None
233
- self.embeddings_weights[:embeddings.shape[0]].copy_(embeddings)
234
-
235
- def forward(self, x: torch.Tensor) -> torch.Tensor:
236
- added_tokens_mask = torch.where(x > self.base_layer.org_vocab_size - 1,
237
- 1, 0)
238
-
239
- # NB: Don't use torch.narrow here. torch.narrow triggers some
240
- # Dynamic Shape specialization in torch.compile
241
- num_tokens = x.shape[0]
242
- indices_1 = self.punica_wrapper._embeddings_indices[1][:num_tokens]
243
- indices_0 = self.punica_wrapper._embeddings_indices[0][:num_tokens]
244
-
245
- full_lora_a_embeddings = F.embedding(
246
- x + indices_1,
247
- self.lora_a_stacked_2d,
248
- )
249
- full_output = self.base_layer.forward(x +
250
- (indices_0 * added_tokens_mask))
251
-
252
- full_output_org = full_output
253
- if full_output.ndim == 3:
254
- full_output = full_output.view(
255
- full_output.shape[0] * full_output.shape[1], -1)
256
- if full_lora_a_embeddings.ndim == 3:
257
- full_lora_a_embeddings = full_lora_a_embeddings.view(
258
- full_lora_a_embeddings.shape[0] *
259
- full_lora_a_embeddings.shape[1],
260
- -1,
261
- )
262
-
263
- lora_output: Optional[
264
- torch.Tensor] = self.punica_wrapper.add_lora_embedding(
265
- full_output,
266
- full_lora_a_embeddings,
267
- self.lora_b_stacked,
268
- add_input=True)
269
-
270
- if not current_platform.can_update_inplace():
271
- full_output = lora_output
272
-
273
- return full_output.view_as(full_output_org)
274
-
275
- @classmethod
276
- def can_replace_layer(
277
- cls,
278
- source_layer: nn.Module,
279
- lora_config: LoRAConfig,
280
- packed_modules_list: list,
281
- model_config: Optional[PretrainedConfig],
282
- ) -> bool:
283
- return type(source_layer) is VocabParallelEmbedding
284
-
285
- @property
286
- def weight(self):
287
- return self.base_layer.weight
288
-
289
-
290
-class BaseLinearLayerWithLoRA(BaseLayerWithLoRA):
291
-
292
- def __init__(self, base_layer: LinearBase):
293
- super().__init__()
294
- self.base_layer = base_layer
295
- self.input_size = self.base_layer.input_size
296
- self.device = _get_lora_device(self.base_layer)
297
- self.lora_bias_stacked: Optional[tuple[torch.Tensor, ...]] = None
298
-
299
- self.output_slices: tuple[int, ...]
300
- self.tp_size: int
301
- self.output_size: int
302
- self.n_slices: int
303
-
304
- def create_lora_weights(
305
- self,
306
- max_loras: int,
307
- lora_config: LoRAConfig,
308
- model_config: Optional[PretrainedConfig] = None,
309
- ) -> None:
310
- self.lora_config = lora_config
311
- #
312
- if isinstance(self.base_layer, ReplicatedLinear):
313
- lora_a_out_size = lora_config.max_lora_rank
314
- lora_b_out_size = self.output_size
315
-
316
- elif isinstance(self.base_layer, ColumnParallelLinear):
317
- lora_a_out_size = (lora_config.max_lora_rank if
318
- not lora_config.fully_sharded_loras else divide(
319
- lora_config.max_lora_rank, self.tp_size))
320
- lora_b_out_size = self.output_size
321
-
322
- elif isinstance(self.base_layer, RowParallelLinear):
323
- lora_a_out_size = lora_config.max_lora_rank
324
- lora_b_out_size = (self.output_size if
325
- not lora_config.fully_sharded_loras else divide(
326
- self.output_size, self.tp_size))
327
- else:
328
- raise NotImplementedError
329
-
330
- self.lora_a_stacked = tuple(
331
- torch.zeros(
332
- max_loras,
333
- 1,
334
- lora_a_out_size,
335
- self.input_size,
336
- dtype=lora_config.lora_dtype,
337
- device=self.device,
338
- ) for _ in range(self.n_slices))
339
- self.lora_b_stacked = tuple(
340
- torch.zeros(
341
- max_loras,
342
- 1,
343
- lora_b_out_size,
344
- lora_config.max_lora_rank,
345
- dtype=lora_config.lora_dtype,
346
- device=self.device,
347
- ) for _ in range(self.n_slices))
348
- if lora_config.bias_enabled:
349
- lora_bias_out_size = lora_b_out_size
350
- self.lora_bias_stacked = tuple(
351
- torch.zeros(
352
- max_loras,
353
- 1,
354
- lora_bias_out_size,
355
- dtype=lora_config.lora_dtype,
356
- device=self.device,
357
- ) for _ in range(self.n_slices))
358
- self.output_slices = (self.lora_b_stacked[0].shape[2], )
359
-
360
- def reset_lora(self, index: int):
361
- for s_index in range(self.n_slices):
362
- self.lora_a_stacked[s_index][index] = 0
363
- self.lora_b_stacked[s_index][index] = 0
364
- if self.lora_config.bias_enabled:
365
- # Make mypy happy
366
- self.lora_bias_stacked = cast(tuple[torch.Tensor, ...],
367
- self.lora_bias_stacked)
368
- self.lora_bias_stacked[s_index][index] = 0
369
-
370
- def set_lora(
371
- self,
372
- index: int,
373
- lora_a: torch.Tensor,
374
- lora_b: torch.Tensor,
375
- embeddings_tensor: Optional[torch.Tensor],
376
- lora_bias: Optional[torch.Tensor] = None,
377
- ):
378
- # Except for QKVParallelLinearWithLoRA and
379
- # MergedColumnParallelLinearWithLoRA, all other linear LoRA layers
380
- # store weights in a tuple of size 1. These two layers will
381
- # override this function.
382
- assert (len(self.lora_a_stacked) == len(self.lora_b_stacked) ==
383
- self.n_slices == 1)
384
-
385
- self.reset_lora(index)
386
- if self.tp_size > 1:
387
- lora_a = self.slice_lora_a(lora_a)
388
- lora_b = self.slice_lora_b(lora_b)
389
- if lora_bias is not None:
390
- lora_bias = self.slice_bias(lora_bias)
391
-
392
- self.lora_a_stacked[0][index,
393
- 0, :lora_a.shape[1], :lora_a.shape[0]].copy_(
394
- lora_a.T, non_blocking=True)
395
- self.lora_b_stacked[0][index,
396
- 0, :lora_b.shape[1], :lora_b.shape[0]].copy_(
397
- lora_b.T, non_blocking=True)
398
- if lora_bias is not None:
399
-
400
- self.lora_bias_stacked = cast(tuple[torch.Tensor, ...],
401
- self.lora_bias_stacked)
402
- assert len(self.lora_bias_stacked)
403
- self.lora_bias_stacked[0][index, 0, :lora_bias.shape[0]].copy_(
404
- lora_bias.T, non_blocking=True)
405
-
406
- def apply(self,
407
- x: torch.Tensor,
408
- bias: Optional[torch.Tensor] = None) -> torch.Tensor:
409
- output = self.base_layer.quant_method.apply(self.base_layer, x, bias)
410
-
411
- # In transformers backend, x and output have extra batch dimension like
412
- # (1, seq_len, hidden_dim), while punica expects (seq_len, hidden_dim),
413
- # therefore we need to flatten the batch dimensions.
414
- if x.ndim == 3 and output.ndim == 3:
415
- output = output.flatten(0, 1)
416
- x = x.flatten(0, 1)
417
-
418
- lora_output: Optional[
419
- torch.Tensor] = self.punica_wrapper.add_lora_linear(
420
- output, x, self.lora_a_stacked, self.lora_b_stacked,
421
- self.lora_bias_stacked, 1.0, self.output_slices)
422
- if not current_platform.can_update_inplace():
423
- output = lora_output
424
-
425
- return output
426
-
427
- @property
428
- def weight(self) -> torch.Tensor:
429
-
430
- # unquantizedLinear
431
- if hasattr(self.base_layer, "weight"):
432
- return self.base_layer.weight
433
- # Compressed Tensor
434
- elif hasattr(self.base_layer, "weight_packed"):
435
- return self.base_layer.weight_packed
436
- # GPTQ/AWQ
437
- elif hasattr(self.base_layer, "qweight"):
438
- return self.base_layer.qweight
439
- # marlin
440
- elif hasattr(self.base_layer, "B"):
441
- return self.base_layer.B
442
- # HQQ marlin
443
- elif hasattr(self.base_layer, "W_q"):
444
- return self.base_layer.W_q
445
- else:
446
- raise ValueError(f"Unsupported base layer: {self.base_layer}")
447
-
448
- @property
449
- def bias(self) -> Optional[torch.Tensor]:
450
- if hasattr(self.base_layer, "bias"):
451
- return self.base_layer.bias
452
- else:
453
- return None
454
-
455
-
456
-class ReplicatedLinearWithLoRA(BaseLinearLayerWithLoRA):
457
-
458
- def __init__(self, base_layer: ReplicatedLinear) -> None:
459
- super().__init__(base_layer, )
460
- # To ensure interface compatibility, set to 1 always.
461
- self.tp_size = 1
462
- self.output_size = self.base_layer.output_size
463
- self.n_slices = 1
464
-
465
- def forward(
466
- self, input_: torch.Tensor
467
- ) -> Union[torch.Tensor, tuple[torch.Tensor, Optional[torch.Tensor]]]:
468
- """Forward of ReplicatedLinearWithLoRA
469
-
470
- Args:
471
- input_: Tensor whose last dimension is `input_size`.
472
-
473
- Returns:
474
- - output
475
- - bias
476
- """
477
- bias = (self.base_layer.bias
478
- if not self.base_layer.skip_bias_add else None)
479
-
480
- # Matrix multiply.
481
- output = self.apply(input_, bias)
482
-
483
- output_bias = (self.base_layer.bias
484
- if self.base_layer.skip_bias_add else None)
485
-
486
- if not self.base_layer.return_bias:
487
- return output
488
-
489
- return output, output_bias
490
-
491
- # ReplicatedLinear should always be replaced, regardless of the fully
492
- # sharded LoRAs setting, because it is, by definition, copied per GPU.
493
- @classmethod
494
- def can_replace_layer(
495
- cls,
496
- source_layer: nn.Module,
497
- lora_config: LoRAConfig,
498
- packed_modules_list: list,
499
- model_config: Optional[PretrainedConfig],
500
- ) -> bool:
501
- return type(source_layer) is ReplicatedLinear
502
-
503
-
504
-class ColumnParallelLinearWithLoRA(BaseLinearLayerWithLoRA):
505
- """
506
- LoRA on top of ColumnParallelLinear layer.
507
- LoRA B is sliced for tensor parallelism.
508
- There are two types for the `base_layer`:
509
- 1. ColumnParallelLinear, e.g.`dense_h_to_4h` in `FalconForCausalLM`.
510
- 2. MergedColumnParallelLinear, e.g.`gate_up_proj` in `Phi3ForCausalLM`.
511
- """
512
-
513
- def __init__(self, base_layer: ColumnParallelLinear) -> None:
514
- super().__init__(base_layer)
515
- # The base_layer type is ColumnParallelLinear or
516
- # MergedColumnParallelLinear, their weight sharding logic is
517
- # inconsistent when TP is greater than 1.
518
- self.is_merged_col_linear = type(
519
- base_layer) is MergedColumnParallelLinear
520
- self.tp_size = get_tensor_model_parallel_world_size()
521
- self.output_size = self.base_layer.output_size_per_partition
522
- # There is only one LoRA layer
523
- self.n_slices = 1
524
-
525
- def slice_lora_a(self, lora_a: torch.Tensor) -> torch.Tensor:
526
- return lora_a
527
-
528
- def slice_lora_b(self, lora_b: torch.Tensor) -> torch.Tensor:
529
- # Applicable to cases where the base_layer is
530
- # MergedColumnParallelLinear.
531
- if self.is_merged_col_linear:
532
- tp_rank = get_tensor_model_parallel_rank()
533
- shard_size = self.output_size // 2
534
- offset = lora_b.shape[-1] // 2
535
-
536
- left_weight = lora_b[:, tp_rank * shard_size:(tp_rank + 1) *
537
- shard_size]
538
- right_weight = lora_b[:, offset + tp_rank * shard_size:offset +
539
- (tp_rank + 1) * shard_size]
540
- lora_b = torch.cat([left_weight, right_weight], dim=1)
541
- # Applicable to cases where the base_layer is
542
- # ColumnParallelLinear.
543
- else:
544
- tensor_model_parallel_rank = get_tensor_model_parallel_rank()
545
- shard_size = self.output_size
546
- start_idx = tensor_model_parallel_rank * shard_size
547
- end_idx = (tensor_model_parallel_rank + 1) * shard_size
548
- lora_b = lora_b[:, start_idx:end_idx]
549
- return lora_b
550
-
551
- def slice_bias(self, bias: torch.Tensor) -> torch.Tensor:
552
- # TODO: Fix the slicing logic of bias.
553
- if bias is None:
554
- return bias
555
- tensor_model_parallel_rank = get_tensor_model_parallel_rank()
556
- shard_size = self.output_size
557
- start_idx = tensor_model_parallel_rank * shard_size
558
- end_idx = (tensor_model_parallel_rank + 1) * shard_size
559
- bias = bias[start_idx:end_idx]
560
- return bias
561
-
562
- def forward(
563
- self, input_: torch.Tensor
564
- ) -> Union[torch.Tensor, tuple[torch.Tensor, Optional[torch.Tensor]]]:
565
- """Forward of ColumnParallelLinear
566
-
567
- Args:
568
- input_: Tensor whose last dimension is `input_size`.
569
-
570
- Returns:
571
- - output
572
- - bias
573
- """
574
- bias = (self.base_layer.bias
575
- if not self.base_layer.skip_bias_add else None)
576
-
577
- # Matrix multiply.
578
- output_parallel = self.apply(input_, bias)
579
- if self.base_layer.gather_output:
580
- # All-gather across the partitions.
581
- output = tensor_model_parallel_all_gather(output_parallel)
582
- else:
583
- output = output_parallel
584
-
585
- if not self.base_layer.return_bias:
586
- return output
587
-
588
- output_bias = (self.base_layer.bias
589
- if self.base_layer.skip_bias_add else None)
590
- return output, output_bias
591
-
592
- @classmethod
593
- @_not_fully_sharded_can_replace
594
- def can_replace_layer(
595
- cls,
596
- source_layer: nn.Module,
597
- lora_config: LoRAConfig,
598
- packed_modules_list: list,
599
- model_config: Optional[PretrainedConfig],
600
- ) -> bool:
601
- return type(source_layer) is ColumnParallelLinear or (
602
- type(source_layer) is MergedColumnParallelLinear
603
- and len(packed_modules_list) == 1)
604
-
605
-
606
-class MergedColumnParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):
607
- """ColumnParallelLinear layer that is composed of 2 sublayers (slices)
608
- packed together (e.g. gate_proj + up_proj -> gate_up_proj).
609
-
610
- This means we have 2 LoRAs, each applied to one half of the layer.
611
-
612
- Both slices must have the same size.
613
- """
614
-
615
- def __init__(
616
- self, base_layer: Union[MergedColumnParallelLinear,
617
- QKVParallelLinear]) -> None:
618
- super().__init__(base_layer)
619
- # There are two LoRA layers
620
- self.tp_size = get_tensor_model_parallel_world_size()
621
- self.tp_rank = get_tensor_model_parallel_rank()
622
- # the output_sizes in MergedColumnParallelLinear is not sharded by tp
623
- # we need to divide it by the tp_size to get correct slices size
624
- output_sizes = self.base_layer.output_sizes
625
- self.output_slices = tuple(
626
- divide(output_size, self.tp_size) for output_size in output_sizes)
627
- self.n_slices = len(self.output_slices)
628
- self.output_ids = (self.tp_rank, ) * self.n_slices
629
-
630
- def create_lora_weights(
631
- self,
632
- max_loras: int,
633
- lora_config: LoRAConfig,
634
- model_config: Optional[PretrainedConfig] = None,
635
- ) -> None:
636
- """
637
- The main reason for overriding this function is to enhance code
638
- maintainability.
639
- """
640
- self.lora_config = lora_config
641
-
642
- lora_a_output_size_per_partition = (
643
- lora_config.max_lora_rank if not lora_config.fully_sharded_loras
644
- else divide(lora_config.max_lora_rank, self.tp_size))
645
-
646
- self.lora_a_stacked = tuple(
647
- torch.zeros(
648
- max_loras,
649
- 1,
650
- lora_a_output_size_per_partition,
651
- self.input_size,
652
- dtype=lora_config.lora_dtype,
653
- device=self.device,
654
- ) for _ in range(self.n_slices))
655
- self.lora_b_stacked = tuple(
656
- torch.zeros(
657
- max_loras,
658
- 1,
659
- output_size,
660
- lora_config.max_lora_rank,
661
- dtype=lora_config.lora_dtype,
662
- device=self.device,
663
- ) for output_size in self.output_slices)
664
- if lora_config.bias_enabled:
665
- self.lora_bias_stacked = tuple(
666
- torch.zeros(
667
- max_loras,
668
- 1,
669
- output_size,
670
- dtype=lora_config.lora_dtype,
671
- device=self.device,
672
- ) for output_size in self.output_slices)
673
-
674
- def slice_lora_a(
675
- self, lora_a: list[Union[torch.Tensor, None]]
676
- ) -> list[Union[torch.Tensor, None]]:
677
- return lora_a
678
-
679
- def slice_lora_b(
680
- self, lora_b: list[Union[torch.Tensor, None]]
681
- ) -> list[Union[torch.Tensor, None]]:
682
- sliced_lora_b = [None] * self.n_slices
683
- for i, (shard_id, shard_size) in enumerate(
684
- zip(self.output_ids, self.output_slices)):
685
- if (lora_b_i := lora_b[i]) is not None:
686
- sliced_lora_b[i] = lora_b_i[:,
687
- shard_size * shard_id:shard_size *
688
- (shard_id + 1)]
689
- return sliced_lora_b
690
-
691
- def slice_bias(
692
- self, bias: list[Union[torch.Tensor,
693
- None]]) -> list[Union[torch.Tensor, None]]:
694
- for i, (shard_id, shard_size) in enumerate(
695
- zip(self.output_ids, self.output_slices)):
696
- if (bias_i := bias[i]) is not None:
697
- bias[i] = bias_i[shard_size * shard_id:shard_size *
698
- (shard_id + 1)]
699
- return bias
700
-
701
- def set_lora(
702
- self,
703
- index: int,
704
- lora_a: torch.Tensor,
705
- lora_b: torch.Tensor,
706
- embeddings_tensor: Optional[torch.Tensor],
707
- lora_bias: Optional[torch.Tensor] = None,
708
- ):
709
- self.reset_lora(index)
710
-
711
- if self.tp_size > 1:
712
- lora_a = self.slice_lora_a(lora_a)
713
- lora_b = self.slice_lora_b(lora_b)
714
- if lora_bias is not None:
715
- lora_bias = self.slice_bias(lora_bias)
716
-
717
- for i in range(self.n_slices):
718
- if (lora_a_i := lora_a[i]) is not None:
719
- self.lora_a_stacked[i][
720
- index, 0, :lora_a_i.shape[1], :lora_a_i.shape[0]].copy_(
721
- lora_a_i.T, non_blocking=True)
722
- if (lora_b_i := lora_b[i]) is not None:
723
- self.lora_b_stacked[i][
724
- index, 0, :lora_b_i.shape[1], :lora_b_i.shape[0]].copy_(
725
- lora_b_i.T, non_blocking=True)
726
-
727
- if lora_bias is not None:
728
- self.lora_bias_stacked = cast(tuple[torch.Tensor, ...],
729
- self.lora_bias_stacked)
730
- for i in range(self.n_slices):
731
- if (lora_bias_i := lora_bias[i]) is not None:
732
- self.lora_bias_stacked[i][index,
733
- 0, :lora_bias_i.shape[0]].copy_(
734
- lora_bias_i.T,
735
- non_blocking=True)
736
-
737
- @classmethod
738
- @_not_fully_sharded_can_replace
739
- def can_replace_layer(
740
- cls,
741
- source_layer: nn.Module,
742
- lora_config: LoRAConfig,
743
- packed_modules_list: list,
744
- model_config: Optional[PretrainedConfig],
745
- ) -> bool:
746
- return (type(source_layer) is MergedColumnParallelLinear
747
- and len(packed_modules_list) == 2)
748
-
749
-
750
-class QKVParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):
751
- """
752
- ColumnParallelLinear layer that is specifically designed for
753
- qkv_proj. Certain models, such as chatglm3 and baichuan-7b,
754
- only contains a single LoRA within their qkv_proj layer.
755
-
756
- During inference with Tensor Parallel, the weights of lora_b
757
- must be accurately partitioned according to the respective ranks.
758
-
759
- Q slice may have different shape than K and V slices (which both have
760
- the same shape).
761
- """
762
-
763
- def __init__(self, base_layer: QKVParallelLinear) -> None:
764
- super().__init__(base_layer)
765
- self.q_proj_total_size = (self.base_layer.total_num_heads *
766
- self.base_layer.head_size)
767
- self.q_proj_shard_size = (self.base_layer.num_heads *
768
- self.base_layer.head_size)
769
- self.kv_proj_shard_size = (self.base_layer.num_kv_heads *
770
- self.base_layer.head_size)
771
- self.kv_proj_total_size = (self.base_layer.total_num_kv_heads *
772
- self.base_layer.head_size)
773
- # There is only one LoRA layer
774
- self.n_slices = 1
775
-
776
- def slice_lora_b(self, lora_b: torch.Tensor) -> torch.Tensor:
777
- tp_rank = get_tensor_model_parallel_rank()
778
- self.q_shard_id = tp_rank
779
- self.kv_shard_id = tp_rank // self.base_layer.num_kv_head_replicas
780
- lora_b_q = lora_b[:, self.q_proj_shard_size *
781
- self.q_shard_id:self.q_proj_shard_size *
782
- (self.q_shard_id + 1)]
783
- k_offset = self.q_proj_total_size
784
- lora_b_k = lora_b[:, k_offset +
785
- self.kv_proj_shard_size * self.kv_shard_id:k_offset +
786
- self.kv_proj_shard_size * (self.kv_shard_id + 1)]
787
- v_offset = k_offset + self.kv_proj_total_size
788
- lora_b_v = lora_b[:, v_offset +
789
- self.kv_proj_shard_size * self.kv_shard_id:v_offset +
790
- self.kv_proj_shard_size * (self.kv_shard_id + 1)]
791
- lora_b = torch.cat([lora_b_q, lora_b_k, lora_b_v], dim=1)
792
- return lora_b
793
-
794
- def slice_bias(self, bias: torch.Tensor) -> torch.Tensor:
795
- bias_q = bias[self.q_proj_shard_size *
796
- self.q_shard_id:self.q_proj_shard_size *
797
- (self.q_shard_id + 1)]
798
- k_offset = self.q_proj_total_size
799
- bias_k = bias[k_offset +
800
- self.kv_proj_shard_size * self.kv_shard_id:k_offset +
801
- self.kv_proj_shard_size * (self.kv_shard_id + 1)]
802
- v_offset = k_offset + self.kv_proj_total_size
803
- bias_v = bias[v_offset +
804
- self.kv_proj_shard_size * self.kv_shard_id:v_offset +
805
- self.kv_proj_shard_size * (self.kv_shard_id + 1)]
806
- bias = torch.cat([bias_q, bias_k, bias_v], dim=1)
807
- return bias
808
-
809
- @classmethod
810
- @_not_fully_sharded_can_replace
811
- def can_replace_layer(cls, source_layer: nn.Module,
812
- lora_config: LoRAConfig, packed_modules_list: list,
813
- model_config: Optional[PretrainedConfig]) -> bool:
814
- return type(source_layer) is QKVParallelLinear and len(
815
- packed_modules_list) == 1
816
-
817
-
818
-class MergedQKVParallelLinearWithLoRA(MergedColumnParallelLinearWithLoRA):
819
- """MergedColumnParallelLinear layer that is composed of 3 sublayers (slices)
820
- packed together in qkv proj fashion
821
- (q_proj + k_proj + v_proj -> qkv_proj).
822
-
823
- This means we have 3 LoRAs, each applied to one slice of the layer.
824
-
825
- Q slice may have different shape than K and V slices (which both have
826
- the same shape).
827
- """
828
-
829
- def __init__(self, base_layer: QKVParallelLinear) -> None:
830
- super().__init__(base_layer)
831
- # There are three LoRA layer.
832
- self.n_slices = len(self.base_layer.output_sizes)
833
- self.tp_size = get_tensor_model_parallel_world_size()
834
- self.tp_rank = get_tensor_model_parallel_rank()
835
-
836
- self.q_proj_shard_size = (self.base_layer.num_heads *
837
- self.base_layer.head_size)
838
- self.kv_proj_shard_size = (self.base_layer.num_kv_heads *
839
- self.base_layer.head_size)
840
- self.q_shard_id = self.tp_rank
841
- self.kv_shard_id = self.tp_rank // self.base_layer.num_kv_head_replicas
842
-
843
- self.output_slices = (
844
- self.q_proj_shard_size,
845
- self.kv_proj_shard_size,
846
- self.kv_proj_shard_size,
847
- )
848
- self.output_ids = (
849
- self.q_shard_id,
850
- self.kv_shard_id,
851
- self.kv_shard_id,
852
- )
853
-
854
- def create_lora_weights(
855
- self,
856
- max_loras: int,
857
- lora_config: LoRAConfig,
858
- model_config: Optional[PretrainedConfig] = None,
859
- ) -> None:
860
- """
861
- The main reason for overloading this function is to handle inconsistent
862
- weight dimensions in qkv lora.
863
- """
864
- super().create_lora_weights(max_loras, lora_config, model_config)
865
-
866
- @classmethod
867
- @_not_fully_sharded_can_replace
868
- def can_replace_layer(
869
- cls,
870
- source_layer: nn.Module,
871
- lora_config: LoRAConfig,
872
- packed_modules_list: list,
873
- model_config: Optional[PretrainedConfig],
874
- ) -> bool:
875
- return (type(source_layer) is QKVParallelLinear
876
- and len(packed_modules_list) == 3)
877
-
878
-
879
-#TODO: Implement this
880
-class QKVCrossParallelLinearWithLoRA(BaseLayerWithLoRA):
881
- pass
882
-
883
-
884
-class RowParallelLinearWithLoRA(BaseLinearLayerWithLoRA):
885
-
886
- def __init__(self, base_layer: RowParallelLinear) -> None:
887
- super().__init__(base_layer)
888
-
889
- self.tp_size = get_tensor_model_parallel_world_size()
890
- # reset input_size
891
- self.input_size = self.base_layer.input_size_per_partition
892
- self.output_size = self.base_layer.output_size
893
-
894
- self.tp_rank = get_tensor_model_parallel_rank()
895
- # There is only one LoRA layer.
896
- self.n_slices = 1
897
-
898
- def slice_lora_a(self, lora_a: torch.Tensor) -> torch.Tensor:
899
-
900
- shard_size = self.input_size
901
- start_idx = self.tp_rank * shard_size
902
- end_idx = (self.tp_rank + 1) * shard_size
903
- lora_a = lora_a[start_idx:end_idx, :]
904
- return lora_a
905
-
906
- def slice_lora_b(self, lora_b: torch.Tensor) -> torch.Tensor:
907
- return lora_b
908
-
909
- def slice_bias(self, bias: torch.Tensor) -> torch.Tensor:
910
- return bias
911
-
912
- def forward(
913
- self, input_: torch.Tensor
914
- ) -> Union[torch.Tensor, tuple[torch.Tensor, Optional[torch.Tensor]]]:
915
- """Forward of RowParallelLinear
916
-
917
- Args:
918
- input_: tensor whose last dimension is `input_size`. If
919
- `input_is_parallel` is set, then the last dimension
920
- is `input_size // tp_size`.
921
-
922
- Returns:
923
- - output
924
- - bias
925
- """
926
- # set up backprop all-reduce.
927
- if self.base_layer.input_is_parallel:
928
- input_parallel = input_
929
- else:
930
- # TODO: simplify code below
931
- splitted_input = split_tensor_along_last_dim(
932
- input_, num_partitions=self.base_layer.tp_size)
933
- input_parallel = splitted_input[self.tp_rank].contiguous()
934
-
935
- # Matrix multiply.
936
- output_parallel = self.apply(input_parallel)
937
- if self.base_layer.reduce_results and self.base_layer.tp_size > 1:
938
- output_ = tensor_model_parallel_all_reduce(output_parallel)
939
- else:
940
- output_ = output_parallel
941
-
942
- if not self.base_layer.skip_bias_add:
943
- output = (output_ + self.base_layer.bias
944
- if self.base_layer.bias is not None else output_)
945
- output_bias = None
946
- else:
947
- output = output_
948
- output_bias = self.base_layer.bias
949
-
950
- if not self.base_layer.return_bias:
951
- return output
952
-
953
- return output, output_bias
954
-
955
- @classmethod
956
- @_not_fully_sharded_can_replace
957
- def can_replace_layer(
958
- cls,
959
- source_layer: nn.Module,
960
- lora_config: LoRAConfig,
961
- packed_modules_list: list,
962
- model_config: Optional[PretrainedConfig],
963
- ) -> bool:
964
- return type(source_layer) is RowParallelLinear
965
-
966
-
967
-class LogitsProcessorWithLoRA(BaseLayerWithLoRA):
968
- """
969
- LoRA wrapper for LogitsProcessor, with extra logic to handle the
970
- application of the LoRA adapter and added LoRA vocabulary.
971
-
972
- Args:
973
- base_layer: LogitsProcessor layer
974
- hidden_size: hidden size of the model
975
- dtype: data type of the model
976
- device: device of the model
977
- sharded_to_full_mapping: index mapping from sharded vocab to full vocab
978
- received from base_layer.get_sharded_to_full_mapping(). If None,
979
- no reindexing will be done.
980
- """
981
-
982
- def __init__(self, base_layer: LogitsProcessor, hidden_size: int,
983
- dtype: torch.dtype, device: torch.device,
984
- sharded_to_full_mapping: Optional[list[int]]) -> None:
985
- super().__init__()
986
- self.base_layer = base_layer
987
- self.hidden_size = hidden_size
988
- self.dtype = dtype
989
- self.device = device
990
- self.tp_size = get_tensor_model_parallel_world_size()
991
- self.tp_rank = get_tensor_model_parallel_rank()
992
- self.sharded_to_full_mapping = sharded_to_full_mapping
993
-
994
- @property
995
- def logits_as_input(self):
996
- return self.base_layer.logits_as_input
997
-
998
- @property
999
- def vocab_size(self):
1000
- return self.base_layer.vocab_size
1001
-
1002
- @property
1003
- def scale(self):
1004
- return self.base_layer.scale
1005
-
1006
- @property
1007
- def soft_cap(self):
1008
- return self.base_layer.soft_cap
1009
-
1010
- @property
1011
- def use_all_gather(self):
1012
- return self.base_layer.use_all_gather
1013
-
1014
- @property
1015
- def org_vocab_size(self):
1016
- return self.base_layer.org_vocab_size
1017
-
1018
- @property
1019
- def include_gpu_probs_tensor(self):
1020
- return self.base_layer.include_gpu_probs_tensor
1021
-
1022
- @property
1023
- def should_modify_greedy_probs_inplace(self):
1024
- return self.base_layer.should_modify_greedy_probs_inplace
1025
-
1026
- def create_lora_weights(
1027
- self,
1028
- max_loras: int,
1029
- lora_config: LoRAConfig,
1030
- model_config: Optional[PretrainedConfig] = None,
1031
- ) -> None:
1032
- # TODO: Verify if this condition can be further relaxed
1033
- if 32000 < self.base_layer.vocab_size > 257024:
1034
- raise ValueError("When using LoRA, vocab size must be "
1035
- "32000 >= vocab_size <= 257024")
1036
- self.lora_a_stacked = torch.zeros(
1037
- (
1038
- max_loras,
1039
- 1,
1040
- lora_config.max_lora_rank,
1041
- self.hidden_size,
1042
- ),
1043
- dtype=lora_config.lora_dtype,
1044
- device=self.device,
1045
- )
1046
- self.lora_b_stacked = torch.zeros(
1047
- (
1048
- max_loras,
1049
- 1,
1050
- # Pad for kernel compatibility
1051
- math.ceil(self.base_layer.vocab_size /
1052
- lora_config.lora_vocab_padding_size) *
1053
- lora_config.lora_vocab_padding_size,
1054
- lora_config.max_lora_rank,
1055
- ),
1056
- dtype=lora_config.lora_dtype,
1057
- device=self.device,
1058
- )
1059
- self.embeddings_tensors = torch.full(
1060
- (max_loras, lora_config.lora_extra_vocab_size, self.hidden_size),
1061
- fill_value=float("-inf"),
1062
- dtype=self.dtype,
1063
- device=self.device,
1064
- )
1065
- if self.sharded_to_full_mapping is not None:
1066
- self.sharded_to_full_mapping_gpu = torch.tensor(
1067
- self.sharded_to_full_mapping,
1068
- device=self.device,
1069
- dtype=torch.long)
1070
- else:
1071
- self.sharded_to_full_mapping_gpu = None
1072
-
1073
- def reset_lora(self, index: int):
1074
- self.lora_a_stacked[index] = 0
1075
- self.lora_b_stacked[index] = 0
1076
- self.embeddings_tensors[index] = float("-inf")
1077
-
1078
- def set_lora(
1079
- self,
1080
- index: int,
1081
- lora_a: torch.Tensor,
1082
- lora_b: torch.Tensor,
1083
- embeddings_tensor: Optional[torch.Tensor],
1084
- bias: Optional[torch.Tensor] = None,
1085
- ):
1086
- self.reset_lora(index)
1087
- self.lora_a_stacked[index,
1088
- 0, :lora_a.shape[1], :lora_a.shape[0]].copy_(
1089
- lora_a.T, non_blocking=True)
1090
- self.lora_b_stacked[index,
1091
- 0, :lora_b.shape[1], :lora_b.shape[0]].copy_(
1092
- lora_b.T, non_blocking=True)
1093
- if embeddings_tensor is not None:
1094
- self.embeddings_tensors[
1095
- index,
1096
- :embeddings_tensor.shape[0],
1097
- :embeddings_tensor.shape[1],
1098
- ] = embeddings_tensor
1099
-
1100
- def _get_logits(
1101
- self,
1102
- hidden_states: torch.Tensor,
1103
- lm_head: VocabParallelEmbedding,
1104
- embedding_bias: Optional[torch.Tensor] = None,
1105
- ) -> Optional[torch.Tensor]:
1106
- # Get the logits for the next tokens.
1107
- logits = lm_head.quant_method.apply(lm_head, hidden_states)
1108
- if embedding_bias is not None:
1109
- logits += embedding_bias
1110
-
1111
- # Gather logits for TP
1112
- logits = self.base_layer._gather_logits(logits)
1113
-
1114
- if logits is None:
1115
- return None
1116
-
1117
- if self.sharded_to_full_mapping_gpu is not None:
1118
- # Reindex full logits tensor to ensure 1:1 mapping between
1119
- # index and token_id
1120
- # Example for:
1121
- # org_vocab_size = 4
1122
- # added_vocab_size = 2
1123
- # pad_to_size = 8
1124
- # tp_size = 2
1125
-
1126
- # indices: [0, 1, 2, 3, 4, 5, 6, 7]
1127
- # token_id: [0, 1, 4, -1, 2, 3, 5, -1]
1128
-
1129
- # Therefore, the mapping is expected to be:
1130
- # [0, 1, 4, 6, 2, 3, 5, 7] so that when we reindex,
1131
- # we get:
1132
- # indices: [0, 1, 2, 3, 4, 5, 6, 7]
1133
- # token_id: [0, 1, 2, 3, 4, 5, -1, -1]
1134
- logits = logits[:, self.sharded_to_full_mapping_gpu]
1135
-
1136
- lora_logits = torch.empty(
1137
- self.embeddings_tensors.shape[0] + 1,
1138
- self.embeddings_tensors.shape[1],
1139
- hidden_states.shape[0],
1140
- dtype=self.embeddings_tensors.dtype,
1141
- device=self.embeddings_tensors.device,
1142
- )
1143
- torch.matmul(self.embeddings_tensors,
1144
- hidden_states.T,
1145
- out=lora_logits[:-1])
1146
-
1147
- neg_inf, pos_inf = current_platform.get_infinity_values(
1148
- lora_logits.dtype)
1149
-
1150
- lora_logits[-1] = neg_inf
1151
- lora_logits = lora_logits.mT
1152
- indices_padded = self.punica_wrapper.sampler_indices_padded
1153
-
1154
- if current_platform.is_tpu() or current_platform.is_xpu():
1155
- indices_padded = indices_padded[:logits.size(0)]
1156
-
1157
- lora_logits = (lora_logits.reshape(
1158
- lora_logits.shape[0] * lora_logits.shape[1],
1159
- lora_logits.shape[2],
1160
- ).index_select(0, indices_padded).nan_to_num_(nan=neg_inf,
1161
- posinf=pos_inf,
1162
- neginf=neg_inf))
1163
-
1164
- logits[:,
1165
- self.base_layer.org_vocab_size:self.base_layer.org_vocab_size +
1166
- lora_logits.shape[1]] = lora_logits
1167
-
1168
- lora_output: Optional[
1169
- torch.Tensor] = self.punica_wrapper.add_lora_logits(
1170
- logits, hidden_states, self.lora_a_stacked,
1171
- self.lora_b_stacked, 1.0)
1172
-
1173
- if not current_platform.can_update_inplace():
1174
- logits = lora_output
1175
-
1176
- # Remove paddings in vocab (if any).
1177
- logits = logits[:, :self.base_layer.vocab_size]
1178
- return logits
1179
-
1180
- def forward(self, *args, **kwargs):
1181
- return type(self.base_layer).forward(self, *args, **kwargs)
1182
-
1183
- @classmethod
1184
- def can_replace_layer(
1185
- cls,
1186
- source_layer: nn.Module,
1187
- lora_config: LoRAConfig,
1188
- packed_modules_list: list,
1189
- model_config: Optional[PretrainedConfig],
1190
- ) -> bool:
1191
- # Special handling for the LogitsProcessor.
1192
- return False
vllm/lora/layers/__init__.py
new
+34
@@ -0,0 +1,34 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from vllm.lora.layers.base import BaseLayerWithLoRA
4
+from vllm.lora.layers.column_parallel_linear import (
5
+ ColumnParallelLinearWithLoRA, ColumnParallelLinearWithShardedLoRA,
6
+ MergedColumnParallelLinearWithLoRA,
7
+ MergedColumnParallelLinearWithShardedLoRA, MergedQKVParallelLinearWithLoRA,
8
+ MergedQKVParallelLinearWithShardedLoRA, QKVParallelLinearWithLoRA,
9
+ QKVParallelLinearWithShardedLoRA)
10
+from vllm.lora.layers.logits_processor import LogitsProcessorWithLoRA
11
+from vllm.lora.layers.replicated_linear import ReplicatedLinearWithLoRA
12
+from vllm.lora.layers.row_parallel_linear import (
13
+ RowParallelLinearWithLoRA, RowParallelLinearWithShardedLoRA)
14
+from vllm.lora.layers.utils import LoRAMapping
15
+from vllm.lora.layers.vocal_parallel_embedding import (
16
+ VocabParallelEmbeddingWithLoRA)
17
+
18
+__all__ = [
19
+ "BaseLayerWithLoRA",
20
+ "VocabParallelEmbeddingWithLoRA",
21
+ "LogitsProcessorWithLoRA",
22
+ "ColumnParallelLinearWithLoRA",
23
+ "ColumnParallelLinearWithShardedLoRA",
24
+ "MergedColumnParallelLinearWithLoRA",
25
+ "MergedColumnParallelLinearWithShardedLoRA",
26
+ "MergedQKVParallelLinearWithLoRA",
27
+ "MergedQKVParallelLinearWithShardedLoRA",
28
+ "QKVParallelLinearWithLoRA",
29
+ "QKVParallelLinearWithShardedLoRA",
30
+ "RowParallelLinearWithLoRA",
31
+ "RowParallelLinearWithShardedLoRA",
32
+ "ReplicatedLinearWithLoRA",
33
+ "LoRAMapping",
34
+]
vllm/lora/layers/base.py
new
+69
@@ -0,0 +1,69 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from typing import TYPE_CHECKING, Optional, Union
5
+
6
+import torch
7
+import torch.nn as nn
8
+from transformers import PretrainedConfig
9
+
10
+from vllm.config import LoRAConfig
11
+
12
+if TYPE_CHECKING:
13
+ from vllm.lora.punica_wrapper import PunicaWrapperBase
14
+
15
+
16
+class BaseLayerWithLoRA(nn.Module):
17
+
18
+ def slice_lora_a(
19
+ self, lora_a: Union[torch.Tensor, list[Union[torch.Tensor, None]]]
20
+ ) -> Union[torch.Tensor, list[Union[torch.Tensor, None]]]:
21
+ """Slice lora a if splitting for tensor parallelism."""
22
+ ...
23
+
24
+ def slice_lora_b(
25
+ self, lora_b: Union[torch.Tensor, list[Union[torch.Tensor, None]]]
26
+ ) -> Union[torch.Tensor, list[Union[torch.Tensor, None]]]:
27
+ """Slice lora b if splitting with tensor parallelism."""
28
+ ...
29
+
30
+ def create_lora_weights(
31
+ self,
32
+ max_loras: int,
33
+ lora_config: LoRAConfig,
34
+ model_config: Optional[PretrainedConfig] = None,
35
+ ) -> None:
36
+ """Initializes lora matrices."""
37
+ ...
38
+
39
+ def reset_lora(self, index: int):
40
+ """Resets the lora weights at index back to 0."""
41
+ ...
42
+
43
+ def set_lora(
44
+ self,
45
+ index: int,
46
+ lora_a: torch.Tensor,
47
+ lora_b: torch.Tensor,
48
+ embeddings_tensor: Optional[torch.Tensor],
49
+ bias: Optional[torch.Tensor] = None,
50
+ ):
51
+ """Overwrites lora tensors at index."""
52
+ ...
53
+
54
+ def set_mapping(
55
+ self,
56
+ punica_wrapper,
57
+ ):
58
+ self.punica_wrapper: PunicaWrapperBase = punica_wrapper
59
+
60
+ @classmethod
61
+ def can_replace_layer(
62
+ cls,
63
+ source_layer: nn.Module,
64
+ lora_config: LoRAConfig,
65
+ packed_modules_list: list,
66
+ model_config: Optional[PretrainedConfig],
67
+ ) -> bool:
68
+ """Returns True if the layer can be replaced by this LoRA layer."""
69
+ raise NotImplementedError
vllm/lora/layers/base_linear.py
new
+184
@@ -0,0 +1,184 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from typing import Optional, cast
5
+
6
+import torch
7
+from transformers import PretrainedConfig
8
+
9
+from vllm.config import LoRAConfig
10
+from vllm.distributed.utils import divide
11
+# yapf: disable
12
+from vllm.model_executor.layers.linear import (ColumnParallelLinear,
13
+ LinearBase, ReplicatedLinear,
14
+ RowParallelLinear)
15
+from vllm.platforms import current_platform
16
+
17
+from .base import BaseLayerWithLoRA
18
+from .utils import _get_lora_device
19
+
20
+
21
+class BaseLinearLayerWithLoRA(BaseLayerWithLoRA):
22
+
23
+ def __init__(self, base_layer: LinearBase):
24
+ super().__init__()
25
+ self.base_layer = base_layer
26
+ self.input_size = self.base_layer.input_size
27
+ self.device = _get_lora_device(self.base_layer)
28
+ self.lora_bias_stacked: Optional[tuple[torch.Tensor, ...]] = None
29
+
30
+ self.output_slices: tuple[int, ...]
31
+ self.tp_size: int
32
+ self.output_size: int
33
+ self.n_slices: int
34
+
35
+ def create_lora_weights(
36
+ self,
37
+ max_loras: int,
38
+ lora_config: LoRAConfig,
39
+ model_config: Optional[PretrainedConfig] = None,
40
+ ) -> None:
41
+ self.lora_config = lora_config
42
+ #
43
+ if isinstance(self.base_layer, ReplicatedLinear):
44
+ lora_a_out_size = lora_config.max_lora_rank
45
+ lora_b_out_size = self.output_size
46
+
47
+ elif isinstance(self.base_layer, ColumnParallelLinear):
48
+ lora_a_out_size = (lora_config.max_lora_rank if
49
+ not lora_config.fully_sharded_loras else divide(
50
+ lora_config.max_lora_rank, self.tp_size))
51
+ lora_b_out_size = self.output_size
52
+
53
+ elif isinstance(self.base_layer, RowParallelLinear):
54
+ lora_a_out_size = lora_config.max_lora_rank
55
+ lora_b_out_size = (self.output_size if
56
+ not lora_config.fully_sharded_loras else divide(
57
+ self.output_size, self.tp_size))
58
+ else:
59
+ raise NotImplementedError
60
+
61
+ self.lora_a_stacked = tuple(
62
+ torch.zeros(
63
+ max_loras,
64
+ 1,
65
+ lora_a_out_size,
66
+ self.input_size,
67
+ dtype=lora_config.lora_dtype,
68
+ device=self.device,
69
+ ) for _ in range(self.n_slices))
70
+ self.lora_b_stacked = tuple(
71
+ torch.zeros(
72
+ max_loras,
73
+ 1,
74
+ lora_b_out_size,
75
+ lora_config.max_lora_rank,
76
+ dtype=lora_config.lora_dtype,
77
+ device=self.device,
78
+ ) for _ in range(self.n_slices))
79
+ if lora_config.bias_enabled:
80
+ lora_bias_out_size = lora_b_out_size
81
+ self.lora_bias_stacked = tuple(
82
+ torch.zeros(
83
+ max_loras,
84
+ 1,
85
+ lora_bias_out_size,
86
+ dtype=lora_config.lora_dtype,
87
+ device=self.device,
88
+ ) for _ in range(self.n_slices))
89
+ self.output_slices = (self.lora_b_stacked[0].shape[2], )
90
+
91
+ def reset_lora(self, index: int):
92
+ for s_index in range(self.n_slices):
93
+ self.lora_a_stacked[s_index][index] = 0
94
+ self.lora_b_stacked[s_index][index] = 0
95
+ if self.lora_config.bias_enabled:
96
+ # Make mypy happy
97
+ self.lora_bias_stacked = cast(tuple[torch.Tensor, ...],
98
+ self.lora_bias_stacked)
99
+ self.lora_bias_stacked[s_index][index] = 0
100
+
101
+ def set_lora(
102
+ self,
103
+ index: int,
104
+ lora_a: torch.Tensor,
105
+ lora_b: torch.Tensor,
106
+ embeddings_tensor: Optional[torch.Tensor],
107
+ lora_bias: Optional[torch.Tensor] = None,
108
+ ):
109
+ # Except for QKVParallelLinearWithLoRA and
110
+ # MergedColumnParallelLinearWithLoRA, all other linear LoRA layers
111
+ # store weights in a tuple of size 1. These two layers will
112
+ # override this function.
113
+ assert (len(self.lora_a_stacked) == len(self.lora_b_stacked) ==
114
+ self.n_slices == 1)
115
+
116
+ self.reset_lora(index)
117
+ if self.tp_size > 1:
118
+ lora_a = self.slice_lora_a(lora_a)
119
+ lora_b = self.slice_lora_b(lora_b)
120
+ if lora_bias is not None:
121
+ lora_bias = self.slice_bias(lora_bias)
122
+
123
+ self.lora_a_stacked[0][index,
124
+ 0, :lora_a.shape[1], :lora_a.shape[0]].copy_(
125
+ lora_a.T, non_blocking=True)
126
+ self.lora_b_stacked[0][index,
127
+ 0, :lora_b.shape[1], :lora_b.shape[0]].copy_(
128
+ lora_b.T, non_blocking=True)
129
+ if lora_bias is not None:
130
+
131
+ self.lora_bias_stacked = cast(tuple[torch.Tensor, ...],
132
+ self.lora_bias_stacked)
133
+ assert len(self.lora_bias_stacked)
134
+ self.lora_bias_stacked[0][index, 0, :lora_bias.shape[0]].copy_(
135
+ lora_bias.T, non_blocking=True)
136
+
137
+ def apply(self,
138
+ x: torch.Tensor,
139
+ bias: Optional[torch.Tensor] = None) -> torch.Tensor:
140
+ output = self.base_layer.quant_method.apply(self.base_layer, x, bias)
141
+
142
+ # In transformers backend, x and output have extra batch dimension like
143
+ # (1, seq_len, hidden_dim), while punica expects (seq_len, hidden_dim),
144
+ # therefore we need to flatten the batch dimensions.
145
+ if x.ndim == 3 and output.ndim == 3:
146
+ output = output.flatten(0, 1)
147
+ x = x.flatten(0, 1)
148
+
149
+ lora_output: Optional[
150
+ torch.Tensor] = self.punica_wrapper.add_lora_linear(
151
+ output, x, self.lora_a_stacked, self.lora_b_stacked,
152
+ self.lora_bias_stacked, 1.0, self.output_slices)
153
+ if not current_platform.can_update_inplace():
154
+ output = lora_output
155
+
156
+ return output
157
+
158
+ @property
159
+ def weight(self) -> torch.Tensor:
160
+
161
+ # unquantizedLinear
162
+ if hasattr(self.base_layer, "weight"):
163
+ return self.base_layer.weight
164
+ # Compressed Tensor
165
+ elif hasattr(self.base_layer, "weight_packed"):
166
+ return self.base_layer.weight_packed
167
+ # GPTQ/AWQ
168
+ elif hasattr(self.base_layer, "qweight"):
169
+ return self.base_layer.qweight
170
+ # marlin
171
+ elif hasattr(self.base_layer, "B"):
172
+ return self.base_layer.B
173
+ # HQQ marlin
174
+ elif hasattr(self.base_layer, "W_q"):
175
+ return self.base_layer.W_q
176
+ else:
177
+ raise ValueError(f"Unsupported base layer: {self.base_layer}")
178
+
179
+ @property
180
+ def bias(self) -> Optional[torch.Tensor]:
181
+ if hasattr(self.base_layer, "bias"):
182
+ return self.base_layer.bias
183
+ else:
184
+ return None
vllm/lora/layers/column_parallel_linear.py
new
+622
@@ -0,0 +1,622 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from typing import Optional, Union, cast
5
+
6
+import torch
7
+import torch.nn as nn
8
+from transformers import PretrainedConfig
9
+
10
+from vllm.config import LoRAConfig
11
+from vllm.distributed import (get_tensor_model_parallel_rank,
12
+ get_tensor_model_parallel_world_size,
13
+ tensor_model_parallel_all_gather)
14
+from vllm.distributed.utils import divide
15
+from vllm.model_executor.layers.linear import (ColumnParallelLinear,
16
+ MergedColumnParallelLinear,
17
+ QKVParallelLinear)
18
+from vllm.platforms import current_platform
19
+
20
+from .base_linear import BaseLinearLayerWithLoRA
21
+from .utils import _fully_sharded_can_replace, _not_fully_sharded_can_replace
22
+
23
+
24
+def _mcp_apply(x, bias, layer: "ColumnParallelLinearWithLoRA"):
25
+ """
26
+ For `ColumnParallelLinearWithLoRA` or classes that inherit from
27
+ `ColumnParallelLinearWithLoRA`, they share the same `apply` logic.
28
+ """
29
+ assert (layer.n_slices == len(layer.lora_a_stacked) == len(
30
+ layer.lora_b_stacked) == len(layer.output_slices))
31
+ if layer.lora_bias_stacked is not None:
32
+ assert layer.n_slices == len(layer.lora_bias_stacked)
33
+
34
+ output = layer.base_layer.quant_method.apply(layer.base_layer, x, bias)
35
+
36
+ x = x.view(-1, x.shape[-1])
37
+ output, out_orig_shape = output.view(-1, output.shape[-1]), output.shape
38
+
39
+ # Since communication is needed, the buffer is directly initialized as a
40
+ # tensor rather than a tuple of tensor.
41
+ buffers = torch.zeros(
42
+ (layer.n_slices, x.shape[0], layer.lora_a_stacked[0].shape[2]),
43
+ dtype=torch.float32,
44
+ device=x.device,
45
+ )
46
+
47
+ shrunk_buffers: Optional[torch.Tensor] = layer.punica_wrapper.add_shrink(
48
+ buffers, x, layer.lora_a_stacked, 1.0)
49
+
50
+ if not current_platform.can_update_inplace():
51
+ buffers = shrunk_buffers
52
+
53
+ buffers = tensor_model_parallel_all_gather(buffers)
54
+
55
+ lora_output: Optional[torch.Tensor] = layer.punica_wrapper.add_expand(
56
+ output,
57
+ buffers,
58
+ layer.lora_b_stacked,
59
+ layer.lora_bias_stacked,
60
+ layer.output_slices,
61
+ offset_start=0,
62
+ add_input=True)
63
+
64
+ if not current_platform.can_update_inplace():
65
+ output = lora_output
66
+
67
+ output = output.view(*out_orig_shape)
68
+ # now have column partitioned and packed output
69
+ return output
70
+
71
+
72
+class ColumnParallelLinearWithLoRA(BaseLinearLayerWithLoRA):
73
+ """
74
+ LoRA on top of ColumnParallelLinear layer.
75
+ LoRA B is sliced for tensor parallelism.
76
+ There are two types for the `base_layer`:
77
+ 1. ColumnParallelLinear, e.g.`dense_h_to_4h` in `FalconForCausalLM`.
78
+ 2. MergedColumnParallelLinear, e.g.`gate_up_proj` in `Phi3ForCausalLM`.
79
+ """
80
+
81
+ def __init__(self, base_layer: ColumnParallelLinear) -> None:
82
+ super().__init__(base_layer)
83
+ # The base_layer type is ColumnParallelLinear or
84
+ # MergedColumnParallelLinear, their weight sharding logic is
85
+ # inconsistent when TP is greater than 1.
86
+ self.is_merged_col_linear = type(
87
+ base_layer) is MergedColumnParallelLinear
88
+ self.tp_size = get_tensor_model_parallel_world_size()
89
+ self.output_size = self.base_layer.output_size_per_partition
90
+ # There is only one LoRA layer
91
+ self.n_slices = 1
92
+
93
+ def slice_lora_a(self, lora_a: torch.Tensor) -> torch.Tensor:
94
+ return lora_a
95
+
96
+ def slice_lora_b(self, lora_b: torch.Tensor) -> torch.Tensor:
97
+ # Applicable to cases where the base_layer is
98
+ # MergedColumnParallelLinear.
99
+ if self.is_merged_col_linear:
100
+ tp_rank = get_tensor_model_parallel_rank()
101
+ shard_size = self.output_size // 2
102
+ offset = lora_b.shape[-1] // 2
103
+
104
+ left_weight = lora_b[:, tp_rank * shard_size:(tp_rank + 1) *
105
+ shard_size]
106
+ right_weight = lora_b[:, offset + tp_rank * shard_size:offset +
107
+ (tp_rank + 1) * shard_size]
108
+ lora_b = torch.cat([left_weight, right_weight], dim=1)
109
+ # Applicable to cases where the base_layer is
110
+ # ColumnParallelLinear.
111
+ else:
112
+ tensor_model_parallel_rank = get_tensor_model_parallel_rank()
113
+ shard_size = self.output_size
114
+ start_idx = tensor_model_parallel_rank * shard_size
115
+ end_idx = (tensor_model_parallel_rank + 1) * shard_size
116
+ lora_b = lora_b[:, start_idx:end_idx]
117
+ return lora_b
118
+
119
+ def slice_bias(self, bias: torch.Tensor) -> torch.Tensor:
120
+ # TODO: Fix the slicing logic of bias.
121
+ if bias is None:
122
+ return bias
123
+ tensor_model_parallel_rank = get_tensor_model_parallel_rank()
124
+ shard_size = self.output_size
125
+ start_idx = tensor_model_parallel_rank * shard_size
126
+ end_idx = (tensor_model_parallel_rank + 1) * shard_size
127
+ bias = bias[start_idx:end_idx]
128
+ return bias
129
+
130
+ def forward(
131
+ self, input_: torch.Tensor
132
+ ) -> Union[torch.Tensor, tuple[torch.Tensor, Optional[torch.Tensor]]]:
133
+ """Forward of ColumnParallelLinear
134
+
135
+ Args:
136
+ input_: Tensor whose last dimension is `input_size`.
137
+
138
+ Returns:
139
+ - output
140
+ - bias
141
+ """
142
+ bias = (self.base_layer.bias
143
+ if not self.base_layer.skip_bias_add else None)
144
+
145
+ # Matrix multiply.
146
+ output_parallel = self.apply(input_, bias)
147
+ if self.base_layer.gather_output:
148
+ # All-gather across the partitions.
149
+ output = tensor_model_parallel_all_gather(output_parallel)
150
+ else:
151
+ output = output_parallel
152
+
153
+ if not self.base_layer.return_bias:
154
+ return output
155
+
156
+ output_bias = (self.base_layer.bias
157
+ if self.base_layer.skip_bias_add else None)
158
+ return output, output_bias
159
+
160
+ @classmethod
161
+ @_not_fully_sharded_can_replace
162
+ def can_replace_layer(
163
+ cls,
164
+ source_layer: nn.Module,
165
+ lora_config: LoRAConfig,
166
+ packed_modules_list: list,
167
+ model_config: Optional[PretrainedConfig],
168
+ ) -> bool:
169
+ return type(source_layer) is ColumnParallelLinear or (
170
+ type(source_layer) is MergedColumnParallelLinear
171
+ and len(packed_modules_list) == 1)
172
+
173
+
174
+class MergedColumnParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):
175
+ """ColumnParallelLinear layer that is composed of 2 sublayers (slices)
176
+ packed together (e.g. gate_proj + up_proj -> gate_up_proj).
177
+
178
+ This means we have 2 LoRAs, each applied to one half of the layer.
179
+
180
+ Both slices must have the same size.
181
+ """
182
+
183
+ def __init__(
184
+ self, base_layer: Union[MergedColumnParallelLinear,
185
+ QKVParallelLinear]) -> None:
186
+ super().__init__(base_layer)
187
+ # There are two LoRA layers
188
+ self.tp_size = get_tensor_model_parallel_world_size()
189
+ self.tp_rank = get_tensor_model_parallel_rank()
190
+ # the output_sizes in MergedColumnParallelLinear is not sharded by tp
191
+ # we need to divide it by the tp_size to get correct slices size
192
+ output_sizes = self.base_layer.output_sizes
193
+ self.output_slices = tuple(
194
+ divide(output_size, self.tp_size) for output_size in output_sizes)
195
+ self.n_slices = len(self.output_slices)
196
+ self.output_ids = (self.tp_rank, ) * self.n_slices
197
+
198
+ def create_lora_weights(
199
+ self,
200
+ max_loras: int,
201
+ lora_config: LoRAConfig,
202
+ model_config: Optional[PretrainedConfig] = None,
203
+ ) -> None:
204
+ """
205
+ The main reason for overriding this function is to enhance code
206
+ maintainability.
207
+ """
208
+ self.lora_config = lora_config
209
+
210
+ lora_a_output_size_per_partition = (
211
+ lora_config.max_lora_rank if not lora_config.fully_sharded_loras
212
+ else divide(lora_config.max_lora_rank, self.tp_size))
213
+
214
+ self.lora_a_stacked = tuple(
215
+ torch.zeros(
216
+ max_loras,
217
+ 1,
218
+ lora_a_output_size_per_partition,
219
+ self.input_size,
220
+ dtype=lora_config.lora_dtype,
221
+ device=self.device,
222
+ ) for _ in range(self.n_slices))
223
+ self.lora_b_stacked = tuple(
224
+ torch.zeros(
225
+ max_loras,
226
+ 1,
227
+ output_size,
228
+ lora_config.max_lora_rank,
229
+ dtype=lora_config.lora_dtype,
230
+ device=self.device,
231
+ ) for output_size in self.output_slices)
232
+ if lora_config.bias_enabled:
233
+ self.lora_bias_stacked = tuple(
234
+ torch.zeros(
235
+ max_loras,
236
+ 1,
237
+ output_size,
238
+ dtype=lora_config.lora_dtype,
239
+ device=self.device,
240
+ ) for output_size in self.output_slices)
241
+
242
+ def slice_lora_a(
243
+ self, lora_a: list[Union[torch.Tensor, None]]
244
+ ) -> list[Union[torch.Tensor, None]]:
245
+ return lora_a
246
+
247
+ def slice_lora_b(
248
+ self, lora_b: list[Union[torch.Tensor, None]]
249
+ ) -> list[Union[torch.Tensor, None]]:
250
+ sliced_lora_b = [None] * self.n_slices
251
+ for i, (shard_id, shard_size) in enumerate(
252
+ zip(self.output_ids, self.output_slices)):
253
+ if (lora_b_i := lora_b[i]) is not None:
254
+ sliced_lora_b[i] = lora_b_i[:,
255
+ shard_size * shard_id:shard_size *
256
+ (shard_id + 1)]
257
+ return sliced_lora_b
258
+
259
+ def slice_bias(
260
+ self, bias: list[Union[torch.Tensor,
261
+ None]]) -> list[Union[torch.Tensor, None]]:
262
+ for i, (shard_id, shard_size) in enumerate(
263
+ zip(self.output_ids, self.output_slices)):
264
+ if (bias_i := bias[i]) is not None:
265
+ bias[i] = bias_i[shard_size * shard_id:shard_size *
266
+ (shard_id + 1)]
267
+ return bias
268
+
269
+ def set_lora(
270
+ self,
271
+ index: int,
272
+ lora_a: torch.Tensor,
273
+ lora_b: torch.Tensor,
274
+ embeddings_tensor: Optional[torch.Tensor],
275
+ lora_bias: Optional[torch.Tensor] = None,
276
+ ):
277
+ self.reset_lora(index)
278
+
279
+ if self.tp_size > 1:
280
+ lora_a = self.slice_lora_a(lora_a)
281
+ lora_b = self.slice_lora_b(lora_b)
282
+ if lora_bias is not None:
283
+ lora_bias = self.slice_bias(lora_bias)
284
+
285
+ for i in range(self.n_slices):
286
+ if (lora_a_i := lora_a[i]) is not None:
287
+ self.lora_a_stacked[i][
288
+ index, 0, :lora_a_i.shape[1], :lora_a_i.shape[0]].copy_(
289
+ lora_a_i.T, non_blocking=True)
290
+ if (lora_b_i := lora_b[i]) is not None:
291
+ self.lora_b_stacked[i][
292
+ index, 0, :lora_b_i.shape[1], :lora_b_i.shape[0]].copy_(
293
+ lora_b_i.T, non_blocking=True)
294
+
295
+ if lora_bias is not None:
296
+ self.lora_bias_stacked = cast(tuple[torch.Tensor, ...],
297
+ self.lora_bias_stacked)
298
+ for i in range(self.n_slices):
299
+ if (lora_bias_i := lora_bias[i]) is not None:
300
+ self.lora_bias_stacked[i][index,
301
+ 0, :lora_bias_i.shape[0]].copy_(
302
+ lora_bias_i.T,
303
+ non_blocking=True)
304
+
305
+ @classmethod
306
+ @_not_fully_sharded_can_replace
307
+ def can_replace_layer(
308
+ cls,
309
+ source_layer: nn.Module,
310
+ lora_config: LoRAConfig,
311
+ packed_modules_list: list,
312
+ model_config: Optional[PretrainedConfig],
313
+ ) -> bool:
314
+ return (type(source_layer) is MergedColumnParallelLinear
315
+ and len(packed_modules_list) == 2)
316
+
317
+
318
+class QKVParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):
319
+ """
320
+ ColumnParallelLinear layer that is specifically designed for
321
+ qkv_proj. Certain models, such as chatglm3 and baichuan-7b,
322
+ only contains a single LoRA within their qkv_proj layer.
323
+
324
+ During inference with Tensor Parallel, the weights of lora_b
325
+ must be accurately partitioned according to the respective ranks.
326
+
327
+ Q slice may have different shape than K and V slices (which both have
328
+ the same shape).
329
+ """
330
+
331
+ def __init__(self, base_layer: QKVParallelLinear) -> None:
332
+ super().__init__(base_layer)
333
+ self.q_proj_total_size = (self.base_layer.total_num_heads *
334
+ self.base_layer.head_size)
335
+ self.q_proj_shard_size = (self.base_layer.num_heads *
336
+ self.base_layer.head_size)
337
+ self.kv_proj_shard_size = (self.base_layer.num_kv_heads *
338
+ self.base_layer.head_size)
339
+ self.kv_proj_total_size = (self.base_layer.total_num_kv_heads *
340
+ self.base_layer.head_size)
341
+ # There is only one LoRA layer
342
+ self.n_slices = 1
343
+
344
+ def slice_lora_b(self, lora_b: torch.Tensor) -> torch.Tensor:
345
+ tp_rank = get_tensor_model_parallel_rank()
346
+ self.q_shard_id = tp_rank
347
+ self.kv_shard_id = tp_rank // self.base_layer.num_kv_head_replicas
348
+ lora_b_q = lora_b[:, self.q_proj_shard_size *
349
+ self.q_shard_id:self.q_proj_shard_size *
350
+ (self.q_shard_id + 1)]
351
+ k_offset = self.q_proj_total_size
352
+ lora_b_k = lora_b[:, k_offset +
353
+ self.kv_proj_shard_size * self.kv_shard_id:k_offset +
354
+ self.kv_proj_shard_size * (self.kv_shard_id + 1)]
355
+ v_offset = k_offset + self.kv_proj_total_size
356
+ lora_b_v = lora_b[:, v_offset +
357
+ self.kv_proj_shard_size * self.kv_shard_id:v_offset +
358
+ self.kv_proj_shard_size * (self.kv_shard_id + 1)]
359
+ lora_b = torch.cat([lora_b_q, lora_b_k, lora_b_v], dim=1)
360
+ return lora_b
361
+
362
+ def slice_bias(self, bias: torch.Tensor) -> torch.Tensor:
363
+ bias_q = bias[self.q_proj_shard_size *
364
+ self.q_shard_id:self.q_proj_shard_size *
365
+ (self.q_shard_id + 1)]
366
+ k_offset = self.q_proj_total_size
367
+ bias_k = bias[k_offset +
368
+ self.kv_proj_shard_size * self.kv_shard_id:k_offset +
369
+ self.kv_proj_shard_size * (self.kv_shard_id + 1)]
370
+ v_offset = k_offset + self.kv_proj_total_size
371
+ bias_v = bias[v_offset +
372
+ self.kv_proj_shard_size * self.kv_shard_id:v_offset +
373
+ self.kv_proj_shard_size * (self.kv_shard_id + 1)]
374
+ bias = torch.cat([bias_q, bias_k, bias_v], dim=1)
375
+ return bias
376
+
377
+ @classmethod
378
+ @_not_fully_sharded_can_replace
379
+ def can_replace_layer(cls, source_layer: nn.Module,
380
+ lora_config: LoRAConfig, packed_modules_list: list,
381
+ model_config: Optional[PretrainedConfig]) -> bool:
382
+ return type(source_layer) is QKVParallelLinear and len(
383
+ packed_modules_list) == 1
384
+
385
+
386
+class MergedQKVParallelLinearWithLoRA(MergedColumnParallelLinearWithLoRA):
387
+ """MergedColumnParallelLinear layer that is composed of 3 sublayers (slices)
388
+ packed together in qkv proj fashion
389
+ (q_proj + k_proj + v_proj -> qkv_proj).
390
+
391
+ This means we have 3 LoRAs, each applied to one slice of the layer.
392
+
393
+ Q slice may have different shape than K and V slices (which both have
394
+ the same shape).
395
+ """
396
+
397
+ def __init__(self, base_layer: QKVParallelLinear) -> None:
398
+ super().__init__(base_layer)
399
+ # There are three LoRA layer.
400
+ self.n_slices = len(self.base_layer.output_sizes)
401
+ self.tp_size = get_tensor_model_parallel_world_size()
402
+ self.tp_rank = get_tensor_model_parallel_rank()
403
+
404
+ self.q_proj_shard_size = (self.base_layer.num_heads *
405
+ self.base_layer.head_size)
406
+ self.kv_proj_shard_size = (self.base_layer.num_kv_heads *
407
+ self.base_layer.head_size)
408
+ self.q_shard_id = self.tp_rank
409
+ self.kv_shard_id = self.tp_rank // self.base_layer.num_kv_head_replicas
410
+
411
+ self.output_slices = (
412
+ self.q_proj_shard_size,
413
+ self.kv_proj_shard_size,
414
+ self.kv_proj_shard_size,
415
+ )
416
+ self.output_ids = (
417
+ self.q_shard_id,
418
+ self.kv_shard_id,
419
+ self.kv_shard_id,
420
+ )
421
+
422
+ def create_lora_weights(
423
+ self,
424
+ max_loras: int,
425
+ lora_config: LoRAConfig,
426
+ model_config: Optional[PretrainedConfig] = None,
427
+ ) -> None:
428
+ """
429
+ The main reason for overloading this function is to handle inconsistent
430
+ weight dimensions in qkv lora.
431
+ """
432
+ super().create_lora_weights(max_loras, lora_config, model_config)
433
+
434
+ @classmethod
435
+ @_not_fully_sharded_can_replace
436
+ def can_replace_layer(
437
+ cls,
438
+ source_layer: nn.Module,
439
+ lora_config: LoRAConfig,
440
+ packed_modules_list: list,
441
+ model_config: Optional[PretrainedConfig],
442
+ ) -> bool:
443
+ return (type(source_layer) is QKVParallelLinear
444
+ and len(packed_modules_list) == 3)
445
+
446
+
447
+# These following layers are based on the tensor parallelism strategy given in
448
+# Y. Sheng et al., S-LoRA: Serving Thousands of Concurrent LoRA Adapters. 2023,
449
+# https://arxiv.org/abs/2311.03285.
450
+
451
+
452
+class ColumnParallelLinearWithShardedLoRA(ColumnParallelLinearWithLoRA):
453
+ """
454
+ Differs from ColumnParallelLinearWithLoRA by slicing LoRA A also.
455
+
456
+ Based on S-LoRA, slicing happens along the rank dim.
457
+ """
458
+
459
+ # For all LoRA layers where the `base_layer` is `ColumnParallelLinear`,
460
+ # their `lora_a` and `lora_b` have different sharding patterns. After
461
+ # completing the `lora_a` GEMM , a gather operation is performed.
462
+ # Therefore, the sharding of `lora_a` only needs to correspond with the
463
+ # gather operation.
464
+ def slice_lora_a(self, lora_a: torch.Tensor) -> torch.Tensor:
465
+ tp_rank = get_tensor_model_parallel_rank()
466
+ shard_size = self.lora_a_stacked[0].shape[2]
467
+ start_idx = tp_rank * shard_size
468
+ lora_a = lora_a[:, start_idx:start_idx + shard_size]
469
+ return lora_a
470
+
471
+ def apply(self,
472
+ x: torch.Tensor,
473
+ bias: Optional[torch.Tensor] = None) -> torch.Tensor:
474
+ return _mcp_apply(x, bias, self)
475
+
476
+ @classmethod
477
+ @_fully_sharded_can_replace
478
+ def can_replace_layer(
479
+ cls,
480
+ source_layer: nn.Module,
481
+ lora_config: LoRAConfig,
482
+ packed_modules_list: list,
483
+ model_config: Optional[PretrainedConfig],
484
+ ) -> bool:
485
+ # specifying kwargs so they can be easily accessed in decorator
486
+ return super().can_replace_layer(
487
+ source_layer=source_layer,
488
+ lora_config=lora_config,
489
+ packed_modules_list=packed_modules_list,
490
+ model_config=model_config,
491
+ decorate=False,
492
+ )
493
+
494
+
495
+class MergedColumnParallelLinearWithShardedLoRA(
496
+ MergedColumnParallelLinearWithLoRA):
497
+ """
498
+ Differs from MergedColumnParallelLinearWithLoRA by slicing the
499
+ LoRA A's also.
500
+
501
+ Based on S-LoRA, slicing happens along the rank dim.
502
+ """
503
+
504
+ def slice_lora_a(
505
+ self, lora_a: list[Union[torch.Tensor, None]]
506
+ ) -> list[Union[torch.Tensor, None]]:
507
+ #NOTE: lora_a contains 2 subloras, and each sublora could be None.
508
+ output_shard_size = self.lora_a_stacked[0].shape[2]
509
+ output_start_idx = self.tp_rank * output_shard_size
510
+ lora_a = [
511
+ lora_a[0][:, output_start_idx:output_start_idx +
512
+ output_shard_size] if lora_a[0] is not None else None,
513
+ lora_a[1][:, output_start_idx:output_start_idx +
514
+ output_shard_size] if lora_a[1] is not None else None,
515
+ ]
516
+ return lora_a
517
+
518
+ def apply(self,
519
+ x: torch.Tensor,
520
+ bias: Optional[torch.Tensor] = None) -> torch.Tensor:
521
+ return _mcp_apply(x, bias, self)
522
+
523
+ @classmethod
524
+ @_fully_sharded_can_replace
525
+ def can_replace_layer(
526
+ cls,
527
+ source_layer: nn.Module,
528
+ lora_config: LoRAConfig,
529
+ packed_modules_list: list,
530
+ model_config: Optional[PretrainedConfig],
531
+ ) -> bool:
532
+ # specifying kwargs so they can be easily accessed in decorator
533
+ return super().can_replace_layer(
534
+ source_layer=source_layer,
535
+ lora_config=lora_config,
536
+ packed_modules_list=packed_modules_list,
537
+ model_config=model_config,
538
+ decorate=False,
539
+ )
540
+
541
+
542
+class QKVParallelLinearWithShardedLoRA(QKVParallelLinearWithLoRA):
543
+ """
544
+ Differs from QKVParallelLinearWithLoRA by slicing the
545
+ LoRA A's also.
546
+
547
+ Based on S-LoRA, slicing happens along the rank dim.
548
+ """
549
+
550
+ def slice_lora_a(self, lora_a: torch.Tensor) -> torch.Tensor:
551
+ tp_rank = get_tensor_model_parallel_rank()
552
+ shard_size = self.lora_a_stacked[0].shape[2]
553
+ start_idx = tp_rank * shard_size
554
+ lora_a = lora_a[:, start_idx:start_idx + shard_size]
555
+ return lora_a
556
+
557
+ def apply(self,
558
+ x: torch.Tensor,
559
+ bias: Optional[torch.Tensor] = None) -> torch.Tensor:
560
+ return _mcp_apply(x, bias, self)
561
+
562
+ @classmethod
563
+ @_fully_sharded_can_replace
564
+ def can_replace_layer(cls, source_layer: nn.Module,
565
+ lora_config: LoRAConfig, packed_modules_list: list,
566
+ model_config: Optional[PretrainedConfig]) -> bool:
567
+ # specifying kwargs so they can be easily accessed in decorator
568
+ return super().can_replace_layer(
569
+ source_layer=source_layer,
570
+ lora_config=lora_config,
571
+ packed_modules_list=packed_modules_list,
572
+ model_config=model_config,
573
+ decorate=False,
574
+ )
575
+
576
+
577
+class MergedQKVParallelLinearWithShardedLoRA(MergedQKVParallelLinearWithLoRA):
578
+ """
579
+ Differs from MergedQKVParallelLinearWithLoRA by slicing the
580
+ LoRA A's also.
581
+
582
+ Based on S-LoRA, slicing happens along the rank dim.
583
+ """
584
+
585
+ def slice_lora_a(
586
+ self, lora_a: list[Union[torch.Tensor, None]]
587
+ ) -> list[Union[torch.Tensor, None]]:
588
+ # NOTE: lora_a contains 3 subloras, and each sublora could be None.
589
+ shard_size = [self.lora_a_stacked[i].shape[2] for i in range(3)]
590
+ start_idx = [self.tp_rank * shard_size[i] for i in range(3)]
591
+ lora_a = [
592
+ lora_a[0][:, start_idx[0]:start_idx[0] +
593
+ shard_size[0]] if lora_a[0] is not None else None,
594
+ lora_a[1][:, start_idx[1]:start_idx[1] +
595
+ shard_size[1]] if lora_a[1] is not None else None,
596
+ lora_a[2][:, start_idx[2]:start_idx[2] +
597
+ shard_size[2]] if lora_a[2] is not None else None,
598
+ ]
599
+ return lora_a
600
+
601
+ def apply(self,
602
+ x: torch.Tensor,
603
+ bias: Optional[torch.Tensor] = None) -> torch.Tensor:
604
+ return _mcp_apply(x, bias, self)
605
+
606
+ @classmethod
607
+ @_fully_sharded_can_replace
608
+ def can_replace_layer(
609
+ cls,
610
+ source_layer: nn.Module,
611
+ lora_config: LoRAConfig,
612
+ packed_modules_list: list,
613
+ model_config: Optional[PretrainedConfig],
614
+ ) -> bool:
615
+ # specifying kwargs so they can be easily accessed in decorator
616
+ return super().can_replace_layer(
617
+ source_layer=source_layer,
618
+ lora_config=lora_config,
619
+ packed_modules_list=packed_modules_list,
620
+ model_config=model_config,
621
+ decorate=False,
622
+ )
vllm/lora/layers/logits_processor.py
new
+247
@@ -0,0 +1,247 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import math
5
+from typing import Optional
6
+
7
+import torch
8
+import torch.nn as nn
9
+from transformers import PretrainedConfig
10
+
11
+from vllm.config import LoRAConfig
12
+from vllm.distributed import (get_tensor_model_parallel_rank,
13
+ get_tensor_model_parallel_world_size)
14
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
15
+from vllm.model_executor.layers.vocab_parallel_embedding import (
16
+ VocabParallelEmbedding)
17
+from vllm.platforms import current_platform
18
+
19
+from .base import BaseLayerWithLoRA
20
+
21
+
22
+class LogitsProcessorWithLoRA(BaseLayerWithLoRA):
23
+ """
24
+ LoRA wrapper for LogitsProcessor, with extra logic to handle the
25
+ application of the LoRA adapter and added LoRA vocabulary.
26
+
27
+ Args:
28
+ base_layer: LogitsProcessor layer
29
+ hidden_size: hidden size of the model
30
+ dtype: data type of the model
31
+ device: device of the model
32
+ sharded_to_full_mapping: index mapping from sharded vocab to full vocab
33
+ received from base_layer.get_sharded_to_full_mapping(). If None,
34
+ no reindexing will be done.
35
+ """
36
+
37
+ def __init__(self, base_layer: LogitsProcessor, hidden_size: int,
38
+ dtype: torch.dtype, device: torch.device,
39
+ sharded_to_full_mapping: Optional[list[int]]) -> None:
40
+ super().__init__()
41
+ self.base_layer = base_layer
42
+ self.hidden_size = hidden_size
43
+ self.dtype = dtype
44
+ self.device = device
45
+ self.tp_size = get_tensor_model_parallel_world_size()
46
+ self.tp_rank = get_tensor_model_parallel_rank()
47
+ self.sharded_to_full_mapping = sharded_to_full_mapping
48
+
49
+ @property
50
+ def logits_as_input(self):
51
+ return self.base_layer.logits_as_input
52
+
53
+ @property
54
+ def vocab_size(self):
55
+ return self.base_layer.vocab_size
56
+
57
+ @property
58
+ def scale(self):
59
+ return self.base_layer.scale
60
+
61
+ @property
62
+ def soft_cap(self):
63
+ return self.base_layer.soft_cap
64
+
65
+ @property
66
+ def use_all_gather(self):
67
+ return self.base_layer.use_all_gather
68
+
69
+ @property
70
+ def org_vocab_size(self):
71
+ return self.base_layer.org_vocab_size
72
+
73
+ @property
74
+ def include_gpu_probs_tensor(self):
75
+ return self.base_layer.include_gpu_probs_tensor
76
+
77
+ @property
78
+ def should_modify_greedy_probs_inplace(self):
79
+ return self.base_layer.should_modify_greedy_probs_inplace
80
+
81
+ def create_lora_weights(
82
+ self,
83
+ max_loras: int,
84
+ lora_config: LoRAConfig,
85
+ model_config: Optional[PretrainedConfig] = None,
86
+ ) -> None:
87
+ # TODO: Verify if this condition can be further relaxed
88
+ if 32000 < self.base_layer.vocab_size > 257024:
89
+ raise ValueError("When using LoRA, vocab size must be "
90
+ "32000 >= vocab_size <= 257024")
91
+ self.lora_a_stacked = torch.zeros(
92
+ (
93
+ max_loras,
94
+ 1,
95
+ lora_config.max_lora_rank,
96
+ self.hidden_size,
97
+ ),
98
+ dtype=lora_config.lora_dtype,
99
+ device=self.device,
100
+ )
101
+ self.lora_b_stacked = torch.zeros(
102
+ (
103
+ max_loras,
104
+ 1,
105
+ # Pad for kernel compatibility
106
+ math.ceil(self.base_layer.vocab_size /
107
+ lora_config.lora_vocab_padding_size) *
108
+ lora_config.lora_vocab_padding_size,
109
+ lora_config.max_lora_rank,
110
+ ),
111
+ dtype=lora_config.lora_dtype,
112
+ device=self.device,
113
+ )
114
+ self.embeddings_tensors = torch.full(
115
+ (max_loras, lora_config.lora_extra_vocab_size, self.hidden_size),
116
+ fill_value=float("-inf"),
117
+ dtype=self.dtype,
118
+ device=self.device,
119
+ )
120
+ if self.sharded_to_full_mapping is not None:
121
+ self.sharded_to_full_mapping_gpu = torch.tensor(
122
+ self.sharded_to_full_mapping,
123
+ device=self.device,
124
+ dtype=torch.long)
125
+ else:
126
+ self.sharded_to_full_mapping_gpu = None
127
+
128
+ def reset_lora(self, index: int):
129
+ self.lora_a_stacked[index] = 0
130
+ self.lora_b_stacked[index] = 0
131
+ self.embeddings_tensors[index] = float("-inf")
132
+
133
+ def set_lora(
134
+ self,
135
+ index: int,
136
+ lora_a: torch.Tensor,
137
+ lora_b: torch.Tensor,
138
+ embeddings_tensor: Optional[torch.Tensor],
139
+ bias: Optional[torch.Tensor] = None,
140
+ ):
141
+ self.reset_lora(index)
142
+ self.lora_a_stacked[index,
143
+ 0, :lora_a.shape[1], :lora_a.shape[0]].copy_(
144
+ lora_a.T, non_blocking=True)
145
+ self.lora_b_stacked[index,
146
+ 0, :lora_b.shape[1], :lora_b.shape[0]].copy_(
147
+ lora_b.T, non_blocking=True)
148
+ if embeddings_tensor is not None:
149
+ self.embeddings_tensors[
150
+ index,
151
+ :embeddings_tensor.shape[0],
152
+ :embeddings_tensor.shape[1],
153
+ ] = embeddings_tensor
154
+
155
+ def _get_logits(
156
+ self,
157
+ hidden_states: torch.Tensor,
158
+ lm_head: VocabParallelEmbedding,
159
+ embedding_bias: Optional[torch.Tensor] = None,
160
+ ) -> Optional[torch.Tensor]:
161
+ # Get the logits for the next tokens.
162
+ logits = lm_head.quant_method.apply(lm_head, hidden_states)
163
+ if embedding_bias is not None:
164
+ logits += embedding_bias
165
+
166
+ # Gather logits for TP
167
+ logits = self.base_layer._gather_logits(logits)
168
+
169
+ if logits is None:
170
+ return None
171
+
172
+ if self.sharded_to_full_mapping_gpu is not None:
173
+ # Reindex full logits tensor to ensure 1:1 mapping between
174
+ # index and token_id
175
+ # Example for:
176
+ # org_vocab_size = 4
177
+ # added_vocab_size = 2
178
+ # pad_to_size = 8
179
+ # tp_size = 2
180
+
181
+ # indices: [0, 1, 2, 3, 4, 5, 6, 7]
182
+ # token_id: [0, 1, 4, -1, 2, 3, 5, -1]
183
+
184
+ # Therefore, the mapping is expected to be:
185
+ # [0, 1, 4, 6, 2, 3, 5, 7] so that when we reindex,
186
+ # we get:
187
+ # indices: [0, 1, 2, 3, 4, 5, 6, 7]
188
+ # token_id: [0, 1, 2, 3, 4, 5, -1, -1]
189
+ logits = logits[:, self.sharded_to_full_mapping_gpu]
190
+
191
+ lora_logits = torch.empty(
192
+ self.embeddings_tensors.shape[0] + 1,
193
+ self.embeddings_tensors.shape[1],
194
+ hidden_states.shape[0],
195
+ dtype=self.embeddings_tensors.dtype,
196
+ device=self.embeddings_tensors.device,
197
+ )
198
+ torch.matmul(self.embeddings_tensors,
199
+ hidden_states.T,
200
+ out=lora_logits[:-1])
201
+
202
+ neg_inf, pos_inf = current_platform.get_infinity_values(
203
+ lora_logits.dtype)
204
+
205
+ lora_logits[-1] = neg_inf
206
+ lora_logits = lora_logits.mT
207
+ indices_padded = self.punica_wrapper.sampler_indices_padded
208
+
209
+ if current_platform.is_tpu() or current_platform.is_xpu():
210
+ indices_padded = indices_padded[:logits.size(0)]
211
+
212
+ lora_logits = (lora_logits.reshape(
213
+ lora_logits.shape[0] * lora_logits.shape[1],
214
+ lora_logits.shape[2],
215
+ ).index_select(0, indices_padded).nan_to_num_(nan=neg_inf,
216
+ posinf=pos_inf,
217
+ neginf=neg_inf))
218
+
219
+ logits[:,
220
+ self.base_layer.org_vocab_size:self.base_layer.org_vocab_size +
221
+ lora_logits.shape[1]] = lora_logits
222
+
223
+ lora_output: Optional[
224
+ torch.Tensor] = self.punica_wrapper.add_lora_logits(
225
+ logits, hidden_states, self.lora_a_stacked,
226
+ self.lora_b_stacked, 1.0)
227
+
228
+ if not current_platform.can_update_inplace():
229
+ logits = lora_output
230
+
231
+ # Remove paddings in vocab (if any).
232
+ logits = logits[:, :self.base_layer.vocab_size]
233
+ return logits
234
+
235
+ def forward(self, *args, **kwargs):
236
+ return type(self.base_layer).forward(self, *args, **kwargs)
237
+
238
+ @classmethod
239
+ def can_replace_layer(
240
+ cls,
241
+ source_layer: nn.Module,
242
+ lora_config: LoRAConfig,
243
+ packed_modules_list: list,
244
+ model_config: Optional[PretrainedConfig],
245
+ ) -> bool:
246
+ # Special handling for the LogitsProcessor.
247
+ return False
vllm/lora/layers/qkv_x_parallel_linear.py
new
+8
@@ -0,0 +1,8 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from .base import BaseLayerWithLoRA
4
+
5
+
6
+#TODO: Implement this
7
+class QKVCrossParallelLinearWithLoRA(BaseLayerWithLoRA):
8
+ pass
vllm/lora/layers/replicated_linear.py
new
+61
@@ -0,0 +1,61 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from typing import Optional, Union
5
+
6
+import torch
7
+import torch.nn as nn
8
+from transformers import PretrainedConfig
9
+
10
+from vllm.config import LoRAConfig
11
+from vllm.model_executor.layers.linear import ReplicatedLinear
12
+
13
+from .base_linear import BaseLinearLayerWithLoRA
14
+
15
+
16
+class ReplicatedLinearWithLoRA(BaseLinearLayerWithLoRA):
17
+
18
+ def __init__(self, base_layer: ReplicatedLinear) -> None:
19
+ super().__init__(base_layer, )
20
+ # To ensure interface compatibility, set to 1 always.
21
+ self.tp_size = 1
22
+ self.output_size = self.base_layer.output_size
23
+ self.n_slices = 1
24
+
25
+ def forward(
26
+ self, input_: torch.Tensor
27
+ ) -> Union[torch.Tensor, tuple[torch.Tensor, Optional[torch.Tensor]]]:
28
+ """Forward of ReplicatedLinearWithLoRA
29
+
30
+ Args:
31
+ input_: Tensor whose last dimension is `input_size`.
32
+
33
+ Returns:
34
+ - output
35
+ - bias
36
+ """
37
+ bias = (self.base_layer.bias
38
+ if not self.base_layer.skip_bias_add else None)
39
+
40
+ # Matrix multiply.
41
+ output = self.apply(input_, bias)
42
+
43
+ output_bias = (self.base_layer.bias
44
+ if self.base_layer.skip_bias_add else None)
45
+
46
+ if not self.base_layer.return_bias:
47
+ return output
48
+
49
+ return output, output_bias
50
+
51
+ # ReplicatedLinear should always be replaced, regardless of the fully
52
+ # sharded LoRAs setting, because it is, by definition, copied per GPU.
53
+ @classmethod
54
+ def can_replace_layer(
55
+ cls,
56
+ source_layer: nn.Module,
57
+ lora_config: LoRAConfig,
58
+ packed_modules_list: list,
59
+ model_config: Optional[PretrainedConfig],
60
+ ) -> bool:
61
+ return type(source_layer) is ReplicatedLinear
vllm/lora/layers/row_parallel_linear.py
new
+201
@@ -0,0 +1,201 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from typing import Optional, Union, cast
5
+
6
+import torch
7
+import torch.nn as nn
8
+from transformers import PretrainedConfig
9
+
10
+from vllm.config import LoRAConfig
11
+from vllm.distributed import (get_tensor_model_parallel_rank,
12
+ get_tensor_model_parallel_world_size,
13
+ split_tensor_along_last_dim,
14
+ tensor_model_parallel_all_reduce)
15
+# yapf: disable
16
+from vllm.model_executor.layers.linear import RowParallelLinear
17
+from vllm.platforms import current_platform
18
+
19
+from .base_linear import BaseLinearLayerWithLoRA
20
+from .utils import _fully_sharded_can_replace, _not_fully_sharded_can_replace
21
+
22
+
23
+class RowParallelLinearWithLoRA(BaseLinearLayerWithLoRA):
24
+
25
+ def __init__(self, base_layer: RowParallelLinear) -> None:
26
+ super().__init__(base_layer)
27
+
28
+ self.tp_size = get_tensor_model_parallel_world_size()
29
+ # reset input_size
30
+ self.input_size = self.base_layer.input_size_per_partition
31
+ self.output_size = self.base_layer.output_size
32
+
33
+ self.tp_rank = get_tensor_model_parallel_rank()
34
+ # There is only one LoRA layer.
35
+ self.n_slices = 1
36
+
37
+ def slice_lora_a(self, lora_a: torch.Tensor) -> torch.Tensor:
38
+
39
+ shard_size = self.input_size
40
+ start_idx = self.tp_rank * shard_size
41
+ end_idx = (self.tp_rank + 1) * shard_size
42
+ lora_a = lora_a[start_idx:end_idx, :]
43
+ return lora_a
44
+
45
+ def slice_lora_b(self, lora_b: torch.Tensor) -> torch.Tensor:
46
+ return lora_b
47
+
48
+ def slice_bias(self, bias: torch.Tensor) -> torch.Tensor:
49
+ return bias
50
+
51
+ def forward(
52
+ self, input_: torch.Tensor
53
+ ) -> Union[torch.Tensor, tuple[torch.Tensor, Optional[torch.Tensor]]]:
54
+ """Forward of RowParallelLinear
55
+
56
+ Args:
57
+ input_: tensor whose last dimension is `input_size`. If
58
+ `input_is_parallel` is set, then the last dimension
59
+ is `input_size // tp_size`.
60
+
61
+ Returns:
62
+ - output
63
+ - bias
64
+ """
65
+ # set up backprop all-reduce.
66
+ if self.base_layer.input_is_parallel:
67
+ input_parallel = input_
68
+ else:
69
+ # TODO: simplify code below
70
+ splitted_input = split_tensor_along_last_dim(
71
+ input_, num_partitions=self.base_layer.tp_size)
72
+ input_parallel = splitted_input[self.tp_rank].contiguous()
73
+
74
+ # Matrix multiply.
75
+ output_parallel = self.apply(input_parallel)
76
+ if self.base_layer.reduce_results and self.base_layer.tp_size > 1:
77
+ output_ = tensor_model_parallel_all_reduce(output_parallel)
78
+ else:
79
+ output_ = output_parallel
80
+
81
+ if not self.base_layer.skip_bias_add:
82
+ output = (output_ + self.base_layer.bias
83
+ if self.base_layer.bias is not None else output_)
84
+ output_bias = None
85
+ else:
86
+ output = output_
87
+ output_bias = self.base_layer.bias
88
+
89
+ if not self.base_layer.return_bias:
90
+ return output
91
+
92
+ return output, output_bias
93
+
94
+ @classmethod
95
+ @_not_fully_sharded_can_replace
96
+ def can_replace_layer(
97
+ cls,
98
+ source_layer: nn.Module,
99
+ lora_config: LoRAConfig,
100
+ packed_modules_list: list,
101
+ model_config: Optional[PretrainedConfig],
102
+ ) -> bool:
103
+ return type(source_layer) is RowParallelLinear
104
+
105
+
106
+
107
+# The following layer is based on the tensor parallelism strategy given in
108
+# Y. Sheng et al., S-LoRA: Serving Thousands of Concurrent LoRA Adapters. 2023,
109
+# https://arxiv.org/abs/2311.03285.
110
+
111
+class RowParallelLinearWithShardedLoRA(RowParallelLinearWithLoRA):
112
+ """
113
+ Differs from RowParallelLinearWithLoRA by slicing the
114
+ LoRA B's also.
115
+
116
+ Based on S-LoRA, slicing happens along the output dim.
117
+ This yields a combined partial sum from the row parallel base
118
+ layer and column partitioned output from the LoRA.
119
+ """
120
+
121
+ def slice_lora_b(self, lora_b: torch.Tensor) -> torch.Tensor:
122
+ shard_size = self.lora_b_stacked[0].shape[2]
123
+ start_idx = self.tp_rank * shard_size
124
+ end_idx = (self.tp_rank + 1) * shard_size
125
+ lora_b = lora_b[:, start_idx:end_idx]
126
+ return lora_b
127
+
128
+ def slice_bias(self, bias: torch.Tensor) -> torch.Tensor:
129
+ if bias is None:
130
+ return bias
131
+ self.lora_bias_stacked = cast(tuple[torch.Tensor, ...],
132
+ self.lora_bias_stacked)
133
+ shard_size = self.lora_bias_stacked[0].shape[2]
134
+ start_idx = self.tp_rank * shard_size
135
+ end_idx = (self.tp_rank + 1) * shard_size
136
+ bias = bias[start_idx:end_idx]
137
+ return bias
138
+
139
+ def apply(self,
140
+ x: torch.Tensor,
141
+ bias: Optional[torch.Tensor] = None) -> torch.Tensor:
142
+ output = self.base_layer.quant_method.apply(self.base_layer, x)
143
+
144
+ x = x.view(-1, x.shape[-1])
145
+ output, out_orig_shape = output.view(-1,
146
+ output.shape[-1]), output.shape
147
+ buffer = torch.zeros(
148
+ (self.n_slices, x.shape[0], self.lora_a_stacked[0].shape[2]),
149
+ dtype=torch.float32,
150
+ device=x.device,
151
+ )
152
+
153
+ shrunk_buffer: Optional[torch.Tensor] = self.punica_wrapper.add_shrink(
154
+ buffer, x, self.lora_a_stacked, 1.0)
155
+ if not current_platform.can_update_inplace():
156
+ buffer = shrunk_buffer
157
+
158
+ buffer = tensor_model_parallel_all_reduce(buffer)
159
+
160
+ # following S-LoRA, allows the fusing of all_gather and all_reduce
161
+ # by adding the column partitioned lora output to a slice of output
162
+ # tensor, which is a partial sum due to row parallel. All that
163
+ # remains is a standard all_reduce. User should be aware though that
164
+ # the output is not the same as a normal row_parallel, it should be
165
+ # reduced before being used
166
+ # NOTE offset are based on the rank.
167
+ shard_size = self.lora_b_stacked[0].shape[2]
168
+ offset_start = self.tp_rank * shard_size
169
+ lora_output: Optional[torch.Tensor] = self.punica_wrapper.add_expand(
170
+ output,
171
+ buffer,
172
+ self.lora_b_stacked,
173
+ self.lora_bias_stacked,
174
+ self.output_slices,
175
+ offset_start=offset_start,
176
+ add_input=True,
177
+ )
178
+
179
+ if not current_platform.can_update_inplace():
180
+ output = lora_output
181
+
182
+ output = output.view(*out_orig_shape)
183
+ return output
184
+
185
+ @classmethod
186
+ @_fully_sharded_can_replace
187
+ def can_replace_layer(
188
+ cls,
189
+ source_layer: nn.Module,
190
+ lora_config: LoRAConfig,
191
+ packed_modules_list: list,
192
+ model_config: Optional[PretrainedConfig],
193
+ ) -> bool:
194
+ # specifying kwargs so they can be easily accessed in decorator
195
+ return super().can_replace_layer(
196
+ source_layer=source_layer,
197
+ lora_config=lora_config,
198
+ packed_modules_list=packed_modules_list,
199
+ model_config=model_config,
200
+ decorate=False,
201
+ )
vllm/lora/layers/utils.py
new
+60
@@ -0,0 +1,60 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from dataclasses import dataclass
4
+
5
+import torch
6
+import torch.nn as nn
7
+
8
+from vllm.adapter_commons.layers import AdapterMapping
9
+
10
+
11
+@dataclass
12
+class LoRAMapping(AdapterMapping):
13
+ is_prefill: bool = False
14
+
15
+
16
+def _get_lora_device(base_layer: nn.Module) -> torch.device:
17
+ # code borrowed from https://github.com/fmmoret/vllm/blob/fm-support-lora-on-quantized-models/vllm/lora/layers.py#L34
18
+ """Returns the device for where to place the LoRA tensors."""
19
+ # unquantizedLinear
20
+ if hasattr(base_layer, "weight"):
21
+ return base_layer.weight.device
22
+ # Compressed Tensor
23
+ elif hasattr(base_layer, "weight_packed"):
24
+ return base_layer.weight_packed.device
25
+ # GPTQ/AWQ
26
+ elif hasattr(base_layer, "qweight"):
27
+ return base_layer.qweight.device
28
+ # HQQ marlin
29
+ elif hasattr(base_layer, "W_q"):
30
+ return base_layer.W_q.device
31
+ else:
32
+ raise ValueError(f"Unsupported base layer: {base_layer}")
33
+
34
+
35
+def _not_fully_sharded_can_replace(can_replace):
36
+ """
37
+ decorator which adds the condition of not using fully sharded loras
38
+ intended to wrap can_replace_layer()
39
+ """
40
+
41
+ def dec(*args, **kwargs):
42
+ decorate = kwargs.pop("decorate") if "decorate" in kwargs else True
43
+ condition = (not kwargs["lora_config"].fully_sharded_loras
44
+ if decorate else True)
45
+ return can_replace(*args, **kwargs) and condition
46
+
47
+ return dec
48
+
49
+
50
+def _fully_sharded_can_replace(can_replace):
51
+ """
52
+ decorator which adds the condition of fully sharded loras
53
+ intended to wrap can_replace_layer()
54
+ """
55
+
56
+ def dec(*args, **kwargs):
57
+ return (can_replace(*args, **kwargs)
58
+ and kwargs["lora_config"].fully_sharded_loras)
59
+
60
+ return dec
vllm/lora/layers/vocal_parallel_embedding.py
new
+172
@@ -0,0 +1,172 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from typing import Optional
5
+
6
+import torch
7
+import torch.nn as nn
8
+import torch.nn.functional as F
9
+from transformers import PretrainedConfig
10
+
11
+from vllm.config import LoRAConfig
12
+from vllm.model_executor.layers.vocab_parallel_embedding import (
13
+ VocabParallelEmbedding)
14
+from vllm.platforms import current_platform
15
+
16
+from .base import BaseLayerWithLoRA
17
+
18
+
19
+class VocabParallelEmbeddingWithLoRA(BaseLayerWithLoRA):
20
+
21
+ def __init__(self, base_layer: VocabParallelEmbedding) -> None:
22
+ super().__init__()
23
+ self.base_layer = base_layer
24
+ self.embeddings_slice: Optional[tuple[int, int]]
25
+ self.embeddings_weights: Optional[torch.Tensor]
26
+
27
+ def create_lora_weights(
28
+ self,
29
+ max_loras: int,
30
+ lora_config: LoRAConfig,
31
+ model_config: Optional[PretrainedConfig] = None) -> None:
32
+
33
+ if self.base_layer.num_added_embeddings_per_partition > 0:
34
+ # We can start adding lora weights
35
+ self.embeddings_weights = self.base_layer.weight.data[
36
+ self.base_layer.num_org_embeddings_per_partition:self.
37
+ base_layer.num_org_embeddings_per_partition +
38
+ self.base_layer.num_added_embeddings_per_partition]
39
+ self.embeddings_slice = (
40
+ self.base_layer.shard_indices.added_vocab_start_index -
41
+ self.base_layer.org_vocab_size,
42
+ self.base_layer.shard_indices.added_vocab_end_index -
43
+ self.base_layer.org_vocab_size)
44
+ self.base_layer.weight.data[
45
+ self.base_layer.num_org_embeddings_per_partition:].fill_(0)
46
+ else:
47
+ self.embeddings_slice = None
48
+ self.embeddings_weights = None
49
+
50
+ self.embeddings_tensors = torch.zeros(
51
+ (
52
+ max_loras,
53
+ lora_config.lora_extra_vocab_size,
54
+ self.base_layer.embedding_dim,
55
+ ),
56
+ dtype=self.base_layer.weight.dtype,
57
+ device=self.base_layer.weight.device,
58
+ )
59
+ self.lora_a_stacked = torch.zeros(
60
+ (
61
+ max_loras,
62
+ self.base_layer.org_vocab_size +
63
+ lora_config.lora_extra_vocab_size,
64
+ lora_config.max_lora_rank,
65
+ ),
66
+ dtype=lora_config.lora_dtype,
67
+ device=self.base_layer.weight.device,
68
+ )
69
+ self.lora_b_stacked = torch.zeros(
70
+ (
71
+ max_loras,
72
+ 1,
73
+ self.base_layer.embedding_dim,
74
+ lora_config.max_lora_rank,
75
+ ),
76
+ dtype=lora_config.lora_dtype,
77
+ device=self.base_layer.weight.device,
78
+ )
79
+ self.lora_a_stacked_2d = self.lora_a_stacked.view(
80
+ self.lora_a_stacked.shape[0] * self.lora_a_stacked.shape[1],
81
+ self.lora_a_stacked.shape[2],
82
+ )
83
+
84
+ def reset_lora(self, index: int):
85
+ self.lora_a_stacked[index] = 0
86
+ self.lora_b_stacked[index] = 0
87
+ self.embeddings_tensors[index] = 0
88
+
89
+ def set_lora(
90
+ self,
91
+ index: int,
92
+ lora_a: torch.Tensor,
93
+ lora_b: torch.Tensor,
94
+ embeddings_tensor: Optional[torch.Tensor],
95
+ bias: Optional[torch.Tensor] = None,
96
+ ):
97
+ self.reset_lora(index)
98
+ self.lora_a_stacked[index, :lora_a.shape[0], :lora_a.shape[1]].copy_(
99
+ lora_a, non_blocking=True)
100
+ self.lora_b_stacked[index,
101
+ 0, :lora_b.shape[1], :lora_b.shape[0]].copy_(
102
+ lora_b.T, non_blocking=True)
103
+ if embeddings_tensor is not None:
104
+ self.embeddings_tensors[
105
+ index,
106
+ :embeddings_tensor.shape[0],
107
+ :embeddings_tensor.shape[1],
108
+ ].copy_(embeddings_tensor, non_blocking=True)
109
+ if self.embeddings_slice is not None:
110
+ # TODO(yard1): Optimize this copy, we don't need to copy
111
+ # everything, just the modified part
112
+ embeddings = self.embeddings_tensors.view(
113
+ self.embeddings_tensors.shape[0] *
114
+ self.embeddings_tensors.shape[1],
115
+ self.embeddings_tensors.shape[2],
116
+ )[self.embeddings_slice[0]:self.embeddings_slice[1]]
117
+ assert self.embeddings_weights is not None
118
+ self.embeddings_weights[:embeddings.shape[0]].copy_(embeddings)
119
+
120
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
121
+ added_tokens_mask = torch.where(x > self.base_layer.org_vocab_size - 1,
122
+ 1, 0)
123
+
124
+ # NB: Don't use torch.narrow here. torch.narrow triggers some
125
+ # Dynamic Shape specialization in torch.compile
126
+ num_tokens = x.shape[0]
127
+ indices_1 = self.punica_wrapper._embeddings_indices[1][:num_tokens]
128
+ indices_0 = self.punica_wrapper._embeddings_indices[0][:num_tokens]
129
+
130
+ full_lora_a_embeddings = F.embedding(
131
+ x + indices_1,
132
+ self.lora_a_stacked_2d,
133
+ )
134
+ full_output = self.base_layer.forward(x +
135
+ (indices_0 * added_tokens_mask))
136
+
137
+ full_output_org = full_output
138
+ if full_output.ndim == 3:
139
+ full_output = full_output.view(
140
+ full_output.shape[0] * full_output.shape[1], -1)
141
+ if full_lora_a_embeddings.ndim == 3:
142
+ full_lora_a_embeddings = full_lora_a_embeddings.view(
143
+ full_lora_a_embeddings.shape[0] *
144
+ full_lora_a_embeddings.shape[1],
145
+ -1,
146
+ )
147
+
148
+ lora_output: Optional[
149
+ torch.Tensor] = self.punica_wrapper.add_lora_embedding(
150
+ full_output,
151
+ full_lora_a_embeddings,
152
+ self.lora_b_stacked,
153
+ add_input=True)
154
+
155
+ if not current_platform.can_update_inplace():
156
+ full_output = lora_output
157
+
158
+ return full_output.view_as(full_output_org)
159
+
160
+ @classmethod
161
+ def can_replace_layer(
162
+ cls,
163
+ source_layer: nn.Module,
164
+ lora_config: LoRAConfig,
165
+ packed_modules_list: list,
166
+ model_config: Optional[PretrainedConfig],
167
+ ) -> bool:
168
+ return type(source_layer) is VocabParallelEmbedding
169
+
170
+ @property
171
+ def weight(self):
172
+ return self.base_layer.weight
vllm/lora/utils.py
+5
-5
@@ -13,21 +13,21 @@ from transformers import PretrainedConfig
13
14
from vllm.config import LoRAConfig
15
from vllm.logger import init_logger
16
-from vllm.lora.fully_sharded_layers import (
17
- ColumnParallelLinearWithShardedLoRA,
18
- MergedColumnParallelLinearWithShardedLoRA,
19
- MergedQKVParallelLinearWithShardedLoRA, QKVParallelLinearWithShardedLoRA,
20
- RowParallelLinearWithShardedLoRA)
16
# being imported for _all_lora_classes below
17
# yapf conflicts with isort for this block
18
# yapf: disable
19
from vllm.lora.layers import (BaseLayerWithLoRA, ColumnParallelLinearWithLoRA,
20
+ ColumnParallelLinearWithShardedLoRA,
21
LogitsProcessorWithLoRA,
22
MergedColumnParallelLinearWithLoRA,
23
+ MergedColumnParallelLinearWithShardedLoRA,
24
MergedQKVParallelLinearWithLoRA,
25
+ MergedQKVParallelLinearWithShardedLoRA,
26
QKVParallelLinearWithLoRA,
27
+ QKVParallelLinearWithShardedLoRA,
28
ReplicatedLinearWithLoRA,
29
RowParallelLinearWithLoRA,
30
+ RowParallelLinearWithShardedLoRA,
31
VocabParallelEmbeddingWithLoRA)
32
from vllm.model_executor.layers.linear import LinearBase
33