94 lines
2.9 KiB
Python
94 lines
2.9 KiB
Python
import sys
|
|
from typing import Any
|
|
|
|
from vllm.logger import logger
|
|
from vllm.model_executor.model_loader.weight_utils import maybe_remap_kv_scale_name
|
|
|
|
|
|
class ImportPatchDecorator:
|
|
"""Import patch decorator"""
|
|
|
|
_patches: dict[str, Any] = {}
|
|
|
|
@classmethod
|
|
def register(cls, module_name):
|
|
"""Decorator for registering module patches"""
|
|
|
|
def decorator(func):
|
|
cls._patches[module_name] = func
|
|
return func
|
|
|
|
return decorator
|
|
|
|
@classmethod
|
|
def apply_patches(cls):
|
|
"""Apply all patches"""
|
|
for module_name, patch_func in cls._patches.items():
|
|
if module_name in sys.modules:
|
|
module = sys.modules[module_name]
|
|
try:
|
|
patch_func(module)
|
|
except Exception as e:
|
|
logger.error("Patch application failed %s: %s", module_name, e)
|
|
|
|
|
|
@ImportPatchDecorator.register("vllm.model_executor.models.deepseek_v2")
|
|
def patch_deepseek(module):
|
|
ori_maybe_remap_kv_scale_name = maybe_remap_kv_scale_name
|
|
|
|
def new_remap(name: str, params_dict: dict):
|
|
name = ori_maybe_remap_kv_scale_name(name, params_dict)
|
|
|
|
replace_scale_names = [
|
|
"fa_q.scale",
|
|
"fa_k.scale",
|
|
"fa_v.scale",
|
|
"fa_q.offset",
|
|
"fa_k.offset",
|
|
"fa_v.offset",
|
|
"indexer.q_rot",
|
|
"indexer.k_rot",
|
|
]
|
|
|
|
for scale_name in replace_scale_names:
|
|
if name.endswith(scale_name):
|
|
remap_name = name.replace(scale_name, f"mla_attn.mla_attn.{scale_name}")
|
|
if remap_name in params_dict:
|
|
return remap_name
|
|
else:
|
|
return remap_name.replace(".mla_attn", "")
|
|
|
|
return name
|
|
|
|
if hasattr(module, "maybe_remap_kv_scale_name"):
|
|
module._original_maybe_remap_kv_scale_name = module.maybe_remap_kv_scale_name
|
|
module.maybe_remap_kv_scale_name = new_remap
|
|
|
|
|
|
@ImportPatchDecorator.register("vllm.model_executor.model_loader.weight_utils")
|
|
def patch_weight_utils(module):
|
|
if "vllm.model_executor.models.deepseek_v2" in sys.modules:
|
|
deepseek = sys.modules["vllm.model_executor.models.deepseek_v2"]
|
|
if hasattr(deepseek, "maybe_remap_kv_scale_name"):
|
|
module.maybe_remap_kv_scale_name = deepseek.maybe_remap_kv_scale_name
|
|
|
|
|
|
original_import = __builtins__["__import__"] # type: ignore
|
|
|
|
|
|
def patched_import(name, globals=None, locals=None, fromlist=(), level=0):
|
|
module = original_import(name, globals, locals, fromlist, level)
|
|
|
|
if name in ImportPatchDecorator._patches:
|
|
try:
|
|
ImportPatchDecorator._patches[name](module)
|
|
except Exception as e:
|
|
logger.error("Patch application failed during import %s: %s", name, e)
|
|
|
|
return module
|
|
|
|
|
|
__builtins__["__import__"] = patched_import
|
|
|
|
ImportPatchDecorator.apply_patches()
|