chore: import upstream snapshot with attribution

2026-07-13 13:34:58 +08:00
commit a203934033
1368 changed files with 175001 additions and 0 deletions
@@ -0,0 +1,52 @@
+# 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,
+    )