column_parallel_linear.py 19.4 KB
Newer Older
Jee Jee Li's avatar
Jee Jee Li committed
1
2
3
4
5
6
7
8
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project


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

9
from vllm.config.lora import LoRAConfig
10
from vllm.distributed import tensor_model_parallel_all_gather
Jee Jee Li's avatar
Jee Jee Li committed
11
from vllm.distributed.utils import divide
12
13
14
15
16
from vllm.model_executor.layers.linear import (
    ColumnParallelLinear,
    MergedColumnParallelLinear,
    QKVParallelLinear,
)
Jee Jee Li's avatar
Jee Jee Li committed
17
18
19
20
21
22
23
from vllm.platforms import current_platform

from .base_linear import BaseLinearLayerWithLoRA
from .utils import _fully_sharded_can_replace, _not_fully_sharded_can_replace


def _mcp_apply(x, bias, layer: "ColumnParallelLinearWithLoRA"):
24
25
    """
    For `ColumnParallelLinearWithLoRA` or classes that inherit from
Jee Jee Li's avatar
Jee Jee Li committed
26
27
    `ColumnParallelLinearWithLoRA`, they share the same `apply` logic.
    """
28
29
30
31
32
33
    assert (
        layer.n_slices
        == len(layer.lora_a_stacked)
        == len(layer.lora_b_stacked)
        == len(layer.output_slices)
    )
Jee Jee Li's avatar
Jee Jee Li committed
34
35
36
37
38
39
40
41
42
43
44
45
46
47

    output = layer.base_layer.quant_method.apply(layer.base_layer, x, bias)

    x = x.view(-1, x.shape[-1])
    output, out_orig_shape = output.view(-1, output.shape[-1]), output.shape

    # Since communication is needed, the buffer is directly initialized as a
    # tensor rather than a tuple of tensor.
    buffers = torch.zeros(
        (layer.n_slices, x.shape[0], layer.lora_a_stacked[0].shape[2]),
        dtype=torch.float32,
        device=x.device,
    )

48
    shrunk_buffers: torch.Tensor | None = layer.punica_wrapper.add_shrink(
49
50
        buffers, x, layer.lora_a_stacked, 1.0
    )
Jee Jee Li's avatar
Jee Jee Li committed
51
52
53
54
55
56

    if not current_platform.can_update_inplace():
        buffers = shrunk_buffers

    buffers = tensor_model_parallel_all_gather(buffers)

57
    lora_output: torch.Tensor | None = layer.punica_wrapper.add_expand(
Jee Jee Li's avatar
Jee Jee Li committed
58
59
60
61
62
        output,
        buffers,
        layer.lora_b_stacked,
        layer.output_slices,
        offset_start=0,
63
64
        add_input=True,
    )
Jee Jee Li's avatar
Jee Jee Li committed
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87

    if not current_platform.can_update_inplace():
        output = lora_output

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


class ColumnParallelLinearWithLoRA(BaseLinearLayerWithLoRA):
    """
    LoRA on top of ColumnParallelLinear layer.
    LoRA B is sliced for tensor parallelism.
    There are two types for the `base_layer`:
    1. ColumnParallelLinear, e.g.`dense_h_to_4h` in `FalconForCausalLM`.
    2. MergedColumnParallelLinear, e.g.`gate_up_proj` in `Phi3ForCausalLM`.
    """

    def __init__(self, base_layer: ColumnParallelLinear) -> None:
        super().__init__(base_layer)
        # The base_layer type is ColumnParallelLinear or
        # MergedColumnParallelLinear, their weight sharding logic is
        # inconsistent when TP is greater than 1.
88
        self.is_merged_col_linear = type(base_layer) is MergedColumnParallelLinear
