89 lines
2.8 KiB
Python
89 lines
2.8 KiB
Python
import vllm
|
|
from torch import nn
|
|
from transformers import PretrainedConfig
|
|
from vllm.config import LoRAConfig
|
|
from vllm.lora.layers import (
|
|
MergedQKVParallelLinearWithLoRA,
|
|
MergedQKVParallelLinearWithShardedLoRA,
|
|
QKVParallelLinearWithLoRA,
|
|
QKVParallelLinearWithShardedLoRA,
|
|
)
|
|
from vllm.lora.layers.utils import _fully_sharded_can_replace, _not_fully_sharded_can_replace
|
|
|
|
from vllm_ascend.lora.fused_moe import (
|
|
AscendFusedMoE3DWithLoRA,
|
|
AscendFusedMoEWithLoRA,
|
|
)
|
|
from vllm_ascend.ops.linear import (
|
|
AscendQKVParallelLinear,
|
|
)
|
|
|
|
|
|
class AscendQKVParallelLinearWithLoRA(QKVParallelLinearWithLoRA):
|
|
@classmethod
|
|
@_not_fully_sharded_can_replace
|
|
def can_replace_layer(
|
|
cls,
|
|
source_layer: nn.Module,
|
|
lora_config: LoRAConfig,
|
|
packed_modules_list: list,
|
|
model_config: PretrainedConfig | None,
|
|
) -> bool:
|
|
return type(source_layer) is AscendQKVParallelLinear and len(packed_modules_list) == 1
|
|
|
|
|
|
class AscendMergedQKVParallelLinearWithLoRA(MergedQKVParallelLinearWithLoRA):
|
|
@classmethod
|
|
@_not_fully_sharded_can_replace
|
|
def can_replace_layer(
|
|
cls,
|
|
source_layer: nn.Module,
|
|
lora_config: LoRAConfig,
|
|
packed_modules_list: list,
|
|
model_config: PretrainedConfig | None,
|
|
) -> bool:
|
|
return type(source_layer) is AscendQKVParallelLinear and len(packed_modules_list) == 3
|
|
|
|
|
|
class AscendMergedQKVParallelLinearWithShardedLoRA(MergedQKVParallelLinearWithShardedLoRA):
|
|
@classmethod
|
|
@_fully_sharded_can_replace
|
|
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 AscendQKVParallelLinear and len(packed_modules_list) == 3
|
|
|
|
|
|
class AscendQKVParallelLinearWithShardedLoRA(QKVParallelLinearWithShardedLoRA):
|
|
@classmethod
|
|
@_fully_sharded_can_replace
|
|
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 AscendQKVParallelLinear and len(packed_modules_list) == 1
|
|
|
|
|
|
def refresh_all_lora_classes():
|
|
ascend_classes = (
|
|
AscendQKVParallelLinearWithLoRA,
|
|
AscendMergedQKVParallelLinearWithLoRA,
|
|
AscendMergedQKVParallelLinearWithShardedLoRA,
|
|
AscendQKVParallelLinearWithShardedLoRA,
|
|
AscendFusedMoEWithLoRA,
|
|
AscendFusedMoE3DWithLoRA,
|
|
)
|
|
# vLLM #35077 changed _all_lora_classes from set to ordered tuple.
|
|
# Append the Ascend classes in a deterministic order.
|
|
vllm.lora.utils._all_lora_classes = (
|
|
*ascend_classes,
|
|
*vllm.lora.utils._all_lora_classes,
|
|
)
|