31 lines
1.0 KiB
Python
31 lines
1.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
import torch.nn as nn
|
|
from transformers import PretrainedConfig
|
|
|
|
from vllm.config.lora import LoRAConfig
|
|
from vllm.model_executor.custom_op import maybe_get_oot_by_class
|
|
from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear
|
|
|
|
from .replicated_linear import ReplicatedLinearWithLoRA
|
|
|
|
|
|
class GateLinearWithLoRA(ReplicatedLinearWithLoRA):
|
|
def __init__(self, base_layer: GateLinear) -> None:
|
|
super().__init__(
|
|
base_layer,
|
|
)
|
|
|
|
# GateLinearWithLoRA should always be replaced, regardless of the fully
|
|
# sharded LoRAs setting, because it is, by definition, copied per GPU.
|
|
@classmethod
|
|
def can_replace_layer(
|
|
cls,
|
|
source_layer: nn.Module,
|
|
lora_config: LoRAConfig,
|
|
packed_modules_list: list,
|
|
model_config: PretrainedConfig | None = None,
|
|
) -> bool:
|
|
return type(source_layer) is maybe_get_oot_by_class(GateLinear)
|