Jee Jee Li's avatar
Jee Jee Li committed
89
90
91
92
93
94
95
96
97
98
99
100
        self.output_size = self.base_layer.output_size_per_partition
        # There is only one LoRA layer
        self.n_slices = 1

    def slice_lora_a(self, lora_a: torch.Tensor) -> torch.Tensor:
        return lora_a

    def slice_lora_b(self, lora_b: torch.Tensor) -> torch.Tensor:
        # Applicable to cases where the base_layer is
        # MergedColumnParallelLinear.
        if self.is_merged_col_linear:
            shard_size = self.output_size // 2
101
            offset = lora_b.shape[0] // 2
Jee Jee Li's avatar
Jee Jee Li committed
102

103
104
105
106
107
108
109
110
            left_weight = lora_b[
                self.tp_rank * shard_size : (self.tp_rank + 1) * shard_size, :
            ]
            right_weight = lora_b[
                offset + self.tp_rank * shard_size : offset
                + (self.tp_rank + 1) * shard_size,
                :,
            ]
111
            lora_b = torch.cat([left_weight, right_weight], dim=0)
Jee Jee Li's avatar
Jee Jee Li committed
112
113
114
115
        # Applicable to cases where the base_layer is
        # ColumnParallelLinear.
        else:
            shard_size = self.output_size
116
117
            start_idx = self.tp_rank * shard_size
            end_idx = (self.tp_rank + 1) * shard_size
118
            lora_b = lora_b[start_idx:end_idx, :]
Jee Jee Li's avatar
Jee Jee Li committed
119
120
121
122
        return lora_b

    def forward(
        self, input_: torch.Tensor
123
    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor | None]:
Jee Jee Li's avatar
Jee Jee Li committed
124
125
126
127
128
129
130
131
132
        """Forward of ColumnParallelLinear

        Args:
            input_: Tensor whose last dimension is `input_size`.

        Returns:
            - output
            - bias
        """
133
        bias = self.base_layer.bias if not self.base_layer.skip_bias_add else None
Jee Jee Li's avatar
Jee Jee Li committed
134
135
136

        # Matrix multiply.
        output_parallel = self.apply(input_, bias)
137
        if self.base_layer.gather_output and self.tp_size > 1:
Jee Jee Li's avatar
Jee Jee Li committed
138
139
140
141
142
143
144
145
            # All-gather across the partitions.
            output = tensor_model_parallel_all_gather(output_parallel)
        else:
            output = output_parallel

        if not self.base_layer.return_bias:
            return output

146
        output_bias = self.base_layer.bias if self.base_layer.skip_bias_add else None
Jee Jee Li's avatar
Jee Jee Li committed
147
148
149
150
151
152
153
154
155
        return output, output_bias

    @classmethod
    @_not_fully_sharded_can_replace
    def can_replace_layer(
        cls,
        source_layer: nn.Module,
        lora_config: LoRAConfig,
        packed_modules_list: list,
156
        model_config: PretrainedConfig | None,
Jee Jee Li's avatar
Jee Jee Li committed
157
158
159
    ) -> bool:
        return type(source_layer) is ColumnParallelLinear or (
            type(source_layer) is MergedColumnParallelLinear
160
161
            and len(packed_modules_list) == 1
        )
Jee Jee Li's avatar
Jee Jee Li committed
162
163
164
165
166
167
168
169
170
171
172
173


class MergedColumnParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):
    """ColumnParallelLinear layer that is composed of 2 sublayers (slices)
    packed together (e.g. gate_proj + up_proj -> gate_up_proj).

    This means we have 2 LoRAs, each applied to one half of the layer.

    Both slices must have the same size.
    """

    def __init__(
174
        self, base_layer: MergedColumnParallelLinear | QKVParallelLinear
175
    ) -> None:
Jee Jee Li's avatar
Jee Jee Li committed
176
177
178
179
180
181
        super().__init__(base_layer)
        # There are two LoRA layers
        # the output_sizes in MergedColumnParallelLinear is not sharded by tp
        # we need to divide it by the tp_size to get correct slices size
        output_sizes = self.base_layer.output_sizes
        self.output_slices = tuple(
182
183
            divide(output_size, self.tp_size) for output_size in output_sizes
        )
