wehub-resource-sync
53 lines
2.2 KiB
Python
53 lines
2.2 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
from swift.utils.logger import get_logger
|
|
|
|
logger = get_logger()
|
|
|
|
_ORIGINAL_MINDSPEED_TE_CP_CLASS = None
|
|
|
|
|
|
def patch_mindspeed_te_cp_implementation(megatron_args: dict[str, Any]) -> None:
|
|
"""
|
|
Route NPU CP to the legacy MindSpeed TE adaptor when the new strategy factory
|
|
only supports kvallgather.
|
|
"""
|
|
# MindSpeed 0.15.3 replaced the TE context-parallel attention class with a
|
|
# new implementation. That new class does not yet cover all CP algorithms,
|
|
# so the default non-kvallgather path can fail during Megatron training.
|
|
# For those algorithms, temporarily route TE attention back to the legacy
|
|
# MindSpeedCPDotProductAttention adaptor. Once MindSpeed's new CP class has
|
|
# feature parity, this compatibility patch can be removed.
|
|
try:
|
|
import mindspeed.te.pytorch.attention.dot_product_attention.dot_product_attention as ms_te_dpa
|
|
from mindspeed.core.context_parallel.adaptor import MindSpeedCPDotProductAttention
|
|
except ImportError as e:
|
|
logger.warning(f'Failed to import MindSpeed CP modules before repatch: {e}')
|
|
return
|
|
|
|
global _ORIGINAL_MINDSPEED_TE_CP_CLASS
|
|
if _ORIGINAL_MINDSPEED_TE_CP_CLASS is None:
|
|
_ORIGINAL_MINDSPEED_TE_CP_CLASS = getattr(ms_te_dpa, 'MindSpeedTEDotProductAttention', None)
|
|
|
|
if _ORIGINAL_MINDSPEED_TE_CP_CLASS is None:
|
|
logger.warning('MindSpeedTEDotProductAttention is unavailable before repatch; skip CP workaround.')
|
|
return
|
|
|
|
cp_algo = megatron_args.get('context_parallel_algo', 'megatron_cp_algo')
|
|
use_legacy_cp_te = int(megatron_args.get('context_parallel_size', 1)) > 1 and cp_algo != 'kvallgather_cp_algo'
|
|
target_cls = MindSpeedCPDotProductAttention if use_legacy_cp_te else _ORIGINAL_MINDSPEED_TE_CP_CLASS
|
|
|
|
if getattr(ms_te_dpa, 'MindSpeedTEDotProductAttention', None) is target_cls:
|
|
return
|
|
|
|
ms_te_dpa.MindSpeedTEDotProductAttention = target_cls
|
|
logger.info(
|
|
'Patched MindSpeedTEDotProductAttention to %s for context_parallel_size=%s, context_parallel_algo=%s.',
|
|
target_cls.__name__,
|
|
megatron_args.get('context_parallel_size', 1),
|
|
cp_algo,
|
|
)
|