diff --git a/vllm/lora/layers/column_parallel_linear.py b/vllm/lora/layers/column_parallel_linear.py index 01bbc9366d..8a86191b89 100644 --- a/vllm/lora/layers/column_parallel_linear.py +++ b/vllm/lora/layers/column_parallel_linear.py @@ -422,7 +422,10 @@ class QKVParallelLinearWithLoRA(ColumnParallelLinearWithLoRA): packed_modules_list: list, model_config: PretrainedConfig | None = None, ) -> bool: - return type(source_layer) is QKVParallelLinear and len(packed_modules_list) == 1 + return ( + type(source_layer) is maybe_get_oot_by_class(QKVParallelLinear) + and len(packed_modules_list) == 1 + ) class MergedQKVParallelLinearWithLoRA(MergedColumnParallelLinearWithLoRA): @@ -480,7 +483,10 @@ class MergedQKVParallelLinearWithLoRA(MergedColumnParallelLinearWithLoRA): packed_modules_list: list, model_config: PretrainedConfig | None = None, ) -> bool: - return type(source_layer) is QKVParallelLinear and len(packed_modules_list) == 3 + return ( + type(source_layer) is maybe_get_oot_by_class(QKVParallelLinear) + and len(packed_modules_list) == 3 + ) # These following layers are based on the tensor parallelism strategy given in diff --git a/vllm/lora/layers/fused_moe.py b/vllm/lora/layers/fused_moe.py index 1ceb34b714..9cb2fe9147 100644 --- a/vllm/lora/layers/fused_moe.py +++ b/vllm/lora/layers/fused_moe.py @@ -9,6 +9,7 @@ from vllm import envs from vllm.config.lora import LoRAConfig from vllm.distributed.utils import divide from vllm.lora.layers.base import BaseLayerWithLoRA +from vllm.model_executor.custom_op import maybe_get_oot_by_class from vllm.model_executor.layers.fused_moe import FusedMoE from vllm.model_executor.layers.fused_moe.experts.lora_context import MoELoRAContext from vllm.model_executor.layers.fused_moe.fused_moe_modular_method import ( @@ -374,7 +375,8 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA): """Returns True if the layer can be replaced by this LoRA layer.""" # source_layer is FusedMoE - return isinstance(source_layer, FusedMoE) and len(packed_modules_list) == 2 + moe_cls = maybe_get_oot_by_class(FusedMoE) + return isinstance(source_layer, moe_cls) and len(packed_modules_list) == 2 class FusedMoE3DWithLoRA(FusedMoEWithLoRA): @@ -537,4 +539,5 @@ class FusedMoE3DWithLoRA(FusedMoEWithLoRA): ) -> bool: """Returns True if the layer can be replaced by this LoRA layer.""" # source_layer is FusedMoE - return isinstance(source_layer, FusedMoE) and len(packed_modules_list) == 1 + moe_cls = maybe_get_oot_by_class(FusedMoE) + return isinstance(source_layer, moe_cls) and len(packed_modules_list) == 1