Jee Jee Li's avatar
Jee Jee Li committed
184
        self.n_slices = len(self.output_slices)
185
        self.output_ids = (self.tp_rank,) * self.n_slices
Jee Jee Li's avatar
Jee Jee Li committed
186
187
188
189
190

    def create_lora_weights(
        self,
        max_loras: int,
        lora_config: LoRAConfig,
191
        model_config: PretrainedConfig | None = None,
Jee Jee Li's avatar
Jee Jee Li committed
192
193
    ) -> None:
        """
194
        The main reason for overriding this function is to enhance  code
Jee Jee Li's avatar
Jee Jee Li committed
195
196
197
198
199
        maintainability.
        """
        self.lora_config = lora_config

        lora_a_output_size_per_partition = (
200
201
202
203
            lora_config.max_lora_rank
            if not lora_config.fully_sharded_loras
            else divide(lora_config.max_lora_rank, self.tp_size)
        )
Jee Jee Li's avatar
Jee Jee Li committed
204
205
206
207
208
209
210
211
212

        self.lora_a_stacked = tuple(
            torch.zeros(
                max_loras,
                1,
                lora_a_output_size_per_partition,
                self.input_size,
                dtype=lora_config.lora_dtype,
                device=self.device,
213
214
215
            )
            for _ in range(self.n_slices)
        )
Jee Jee Li's avatar
Jee Jee Li committed
216
217
218
219
220
221
222
223
        self.lora_b_stacked = tuple(
            torch.zeros(
                max_loras,
                1,
                output_size,
                lora_config.max_lora_rank,
                dtype=lora_config.lora_dtype,
                device=self.device,
224
225
226
            )
            for output_size in self.output_slices
        )
Jee Jee Li's avatar
Jee Jee Li committed
227
228

    def slice_lora_a(
229
230
        self, lora_a: list[torch.Tensor | None]
    ) -> list[torch.Tensor | None]:
Jee Jee Li's avatar
Jee Jee Li committed
231
232
233
        return lora_a

    def slice_lora_b(
234
235
        self, lora_b: list[torch.Tensor | None]
    ) -> list[torch.Tensor | None]:
Jee Jee Li's avatar
Jee Jee Li committed
236
237
        sliced_lora_b = [None] * self.n_slices
        for i, (shard_id, shard_size) in enumerate(
238
239
            zip(self.output_ids, self.output_slices)
        ):
Jee Jee Li's avatar
Jee Jee Li committed
240
            if (lora_b_i := lora_b[i]) is not None:
241
242
243
                sliced_lora_b[i] = lora_b_i[
                    shard_size * shard_id : shard_size * (shard_id + 1), :
                ]
Jee Jee Li's avatar
Jee Jee Li committed
244
245
246
247
248
        return sliced_lora_b

    def set_lora(
        self,
        index: int,
249
250
        lora_a: torch.Tensor | list[torch.Tensor],
        lora_b: torch.Tensor | list[torch.Tensor],
Jee Jee Li's avatar
Jee Jee Li committed
251
252
253
254
255
256
257
258
259
260
    ):
        self.reset_lora(index)

        if self.tp_size > 1:
            lora_a = self.slice_lora_a(lora_a)
            lora_b = self.slice_lora_b(lora_b)

        for i in range(self.n_slices):
            if (lora_a_i := lora_a[i]) is not None:
                self.lora_a_stacked[i][
261
262
                    index, 0, : lora_a_i.shape[0], : lora_a_i.shape[1]
                ].copy_(lora_a_i, non_blocking=True)
Jee Jee Li's avatar
Jee Jee Li committed
263
264
            if (lora_b_i := lora_b[i]) is not None:
                self.lora_b_stacked[i][
265
266
                    index, 0, : lora_b_i.shape[0], : lora_b_i.shape[1]
                ].copy_(lora_b_i, non_blocking=True)
