wehub-resource-sync
64 lines
2.6 KiB
Python
64 lines
2.6 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
"""Facade for SWIFT's vLLM-Ascend NPU compatibility patches.
|
|
|
|
Keep this file thin. The real patches are split by responsibility:
|
|
|
|
* ``vllm_ascend_moe``: MoE routing and GRPO weight-sync layout handling.
|
|
* ``vllm_ascend_memory``: small torch-npu/vLLM-Ascend memory API compatibility.
|
|
Callers should import from this module so the public entrypoints stay stable,
|
|
while reviewers can audit each patch family in its own file. The caller is
|
|
still responsible for guarding these entrypoints with an NPU/device check.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import sys
|
|
|
|
from swift.model.npu_patch.vllm_ascend_memory import patch_vllm_ascend_memory_runtime
|
|
from swift.model.npu_patch.vllm_ascend_moe import (patch_vllm_ascend_moe_expert_weight_loader,
|
|
patch_vllm_ascend_moe_runtime, should_skip_vllm_ascend_moe_post_load,
|
|
use_vllm_ascend_moe_preprocessed_weight)
|
|
from swift.utils.logger import get_logger
|
|
|
|
logger = get_logger()
|
|
|
|
|
|
def _patch_flash_attn_optional_import() -> None:
|
|
"""Clear a stub ``flash_attn`` module that can block optional imports.
|
|
|
|
Some stacks insert a non-package ``flash_attn`` placeholder into
|
|
``sys.modules``. vLLM import paths then treat it as the real package and
|
|
fail on submodule imports. Removing the placeholder lets normal optional
|
|
dependency checks proceed.
|
|
"""
|
|
module = sys.modules.get('flash_attn')
|
|
if module is None or hasattr(module, '__path__'):
|
|
return
|
|
for module_name in list(sys.modules):
|
|
if module_name == 'flash_attn' or module_name.startswith('flash_attn.'):
|
|
sys.modules.pop(module_name, None)
|
|
|
|
|
|
def patch_vllm_ascend_runtime(*, colocate: bool = False) -> None:
|
|
"""Apply vLLM-Ascend patches needed by SWIFT NPU rollout.
|
|
|
|
``colocate=False`` covers patches that are also safe for standalone
|
|
vLLM-Ascend server/native inference, such as optional import cleanup, MoE
|
|
routing, and ``mem_get_info`` binding compatibility.
|
|
|
|
``colocate`` is kept in the public signature for callers that share this
|
|
entrypoint between server and colocate modes. Process-group creation is
|
|
left to upstream vLLM/vLLM-Ascend; SWIFT only keeps the narrow runtime
|
|
compatibility patches below.
|
|
"""
|
|
_patch_flash_attn_optional_import()
|
|
patch_vllm_ascend_moe_runtime()
|
|
patch_vllm_ascend_memory_runtime()
|
|
|
|
|
|
__all__ = [
|
|
'patch_vllm_ascend_moe_expert_weight_loader',
|
|
'patch_vllm_ascend_runtime',
|
|
'should_skip_vllm_ascend_moe_post_load',
|
|
'use_vllm_ascend_moe_preprocessed_weight',
|
|
]
|