fully_sharded_layers.py 12.2 KB
Newer Older
1
# pylint: disable=unused-argument
2
from typing import TYPE_CHECKING, List, Optional, Union
3
4
5
6
7
8
9
10
11
12
13
14

import torch
import torch.nn as nn
from transformers import PretrainedConfig

from vllm.config import LoRAConfig
from vllm.distributed.communication_op import (
    tensor_model_parallel_all_gather, tensor_model_parallel_all_reduce)
from vllm.distributed.parallel_state import get_tensor_model_parallel_rank
from vllm.lora.layers import (ColumnParallelLinearWithLoRA,
                              MergedColumnParallelLinearWithLoRA,
                              MergedQKVParallelLinearWithLora,
15
                              QKVParallelLinearWithLora,
16
17
18
19
20
21
22
23
24
25
26
27
28
29
                              RowParallelLinearWithLoRA)

if TYPE_CHECKING:
    pass


def _fully_sharded_can_replace(can_replace):
    """
    decorator which adds the condition of fully sharded loras
    intended to wrap can_replace_layer()
    """

    def dec(*args, **kwargs):
        return (can_replace(*args, **kwargs)
30
                and kwargs["lora_config"].fully_sharded_loras)
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53

    return dec


# these layers are based on the tensor parallelism strategy given in
# Y. Sheng et al., S-LoRA: Serving Thousands of Concurrent LoRA Adapters. 2023,
# https://arxiv.org/abs/2311.03285.


class ColumnParallelLinearWithShardedLoRA(ColumnParallelLinearWithLoRA):
    """
    Differs from ColumnParallelLinearWithLoRA by slicing LoRA A also.

    Based on S-LoRA, slicing happens along the rank dim.
    """

    def slice_lora_a(self, lora_a: torch.Tensor) -> torch.Tensor:
        tp_rank = get_tensor_model_parallel_rank()
        shard_size = self.lora_a_stacked.shape[2]
        start_idx = tp_rank * shard_size
        lora_a = lora_a[:, start_idx:start_idx + shard_size]
        return lora_a

54
55
56
    def apply(self, x: torch.Tensor,
              bias: Optional[torch.Tensor]) -> torch.Tensor:
        output = self.base_layer.quant_method.apply(self.base_layer, x, bias)
57
58
59
60

        x = x.view(-1, x.shape[-1])
        output, out_orig_shape = output.view(-1,
                                             output.shape[-1]), output.shape
61
62
63
64
65
66
        buffer = torch.zeros(
            (x.shape[0], self.lora_a_stacked.shape[2]),
            dtype=torch.float32,
            device=x.device,
        )
        self.punica_wrapper.add_shrink(buffer, x, self.lora_a_stacked, 1.0)
67
        buffer = tensor_model_parallel_all_gather(buffer)
68
69
70
71
        self.punica_wrapper.add_expand(output,
                                       buffer,
                                       self.lora_b_stacked,
                                       add_input=True)
72
73
74
75
76
77
        # now have column partitioned output
        output = output.view(*out_orig_shape)
        return output

    @classmethod
    @_fully_sharded_can_replace
78
79
80
81
82
83
84
    def can_replace_layer(
        cls,
        source_layer: nn.Module,
        lora_config: LoRAConfig,
        packed_modules_list: List,
        model_config: Optional[PretrainedConfig],
    ) -> bool:
85
86
87
88
89
90
91
92
93
94
        # specifying kwargs so they can be easily accessed in decorator
        return super().can_replace_layer(
            source_layer=source_layer,
            lora_config=lora_config,
            packed_modules_list=packed_modules_list,
            model_config=model_config,
            decorate=False,
        )


95
def _mcp_apply(x, bias, layer: QKVParallelLinearWithLora):
96
    """
97
98
    MergedColumnParallelLinearWithShardedLoRA and
    MergedQKVParallelLinearWithShardedLora share the same
99
100
101
    LoRa weight application method.
    
    The main difference is the step by shard_size for lora_b which can
102
    vary for MergedQKVParallelLinearWithShardedLora but is constant for
103
104
105
106
    MergedColumnParallelLinearWithShardedLoRA.
    """
    # expecting 2 for column parallel and 3 for qkv
    n = len(layer.lora_a_stacked)
107
    output = layer.base_layer.quant_method.apply(layer.base_layer, x, bias)
108
109
110

    x = x.view(-1, x.shape[-1])
    output, out_orig_shape = output.view(-1, output.shape[-1]), output.shape
111
112
113
114
115
    buffers = torch.zeros(
        (n, x.shape[0], layer.lora_a_stacked[0].shape[2]),
        dtype=torch.float32,
        device=x.device,
    )