Jee Jee Li's avatar
Jee Jee Li committed
267
268
269
270
271
272
273
274

    @classmethod
    @_not_fully_sharded_can_replace
    def can_replace_layer(
        cls,
        source_layer: nn.Module,
        lora_config: LoRAConfig,
        packed_modules_list: list,
275
        model_config: PretrainedConfig | None,
Jee Jee Li's avatar
Jee Jee Li committed
276
    ) -> bool:
277
278
279
280
        return (
            type(source_layer) is MergedColumnParallelLinear
            and len(packed_modules_list) == 2
        )
Jee Jee Li's avatar
Jee Jee Li committed
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297


class QKVParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):
    """
    ColumnParallelLinear layer that is specifically designed for
    qkv_proj. Certain models, such as chatglm3 and baichuan-7b,
    only contains a single LoRA within their qkv_proj layer.

    During inference with Tensor Parallel, the weights of lora_b
    must be accurately partitioned according to the respective ranks.

    Q slice may have different shape than K and V slices (which both have
    the same shape).
    """

    def __init__(self, base_layer: QKVParallelLinear) -> None:
        super().__init__(base_layer)
298
299
300
301
302
303
304
305
306
307
        self.q_proj_total_size = (
            self.base_layer.total_num_heads * self.base_layer.head_size
        )
        self.q_proj_shard_size = self.base_layer.num_heads * self.base_layer.head_size
        self.kv_proj_shard_size = (
            self.base_layer.num_kv_heads * self.base_layer.head_size
        )
        self.kv_proj_total_size = (
            self.base_layer.total_num_kv_heads * self.base_layer.head_size
        )
Jee Jee Li's avatar
Jee Jee Li committed
308
309
310
311
        # There is only one LoRA layer
        self.n_slices = 1

    def slice_lora_b(self, lora_b: torch.Tensor) -> torch.Tensor:
312
313
        self.q_shard_id = self.tp_rank
        self.kv_shard_id = self.tp_rank // self.base_layer.num_kv_head_replicas
314
315
316
317
318
        lora_b_q = lora_b[
            self.q_proj_shard_size * self.q_shard_id : self.q_proj_shard_size
            * (self.q_shard_id + 1),
            :,
        ]
Jee Jee Li's avatar
Jee Jee Li committed
319
        k_offset = self.q_proj_total_size
320
321
322
323
324
        lora_b_k = lora_b[
            k_offset + self.kv_proj_shard_size * self.kv_shard_id : k_offset
            + self.kv_proj_shard_size * (self.kv_shard_id + 1),
            :,
        ]
Jee Jee Li's avatar
Jee Jee Li committed
325
        v_offset = k_offset + self.kv_proj_total_size
326
327
328
329
330
        lora_b_v = lora_b[
            v_offset + self.kv_proj_shard_size * self.kv_shard_id : v_offset
            + self.kv_proj_shard_size * (self.kv_shard_id + 1),
            :,
        ]
331
        lora_b = torch.cat([lora_b_q, lora_b_k, lora_b_v], dim=0)
Jee Jee Li's avatar
Jee Jee Li committed
332
333
334
335
        return lora_b

    @classmethod
    @_not_fully_sharded_can_replace
336
337
338
339
340
    def can_replace_layer(
        cls,
        source_layer: nn.Module,
        lora_config: LoRAConfig,
        packed_modules_list: list,
341
        model_config: PretrainedConfig | None,
342
343
    ) -> bool:
        return type(source_layer) is QKVParallelLinear and len(packed_modules_list) == 1
Jee Jee Li's avatar
Jee Jee Li committed
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361


class MergedQKVParallelLinearWithLoRA(MergedColumnParallelLinearWithLoRA):
    """MergedColumnParallelLinear layer that is composed of 3 sublayers (slices)
    packed together in qkv proj fashion
    (q_proj + k_proj + v_proj -> qkv_proj).

    This means we have 3 LoRAs, each applied to one slice of the layer.

    Q slice may have different shape than K and V slices (which both have
    the same shape).
    """

    def __init__(self, base_layer: QKVParallelLinear) -> None:
        super().__init__(base_layer)
        # There are three LoRA layer.
        self.n_slices = len(self.base_layer.output_sizes)

362
363
364
365
        self.q_proj_shard_size = self.base_layer.num_heads * self.base_layer.head_size
        self.kv_proj_shard_size = (
            self.base_layer.num_kv_heads * self.base_layer.head_size
        )
Jee Jee Li's avatar
Jee Jee Li committed
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
        self.q_shard_id = self.tp_rank
        self.kv_shard_id = self.tp_rank // self.base_layer.num_kv_head_replicas

        self.output_slices = (
            self.q_proj_shard_size,
            self.kv_proj_shard_size,
            self.kv_proj_shard_size,
        )
        self.output_ids = (
            self.q_shard_id,
            self.kv_shard_id,
            self.kv_shard_id,
        )

    def create_lora_weights(
        self,
        max_loras: int,
        lora_config: LoRAConfig,
384
        model_config: PretrainedConfig | None = None,
Jee Jee Li's avatar
Jee Jee Li committed
385
386
    ) -> None:
        """
387
        The main reason for overloading this function is to handle inconsistent
Jee Jee Li's avatar
Jee Jee Li committed
388
389
390
391
392
393
394
395
396
397
398
        weight dimensions in qkv lora.
        """
        super().create_lora_weights(max_loras, lora_config, model_config)

    @classmethod
    @_not_fully_sharded_can_replace
    def can_replace_layer(
        cls,
        source_layer: nn.Module,
        lora_config: LoRAConfig,
        packed_modules_list: list,
399
        model_config: PretrainedConfig | None,
Jee Jee Li's avatar
Jee Jee Li committed
400
    ) -> bool:
401
        return type(source_layer) is QKVParallelLinear and len(packed_modules_list) == 3
Jee Jee Li's avatar
Jee Jee Li committed
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422


# These following 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.
    """

    # For all LoRA layers where the `base_layer` is `ColumnParallelLinear`,
    # their `lora_a` and `lora_b` have different sharding patterns. After
    # completing the `lora_a` GEMM , a gather operation is performed.
    # Therefore, the sharding of `lora_a` only needs to correspond with the
    # gather operation.
    def slice_lora_a(self, lora_a: torch.Tensor) -> torch.Tensor:
        shard_size = self.lora_a_stacked[0].shape[2]
423
        start_idx = self.tp_rank * shard_size
424
        lora_a = lora_a[start_idx : start_idx + shard_size, :]
Jee Jee Li's avatar
Jee Jee Li committed
425
426
        return lora_a

427
    def apply(self, x: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor:
Jee Jee Li's avatar
Jee Jee Li committed
428
429
430
431
432
433
434
435
436
        return _mcp_apply(x, bias, self)

    @classmethod
    @_fully_sharded_can_replace
    def can_replace_layer(
        cls,
        source_layer: nn.Module,
        lora_config: LoRAConfig,
        packed_modules_list: list,
437
        model_config: PretrainedConfig | None,
Jee Jee Li's avatar
Jee Jee Li committed
438
439
440
441
442
443
444
445
446
447
448
    ) -> 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,
        )


449
class MergedColumnParallelLinearWithShardedLoRA(MergedColumnParallelLinearWithLoRA):
Jee Jee Li's avatar
Jee Jee Li committed
450
451
452
453
454
455
456
457
    """
    Differs from MergedColumnParallelLinearWithLoRA by slicing the
    LoRA A's also.

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

    def slice_lora_a(
458
459
        self, lora_a: list[torch.Tensor | None]
    ) -> list[torch.Tensor | None]:
460
        # NOTE: lora_a contains 2 subloras, and each sublora could be None.