116
    for idx in range(n):
117
118
        layer.punica_wrapper.add_shrink(buffers[idx], x,
                                        layer.lora_a_stacked[idx], 1.0)
119
120
121
122
123

    buffers = tensor_model_parallel_all_gather(buffers)
    left_offset = 0
    for idx in range(n):
        shard_size = layer.lora_b_stacked[idx].shape[2]
124
125
126
127
128
129
130
131
        layer.punica_wrapper.add_expand_slice(
            output,
            buffers[idx],
            layer.lora_b_stacked[idx],
            left_offset,
            shard_size,
            add_input=True,
        )
132
133
134
135
136
137
138
139
140
141
        left_offset += shard_size

    output = output.view(*out_orig_shape)
    # now have column partitioned and packed output
    return output


class MergedColumnParallelLinearWithShardedLoRA(
        MergedColumnParallelLinearWithLoRA):
    """
142
    Differs from MergedColumnParallelLinearWithLoRA by slicing the
143
144
145
146
147
    LoRA A's also.

    Based on S-LoRA, slicing happens along the rank dim.
    """

148
149
150
151
152
    def slice_lora_a(
        self, lora_a: List[Union[torch.Tensor, None]]
    ) -> List[Union[torch.Tensor, None]]:
        if lora_a[0] is None or lora_a[1] is None:
            return lora_a
153
154
155
        output_shard_size = self.lora_a_stacked[0].shape[2]
        output_start_idx = self.tp_rank * output_shard_size
        lora_a = [
156
157
            lora_a[0][:,
                      output_start_idx:output_start_idx + output_shard_size],
158
159
            lora_a[1][:,
                      output_start_idx:output_start_idx + output_shard_size],
160
161
162
        ]
        return lora_a

163
164
165
    def apply(self, x: torch.Tensor,
              bias: Optional[torch.Tensor]) -> torch.Tensor:
        return _mcp_apply(x, bias, self)
166
167
168

    @classmethod
    @_fully_sharded_can_replace
169
170
171
172
173
174
175
    def can_replace_layer(
        cls,
        source_layer: nn.Module,
        lora_config: LoRAConfig,
        packed_modules_list: List,
        model_config: Optional[PretrainedConfig],
    ) -> bool:
176
177
178
179
180
181
182
183
184
185
        # specifying kwargs so they can be easily accessed in decorator
        return super().can_replace_layer(
            source_layer=source_layer,
            lora_config=lora_config,
            packed_modules_list=packed_modules_list,
            model_config=model_config,
            decorate=False,
        )


186
class QKVParallelLinearWithShardedLora(QKVParallelLinearWithLora):
187
    """
188
    Differs from QKVParallelLinearWithLora by slicing the
189
190
191
192
193
    LoRA A's also.

    Based on S-LoRA, slicing happens along the rank dim.
    """

194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
    def slice_lora_a(self, lora_a: torch.Tensor) -> torch.Tensor:
        tp_rank = get_tensor_model_parallel_rank()
        shard_size = self.lora_a_stacked.shape[2]
        start_idx = tp_rank * shard_size
        lora_a = lora_a[:, start_idx:start_idx + shard_size]
        return lora_a

    def apply(self, x: torch.Tensor,
              bias: Optional[torch.Tensor]) -> torch.Tensor:
        output = self.base_layer.quant_method.apply(self.base_layer, x, bias)

        x = x.view(-1, x.shape[-1])
        output, out_orig_shape = output.view(-1,
                                             output.shape[-1]), output.shape
        buffer = torch.zeros((x.shape[0], self.lora_a_stacked.shape[2]),
                             dtype=torch.float32,
                             device=x.device)
211
        self.punica_wrapper.add_shrink(buffer, x, self.lora_a_stacked, 1.0)
212
        buffer = tensor_model_parallel_all_gather(buffer)
213
214
215
216
        self.punica_wrapper.add_expand(output,
                                       buffer,
                                       self.lora_b_stacked,
                                       add_input=True)
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
        # now have column partitioned output
        output = output.view(*out_orig_shape)
        return output

    @classmethod
    @_fully_sharded_can_replace
    def can_replace_layer(cls, source_layer: nn.Module,
                          lora_config: LoRAConfig, packed_modules_list: List,
                          model_config: Optional[PretrainedConfig]) -> bool:
        # specifying kwargs so they can be easily accessed in decorator
        return super().can_replace_layer(
            source_layer=source_layer,
            lora_config=lora_config,
            packed_modules_list=packed_modules_list,
            model_config=model_config,
            decorate=False,
        )


class MergedQKVParallelLinearWithShardedLora(MergedQKVParallelLinearWithLora):
    """
    Differs from MergedQKVParallelLinearWithLora by slicing the 
    LoRA A's also.

    Based on S-LoRA, slicing happens along the rank dim.
    """

244
245
246
247
248
    def slice_lora_a(
        self, lora_a: List[Union[torch.Tensor, None]]
    ) -> List[Union[torch.Tensor, None]]:
        if lora_a[0] is None or lora_a[1] is None or lora_a[2] is None:
            return lora_a
249
250
251
        shard_size = [self.lora_a_stacked[i].shape[2] for i in range(3)]
        start_idx = [self.tp_rank * shard_size[i] for i in range(3)]
        lora_a = [
252
253
            lora_a[0][:, start_idx[0]:start_idx[0] + shard_size[0]],
            lora_a[1][:, start_idx[1]:start_idx[1] + shard_size[1]],
254
            lora_a[2][:, start_idx[2]:start_idx[2] + shard_size[2]],
255
256
257
        ]
        return lora_a

258
259
260
    def apply(self, x: torch.Tensor,
              bias: Optional[torch.Tensor]) -> torch.Tensor:
        return _mcp_apply(x, bias, self)
261
262
263

    @classmethod
    @_fully_sharded_can_replace
264
265
266
267
268
269
270
    def can_replace_layer(
        cls,
        source_layer: nn.Module,
        lora_config: LoRAConfig,
        packed_modules_list: List,
        model_config: Optional[PretrainedConfig],
    ) -> bool:
271
272
273
274
275
276
277
278
279
280
281
282
        # specifying kwargs so they can be easily accessed in decorator
        return super().can_replace_layer(
            source_layer=source_layer,
            lora_config=lora_config,
            packed_modules_list=packed_modules_list,
            model_config=model_config,
            decorate=False,
        )


class RowParallelLinearWithShardedLoRA(RowParallelLinearWithLoRA):
    """
283
    Differs from RowParallelLinearWithLoRA by slicing the
284
285
286
    LoRA B's also.

    Based on S-LoRA, slicing happens along the output dim.
287
    This yields a combined partial sum from the row parallel base
288
289
290
291
292
293
294
295
296
297
    layer and column partitioned output from the LoRA.
    """

    def slice_lora_b(self, lora_b: torch.Tensor) -> torch.Tensor:
        shard_size = self.lora_b_stacked.shape[2]
        start_idx = self.tp_rank * shard_size
        end_idx = (self.tp_rank + 1) * shard_size
        lora_b = lora_b[:, start_idx:end_idx]
        return lora_b

298
299
    def apply(self, x: torch.Tensor) -> torch.Tensor:
        output = self.base_layer.quant_method.apply(self.base_layer, x)
300
301
302
303

        x = x.view(-1, x.shape[-1])
        output, out_orig_shape = output.view(-1,
                                             output.shape[-1]), output.shape
304
305
306
307
308
309
310
        buffer = torch.zeros(
            (x.shape[0], self.lora_a_stacked.shape[2]),
            dtype=torch.float32,
            device=x.device,
        )

        self.punica_wrapper.add_shrink(buffer, x, self.lora_a_stacked, 1.0)
311
312
313
314
315
316
317
318
319
320
        buffer = tensor_model_parallel_all_reduce(buffer)

        # following S-LoRA, allows the fusing of all_gather and all_reduce
        # by adding the column partitioned lora output to a slice of output
        # tensor, which is a partial sum due to row parallel. All that
        # remains is a standard all_reduce. User should be aware though that
        # the output is not the same as a normal row_parallel, it should be
        # reduced before being used
        shard_size = self.lora_b_stacked.shape[2]
        start_idx = self.tp_rank * shard_size
321
322
323
        self.punica_wrapper.add_expand_slice(output, buffer,
                                             self.lora_b_stacked, start_idx,
                                             shard_size)
324
325
326
327
328
        output = output.view(*out_orig_shape)
        return output

    @classmethod
    @_fully_sharded_can_replace
329
330
331
332
333
334
335
    def can_replace_layer(
        cls,
        source_layer: nn.Module,
        lora_config: LoRAConfig,
        packed_modules_list: List,
        model_config: Optional[PretrainedConfig],
    ) -> bool:
336
337
338
339
340
341
342
343
        # specifying kwargs so they can be easily accessed in decorator
        return super().can_replace_layer(
            source_layer=source_layer,
            lora_config=lora_config,
            packed_modules_list=packed_modules_list,
            model_config=model_config,
            decorate=False,
        )