Jee Jee Li's avatar
Jee Jee Li committed
461
462
463
        output_shard_size = self.lora_a_stacked[0].shape[2]
        output_start_idx = self.tp_rank * output_shard_size
        lora_a = [
464
465
466
467
468
469
            lora_a[0][output_start_idx : output_start_idx + output_shard_size, :]
            if lora_a[0] is not None
            else None,
            lora_a[1][output_start_idx : output_start_idx + output_shard_size, :]
            if lora_a[1] is not None
            else None,
Jee Jee Li's avatar
Jee Jee Li committed
470
471
472
        ]
        return lora_a

473
    def apply(self, x: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor:
Jee Jee Li's avatar
Jee Jee Li committed
474
475
476
477
478
479
480
481
482
        return _mcp_apply(x, bias, self)

    @classmethod
    @_fully_sharded_can_replace
    def can_replace_layer(
        cls,
        source_layer: nn.Module,
        lora_config: LoRAConfig,
        packed_modules_list: list,
483
        model_config: PretrainedConfig | None,
Jee Jee Li's avatar
Jee Jee Li committed
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
    ) -> 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 QKVParallelLinearWithShardedLoRA(QKVParallelLinearWithLoRA):
    """
    Differs from QKVParallelLinearWithLoRA by slicing the
    LoRA A's also.

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

    def slice_lora_a(self, lora_a: torch.Tensor) -> torch.Tensor:
        shard_size = self.lora_a_stacked[0].shape[2]
505
        start_idx = self.tp_rank * shard_size
506
        lora_a = lora_a[start_idx : start_idx + shard_size, :]
Jee Jee Li's avatar
Jee Jee Li committed
507
508
        return lora_a

509
    def apply(self, x: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor:
Jee Jee Li's avatar
Jee Jee Li committed
510
511
512
513
        return _mcp_apply(x, bias, self)

    @classmethod
    @_fully_sharded_can_replace
514
515
516
517
518
    def can_replace_layer(
        cls,
        source_layer: nn.Module,
        lora_config: LoRAConfig,
        packed_modules_list: list,
519
        model_config: PretrainedConfig | None,
520
    ) -> bool:
Jee Jee Li's avatar
Jee Jee Li committed
521
522
523
524
525
526
527
528
529
530
531
532
        # 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):
    """
533
    Differs from MergedQKVParallelLinearWithLoRA by slicing the
Jee Jee Li's avatar
Jee Jee Li committed
534
535
536
537
538
539
    LoRA A's also.

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

    def slice_lora_a(
540
541
        self, lora_a: list[torch.Tensor | None]
    ) -> list[torch.Tensor | None]:
Jee Jee Li's avatar
Jee Jee Li committed
542
543
544
545
        # NOTE: lora_a contains 3 subloras, and each sublora could be None.
        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 = [
546
547
548
549
550
551
552
553
554
            lora_a[0][start_idx[0] : start_idx[0] + shard_size[0], :]
            if lora_a[0] is not None
            else None,
            lora_a[1][start_idx[1] : start_idx[1] + shard_size[1], :]
            if lora_a[1] is not None
            else None,
            lora_a[2][start_idx[2] : start_idx[2] + shard_size[2], :]
            if lora_a[2] is not None
            else None,
Jee Jee Li's avatar
Jee Jee Li committed
555
556
557
        ]
        return lora_a

558
    def apply(self, x: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor:
Jee Jee Li's avatar
Jee Jee Li committed
559
560
561
562
563
564
565
566
567
        return _mcp_apply(x, bias, self)

    @classmethod
    @_fully_sharded_can_replace
    def can_replace_layer(
        cls,
        source_layer: nn.Module,
        lora_config: LoRAConfig,
        packed_modules_list: list,
568
        model_config: PretrainedConfig | None,
Jee Jee Li's avatar
Jee Jee Li committed
569
570
571
572
573
574
575
576
577
    ) -> 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,
        )