Star83
364
代码
代码
Issues57
Pull Requests151
项目讨论
Wiki
分析
Star83
364
已合并
add MindSpeedFeaturesManager #2228
jializheng创建于 2025年4月21日
add MindSpeedFeaturesManager #2228
已合并
jializheng创建于 2025年4月21日
从refs/pull/2228/head合入到master
共 16 个文件变更+348-286
Dmindspeed/arguments_v2.py+0-106
@@ -1,106 +0,0 @@
1-"""Handle cli arguments by features granularity.
2- 
3-Copyright (c) 2025, Huawei Technologies Co., Ltd. All rights reserved.
4-"""
5- 
6-from argparse import ArgumentParser, Namespace
7-from typing import List
8-from functools import wraps
9- 
10-from mindspeed.features_manager import FEATURES_LIST_V2
11-from .arguments import process_args
12- 
13-_ARGS = None
14- 
15- 
16-def extra_args_provider_decorator(extra_args_provider):
17- """Make a extra args parser for magatron."""
18- @wraps(extra_args_provider)
19- def wrapper(parser):
20- if extra_args_provider is not None:
21- parser = extra_args_provider(parser)
22- parser = process_args(parser)
23- for feature in FEATURES_LIST_V2:
24- feature.register_args(parser)
25- return parser
26- 
27- return wrapper
28- 
29- 
30-def parse_args_wrapper(parse_args):
31- """Decorate parse_args function of megatron."""
32- 
33- @wraps(parse_args)
34- def wrapper(extra_args_provider=None, ignore_unknown_args=False):
35- decorated_provider = extra_args_provider_decorator(extra_args_provider)
36- return parse_args(decorated_provider, ignore_unknown_args)
37- 
38- return wrapper
39- 
40- 
41-def core_transformer_config_from_args_wrapper(fn):
42- """A decorator for transformer config args."""
43- @wraps(fn)
44- def wrapper(args):
45- config = fn(args)
46- config.context_parallel_algo = args.context_parallel_algo
47- config.batch_p2p_comm = False
48- if args.use_multiparameter_pipeline_model_parallel:
49- config.deallocate_pipeline_outputs = False
50- return config
51- 
52- return wrapper
53- 
54- 
55-def get_mindspeed_args() -> Namespace:
56- """Get cli arguments of mindspeed."""
57- global _ARGS
58- 
59- if not _ARGS:
60- parser = ArgumentParser(
61- description="MindSpeed Arguments",
62- allow_abbrev=False,
63- )
64- parser = process_args(parser)
65- for feature in FEATURES_LIST_V2:
66- feature.register_args(parser)
67- _ARGS, unknown = parser.parse_known_args()
68- parse_unknown_args(_ARGS, unknown)
69- 
70- return _ARGS
71- 
72- 
73-def add_args(args, key, value):
74- """Add args to parser."""
75- if key is not None:
76- key = key[2:].replace("-", "_")
77- if value is None:
78- value = True
79- elif len(value) == 1:
80- value = value[0]
81- setattr(args, key, value)
82- 
83- 
84-def parse_unknown_args(args: Namespace, unknown: List[str]):
85- """Parse special unknown args.
86- 
87- Args:
88- args (Namespace): regular arguments.
89- unknown (List[str]): special arguments string.
90- """
91- i = 0
92- key, value = None, None
93- while i < len(unknown):
94- if unknown[i].startswith("--"):
95- add_args(args, key, value)
96- key = unknown[i]
97- value = None
98- else:
99- if value is None:
100- value = [unknown[i]]
101- else:
102- value.append(unknown[i])
103- i += 1
104- add_args(args, key, value)
105- 
106- 
Mmindspeed/core/megatron_basic/arguments_basic.py+5-19
@@ -6,7 +6,7 @@ from dataclasses import make_dataclass, field
6from megatron.training import get_args6from megatron.training import get_args
7from megatron.training.arguments import _print_args7from megatron.training.arguments import _print_args
8 8 
9-from mindspeed.features_manager import FEATURES_LIST_V29+from mindspeed.features_manager.features_manager import MindSpeedFeaturesManager
10 10 
11 11 
12def extra_args_provider_decorator(extra_args_provider):12def extra_args_provider_decorator(extra_args_provider):
@@ -15,8 +15,7 @@ def extra_args_provider_decorator(extra_args_provider):
15 def wrapper(parser):15 def wrapper(parser):
16 if extra_args_provider is not None:16 if extra_args_provider is not None:
17 parser = extra_args_provider(parser)17 parser = extra_args_provider(parser)
18- for feature in FEATURES_LIST_V2:18+ MindSpeedFeaturesManager.register_features_args(parser)
19- feature.register_args(parser)
20 return parser19 return parser
21 20 
22 return wrapper21 return wrapper
@@ -41,20 +40,15 @@ def validate_args_wrapper(validate_args):
41 if defaults is None:40 if defaults is None:
42 defaults = {}41 defaults = {}
43 # make prev validation and copy some args.42 # make prev validation and copy some args.
44- origin = _pre_validate(args)43+ MindSpeedFeaturesManager.pre_validate_features_args(args)
45- for feature in FEATURES_LIST_V2:
46- feature.pre_validate_args(args)
47 44 
48 # make megatron args validation then restore args thar are copied.45 # make megatron args validation then restore args thar are copied.
49 args = validate_args(args, defaults)46 args = validate_args(args, defaults)
50 47 
51 # make post validation after megatron validation.48 # make post validation after megatron validation.
52- _post_validate(args, origin)49+ MindSpeedFeaturesManager.post_validate_features_args(args=args)
53- for feature in FEATURES_LIST_V2:
54- feature.post_validate_args(args=args)
55 50 
56- for feature in FEATURES_LIST_V2:51+ MindSpeedFeaturesManager.validate_features_args(args=args)
57- feature.validate_args(args=args)
58 52 
59 args.create_attention_mask_in_dataloader = False53 args.create_attention_mask_in_dataloader = False
60 args.reduce_recompute_for_last_chunk = False54 args.reduce_recompute_for_last_chunk = False
@@ -67,14 +61,6 @@ def validate_args_wrapper(validate_args):
67 return wrapper61 return wrapper
68 62 
69 63 
70-def _pre_validate(_args: Namespace):
71- return (None,)
72- 
73- 
74-def _post_validate(_args: Namespace, _origin):
75- pass
76- 
77- 
78def print_args_wrapper(fn):64def print_args_wrapper(fn):
79 @wraps(fn)65 @wraps(fn)
80 def wrapper(title, args, after_validate=False):66 def wrapper(title, args, after_validate=False):
Mmindspeed/core/megatron_basic/megatron_basic.py+17-0
@@ -1,4 +1,5 @@
1import time1import time
2+from functools import wraps
2 3 
3import torch4import torch
4from torch import _C5from torch import _C
@@ -52,3 +53,19 @@ class PTNorm:
52 raise Exception('Only LayerNorm and RMSNorm are curently supported')53 raise Exception('Only LayerNorm and RMSNorm are curently supported')
53 54 
54 return instance55 return instance
56+ 
57+ 
58+def get_device_wrapper(func):
JialiZheng
JialiZhengJialiZheng2025年4月22日

transfer_to_npu不能cover这种场景么

迁移工具可以将 nccl 修改为 hccl，无法把 hccl 修改为 nccl

likedislike
59+ @wraps(func)
60+ def wrapper(*args, **kwargs):
61+ backend = torch.distributed.get_backend()
62+ local_rank = args[0]
63+ if backend == 'hccl':
64+ if local_rank is None:
65+ device = torch.device('cuda')
66+ else:
67+ device = torch.device(f'cuda:{local_rank}')
68+ else:
69+ device = func(*args, **kwargs)
70+ return device
71+ return wrapper
Mmindspeed/deprecate.py+11-5
@@ -124,15 +124,21 @@ class Deprecated:
124 )124 )
125 125 
126 126 
127-class DisableExecution:127+class AutoExecuteFunction:
128- """A decorator to control a function's execution."""128+ AUTO_EXECUTE = True
129- 
130- DISABLE = False
131 129 
132 def __init__(self, func: Callable):130 def __init__(self, func: Callable):
133 self._func = func131 self._func = func
134 132 
135 def __call__(self, *args, **kwargs):133 def __call__(self, *args, **kwargs):
136- if self.DISABLE:134+ if not AutoExecuteFunction.AUTO_EXECUTE:
137 return None135 return None
138 return self._func(*args, **kwargs)136 return self._func(*args, **kwargs)
137+ 
138+ 
139+class NoExecuteFunction:
140+ def __enter__(self):
141+ AutoExecuteFunction.AUTO_EXECUTE = False
142+ 
143+ def __exit__(self, exc_type, exc_val, exc_tb):
144+ AutoExecuteFunction.AUTO_EXECUTE = True
Mmindspeed/features_manager/__init__.py+137-108
@@ -1,83 +1,45 @@
1-from mindspeed.features_manager.functional.profiler_default import (1+from typing import List
2- ProfilerDefaultFeature,
3-)
4-from mindspeed.features_manager.functional.npu_deterministic import (
5- NPUDeterministicFeature,
6-)
7-from mindspeed.features_manager.functional.tflops_calculate import (
8- TflopsCalculateFeature,
9-)
10-from mindspeed.features_manager.functional.profile import ProfileFeature
11-from mindspeed.features_manager.pipeline_parallel.ripipe_schedules_feature import (
12- RiPipeSchedulesBubbleFeature,
13- RiPipeSchedulesAdvanceFeature,
14-)
15 2 
16-from mindspeed.features_manager.functional.profiler_default import (3+from mindspeed.deprecate import AutoExecuteFunction
17- ProfilerDefaultFeature,4+ 
18-)5+from mindspeed.features_manager.feature import MindSpeedFeature
19-from mindspeed.features_manager.functional.npu_deterministic import (6+from mindspeed.features_manager.features_manager import MindSpeedFeaturesManager
20- NPUDeterministicFeature,7+from mindspeed.features_manager.functional.profile import ProfileFeature
21-)8+from mindspeed.features_manager.functional.profiler_default import ProfilerDefaultFeature
9+from mindspeed.features_manager.functional.npu_deterministic import NPUDeterministicFeature
10+from mindspeed.features_manager.functional.tflops_calculate import TflopsCalculateFeature
22 11 
23from mindspeed.features_manager.fusions.grouped_matmul import GroupedMatmulFeature12from mindspeed.features_manager.fusions.grouped_matmul import GroupedMatmulFeature
24from mindspeed.features_manager.fusions.fused_bias_swiglu import FusedSwigluFeature13from mindspeed.features_manager.fusions.fused_bias_swiglu import FusedSwigluFeature
14+from mindspeed.features_manager.fusions.fused_softmax import FusedSoftmaxFeature
15+from mindspeed.features_manager.hccl_buffer.hccl_buffer_adaptive import HcclBufferAdaptiveFeature
25 16 
26-from mindspeed.features_manager.megatron_basic.requirements_basic import (17+from mindspeed.features_manager.megatron_basic.requirements_basic import RequirementsBasicFeature
27- RequirementsBasicFeature,18+from mindspeed.features_manager.megatron_basic.megatron_basic import MegatronBasicFeature
28-)19+from mindspeed.features_manager.pipeline_parallel import NoopLayersFeature
29-from mindspeed.features_manager.megatron_basic.megatron_basic import (20+from mindspeed.features_manager.pipeline_parallel.ripipe_schedules_feature import RiPipeSchedulesBubbleFeature, \
30- MegatronBasicFeature,21+ RiPipeSchedulesAdvanceFeature
31-)
32 22 
33-from mindspeed.features_manager.tensor_parallel.unaligned_linear_feature import (23+from mindspeed.features_manager.tensor_parallel.unaligned_linear_feature import UnalignedLinearFeature
34- UnalignedLinearFeature,24+from mindspeed.features_manager.pipeline_parallel.unaligned_pineline_feature import UnalignedPipelineFeature
35-)
36-from mindspeed.features_manager.pipeline_parallel.unaligned_pineline_feature import (
37- UnalignedPipelineFeature,
38-)
39 25 
40from mindspeed.features_manager.llava.llava_multimodal import LlavaModel26from mindspeed.features_manager.llava.llava_multimodal import LlavaModel
41-from mindspeed.features_manager.megatron_basic.megatron_basic import (27+from mindspeed.features_manager.megatron_basic.megatron_basic import MegatronBasicFeature
42- MegatronBasicFeature,28+from mindspeed.features_manager.megatron_basic.requirements_basic import RequirementsBasicFeature
43-)29+from mindspeed.features_manager.recompute.activation_function import RecomputeActivationFeature
44-from mindspeed.features_manager.megatron_basic.requirements_basic import (
45- RequirementsBasicFeature,
46-)
47-from mindspeed.features_manager.recompute.activation_function import (
48- RecomputeActivationFeature,
49-)
50from mindspeed.features_manager.recompute.norm_function import RecomputeNormFeature30from mindspeed.features_manager.recompute.norm_function import RecomputeNormFeature
51-from mindspeed.features_manager.tensor_parallel.unaligned_linear_feature import (31+from mindspeed.features_manager.tensor_parallel.unaligned_linear_feature import UnalignedLinearFeature
52- UnalignedLinearFeature,
53-)
54 32 
55from mindspeed.features_manager.tensor_parallel.mc2 import MC2Feature33from mindspeed.features_manager.tensor_parallel.mc2 import MC2Feature
56 34 
57from mindspeed.features_manager.moe.tp_extend_ep import MoETpExtendEpFeature35from mindspeed.features_manager.moe.tp_extend_ep import MoETpExtendEpFeature
58from mindspeed.features_manager.moe.gmm import MoEGmmFeature36from mindspeed.features_manager.moe.gmm import MoEGmmFeature
59from mindspeed.features_manager.moe.shared_expert import MoESharedExpertsFeature37from mindspeed.features_manager.moe.shared_expert import MoESharedExpertsFeature
60-from mindspeed.features_manager.optimizer.virtual_optimizer import (38+from mindspeed.features_manager.optimizer.virtual_optimizer import VirtualOptimizerFeature
61- VirtualOptimizerFeature,39+from mindspeed.features_manager.transformer.flash_attention.alibi_feature import AlibiFeature
62-)40+from mindspeed.features_manager.transformer.flash_attention.fusion_attention_v2_feature import FusionAttentionV2Feature
63-from mindspeed.features_manager.pipeline_parallel.noop_layers import NoopLayersFeature41+from mindspeed.features_manager.transformer.flash_attention.generate_mask_feature import GenerateMaskFeature
64- 42+from mindspeed.features_manager.pipeline_parallel.variable_seq_length import VariableSequenceLengthFeature
65-from mindspeed.features_manager.hccl_buffer.hccl_buffer_adaptive import (
66- HcclBufferFAdaptiveFeature,
67-)
68- 
69-from mindspeed.features_manager.transformer.flash_attention.fusion_attention_v2_feature import (
70- FusionAttentionV2Feature,
71-)
72-from mindspeed.features_manager.transformer.flash_attention.alibi_feature import (
73- AlibiFeature,
74-)
75-from mindspeed.features_manager.transformer.flash_attention.generate_mask_feature import (
76- GenerateMaskFeature,
77-)
78-from mindspeed.features_manager.pipeline_parallel.variable_seq_length import (
79- VariableSequenceLengthFeature,
80-)
81 43 
82from mindspeed.features_manager.dist_train.dist_train_feature import DistTrainFeature44from mindspeed.features_manager.dist_train.dist_train_feature import DistTrainFeature
83 45 
@@ -90,51 +52,118 @@ FEATURES_LIST = [
90 UnalignedLinearFeature(),52 UnalignedLinearFeature(),
91 # llava-multimodal53 # llava-multimodal
92 LlavaModel(),54 LlavaModel(),
93- UnalignedPipelineFeature(),55+ UnalignedPipelineFeature()
94]56]
95 57 
58+ 
96# this list is for reconstruction of mindspeed59# this list is for reconstruction of mindspeed
97-FEATURES_LIST_V2 = (60+def add_megatron_basic_features(features_list: List[MindSpeedFeature]):
98- # Hccl Buffer61+ features_list.extend([
99- HcclBufferFAdaptiveFeature(),62+ RequirementsBasicFeature(),
100- # Recompute63+ MegatronBasicFeature(),
101- RecomputeActivationFeature(),64+ ])
102- RecomputeNormFeature(),
103- # Functional features
104- ProfilerDefaultFeature(),
105- NPUDeterministicFeature(),
106- TflopsCalculateFeature(),
107- ProfileFeature(),
108- # Fusions features
109- GroupedMatmulFeature(),
110- FusedSwigluFeature(),
111- # Megatron basic
112- RequirementsBasicFeature(),
113- MegatronBasicFeature(),
114- # Tensor parallel features
115- UnalignedLinearFeature(),
116- # ripipe
117- RiPipeSchedulesBubbleFeature(),
118- RiPipeSchedulesAdvanceFeature(),
119- # llava-multimodal
120- LlavaModel(),
121 65 
122- # coc_overlap: mc2
123- MC2Feature(),
124 66 
125- # Transformer flash attention features67+def add_fusions_features(features_list: List[MindSpeedFeature]):
126- FusionAttentionV2Feature(),68+ features_list.extend([
127- AlibiFeature(),69+ GroupedMatmulFeature(),
128- GenerateMaskFeature(),70+ FusedSwigluFeature(),
129- # MoeExperts use gemm71+ FusedSoftmaxFeature(),
130- MoEGmmFeature(),72+ ])
131- # MoeTp2EpFeature73+ 
132- MoETpExtendEpFeature(),74+ 
133- # MoeSharedExpertsFeature75+def add_functional_features(features_list: List[MindSpeedFeature]):
134- MoESharedExpertsFeature(),76+ features_list.extend([
135- # pipeline parallel77+ ProfilerDefaultFeature(),
136- NoopLayersFeature(),78+ NPUDeterministicFeature(),
137- VariableSequenceLengthFeature(),79+ TflopsCalculateFeature(),
138- # DistTrainFeature80+ ProfileFeature(),
139- DistTrainFeature()81+ ])
140-)82+ 
83+ 
84+def add_recompute_features(features_list: List[MindSpeedFeature]):
85+ features_list.extend([
86+ RecomputeActivationFeature(),
87+ RecomputeNormFeature(),
88+ ])
89+ 
90+ 
91+def add_tensor_parallel_features(features_list: List[MindSpeedFeature]):
92+ features_list.extend([
93+ UnalignedLinearFeature(),
94+ MC2Feature()
95+ ])
96+ 
97+ 
98+def add_pipeline_parallel_features(features_list: List[MindSpeedFeature]):
99+ features_list.extend([
100+ RiPipeSchedulesBubbleFeature(),
101+ RiPipeSchedulesAdvanceFeature(),
102+ NoopLayersFeature(),
103+ VariableSequenceLengthFeature()
104+ ])
105+ 
106+ 
107+def add_transformer_features(features_list: List[MindSpeedFeature]):
108+ features_list.extend([
109+ FusionAttentionV2Feature(),
110+ AlibiFeature(),
111+ GenerateMaskFeature(),
112+ ])
113+ 
114+ 
115+def add_moe_features(features_list: List[MindSpeedFeature]):
116+ features_list.extend([
117+ MoEGmmFeature(),
118+ MoETpExtendEpFeature(),
119+ MoESharedExpertsFeature()
120+ ])
121+ 
122+ 
123+def add_hccl_buffer_features(features_list: List[MindSpeedFeature]):
124+ features_list.extend([
125+ HcclBufferAdaptiveFeature(),
126+ ])
127+ 
128+ 
129+def add_optimizer_features(features_list: List[MindSpeedFeature]):
130+ features_list.extend([
131+ VirtualOptimizerFeature(),
132+ ])
133+ 
134+ 
135+def add_llava_features(features_list: List[MindSpeedFeature]):
136+ features_list.extend([
137+ LlavaModel()
138+ ])
139+ 
140+ 
141+def add_dist_train_features(features_list: List[MindSpeedFeature]):
142+ features_list.extend([
143+ DistTrainFeature()
144+ ])
145+ 
146+ 
147+ 
148+def create_features_list():
149+ features_list = []
150+ add_megatron_basic_features(features_list)
151+ add_fusions_features(features_list)
152+ add_functional_features(features_list)
153+ add_recompute_features(features_list)
154+ add_tensor_parallel_features(features_list)
155+ add_pipeline_parallel_features(features_list)
156+ add_moe_features(features_list)
157+ add_hccl_buffer_features(features_list)
158+ add_optimizer_features(features_list)
159+ add_llava_features(features_list)
160+ add_dist_train_features(features_list)
161+ return features_list
162+ 
163+ 
164+@AutoExecuteFunction
165+def set_default_features_list():
166+ MindSpeedFeaturesManager.set_features_list(create_features_list())
167+ 
168+ 
169+set_default_features_list()
Amindspeed/features_manager/features_manager.py+63-0
@@ -0,0 +1,63 @@
1+from typing import List
2+ 
3+from mindspeed.features_manager import MindSpeedFeature
4+from mindspeed.patch_utils import MindSpeedPatchesManager
5+ 
6+ 
7+class MindSpeedFeaturesManager:
8+ FEATURES_LIST = []
9+ 
10+ @classmethod
11+ def set_features_list(cls, features_list: List[MindSpeedFeature]):
12+ """Set features list"""
13+ cls.FEATURES_LIST[:] = features_list
14+ 
15+ @classmethod
16+ def apply_features_pre_patches(cls, mindspeed_args):
17+ """Apply pre patches of all features."""
18+ for feature in cls.FEATURES_LIST:
19+ if feature.is_need_apply(mindspeed_args):
20+ feature.pre_register_patches(MindSpeedPatchesManager, mindspeed_args)
21+ MindSpeedPatchesManager.apply_patches()
22+ 
23+ @classmethod
24+ def apply_features_patches(cls, mindspeed_args):
25+ """Apply patches of all features."""
26+ for feature in cls.FEATURES_LIST:
27+ if feature.is_need_apply(mindspeed_args):
28+ feature.register_patches(MindSpeedPatchesManager, mindspeed_args)
29+ MindSpeedPatchesManager.apply_patches()
30+ 
31+ @classmethod
32+ def register_features_args(cls, parser):
33+ """Parse arguments of all features."""
34+ for feature in cls.FEATURES_LIST:
35+ feature.register_args(parser)
36+ 
37+ @classmethod
38+ def pre_validate_features_args(cls, args):
39+ """Pre-validate arguments of all features. Used to bypass megatron arguments validation.
40+ Example:
41+ pre_validate_features_args(args) # old_x = args.x; args.x = new_x
42+ args = validate_args(args, defaults) # bypass args.x validation
43+ post_validate_features_args(args=args) # args.x = old_x
44+ """
45+ for feature in cls.FEATURES_LIST:
46+ feature.pre_validate_args(args)
47+ 
48+ @classmethod
49+ def post_validate_features_args(cls, args):
50+ """Post-validate arguments of all features. Used to bypass megatron arguments validation.
51+ Example:
52+ pre_validate_features_args(args) # old_x = args.x; args.x = new_x
53+ args = validate_args(args, defaults) # bypass args.x validation
54+ post_validate_features_args(args=args) # args.x = old_x
55+ """
56+ for feature in cls.FEATURES_LIST:
57+ feature.post_validate_args(args)
58+ 
59+ @classmethod
60+ def validate_features_args(cls, args):
61+ """Validate arguments of all features."""
62+ for feature in cls.FEATURES_LIST:
63+ feature.validate_args(args)
Amindspeed/features_manager/fusions/fused_softmax.py+18-0
@@ -0,0 +1,18 @@
1+from mindspeed.features_manager.feature import MindSpeedFeature
2+ 
3+ 
4+class FusedSoftmaxFeature(MindSpeedFeature):
5+ def __init__(self):
6+ super().__init__('fused-softmax', optimization_level=0)
7+ 
8+ def register_patches(self, pm, args):
9+ from mindspeed.core.fusions.fused_softmax import is_kernel_available, ScaledUpperTriangMaskedSoftmax, \
10+ ScaledMaskedSoftmax, ScaledSoftmax, forward_fused_softmax
11+ pm.register_patch('megatron.core.fusions.fused_softmax.ScaledUpperTriangMaskedSoftmax',
12+ ScaledUpperTriangMaskedSoftmax)
13+ pm.register_patch('megatron.core.fusions.fused_softmax.ScaledMaskedSoftmax', ScaledMaskedSoftmax)
14+ pm.register_patch('megatron.core.fusions.fused_softmax.ScaledSoftmax', ScaledSoftmax)
15+ pm.register_patch('megatron.core.fusions.fused_softmax.FusedScaleMaskSoftmax.is_kernel_available',
16+ is_kernel_available)
17+ pm.register_patch('megatron.core.fusions.fused_softmax.FusedScaleMaskSoftmax.forward_fused_softmax',
18+ forward_fused_softmax)
Mmindspeed/features_manager/hccl_buffer/hccl_buffer_adaptive.py+1-1
@@ -2,7 +2,7 @@
2from mindspeed.features_manager.feature import MindSpeedFeature2from mindspeed.features_manager.feature import MindSpeedFeature
3 3 
4 4 
5-class HcclBufferFAdaptiveFeature(MindSpeedFeature):5+class HcclBufferAdaptiveFeature(MindSpeedFeature):
6 def __init__(self):6 def __init__(self):
7 super().__init__('hccl-group-buffer-adaptive')7 super().__init__('hccl-group-buffer-adaptive')
8 8 
Mmindspeed/features_manager/megatron_basic/megatron_basic.py+2-1
@@ -23,9 +23,10 @@ class MegatronBasicFeature(MindSpeedFeature):
23 pm.register_patch("megatron.core.transformer.transformer_config.TransformerConfig.__post_init__", transformer_config_post_init_wrapper)23 pm.register_patch("megatron.core.transformer.transformer_config.TransformerConfig.__post_init__", transformer_config_post_init_wrapper)
24 24 
25 # initialization patches25 # initialization patches
26- from mindspeed.core.megatron_basic.megatron_basic import _set_cuda_rng_state, _compile_dependencies26+ from mindspeed.core.megatron_basic.megatron_basic import _set_cuda_rng_state, _compile_dependencies, get_device_wrapper
27 pm.register_patch('megatron.core.tensor_parallel.random._set_cuda_rng_state', _set_cuda_rng_state)27 pm.register_patch('megatron.core.tensor_parallel.random._set_cuda_rng_state', _set_cuda_rng_state)
28 pm.register_patch('megatron.training.initialize._compile_dependencies', _compile_dependencies)28 pm.register_patch('megatron.training.initialize._compile_dependencies', _compile_dependencies)
29+ pm.register_patch('megatron.training.dist_signal_handler.get_device', get_device_wrapper)
29 30 
30 # norm patches31 # norm patches
31 from mindspeed.core.megatron_basic.megatron_basic import PTNorm32 from mindspeed.core.megatron_basic.megatron_basic import PTNorm
Mmindspeed/features_manager/megatron_basic/requirements_basic.py+11-6
@@ -1,4 +1,5 @@
1import sys1import sys
2+from argparse import ArgumentParser
2 3 
3import torch4import torch
4from mindspeed.features_manager.feature import MindSpeedFeature5from mindspeed.features_manager.feature import MindSpeedFeature
@@ -8,6 +9,12 @@ class RequirementsBasicFeature(MindSpeedFeature):
8 def __init__(self):9 def __init__(self):
9 super().__init__('requirements-basic', optimization_level=0)10 super().__init__('requirements-basic', optimization_level=0)
10 11 
12+ def register_args(self, parser: ArgumentParser):
13+ group = parser.add_argument_group(title=self.feature_name)
14+ group.add_argument('--optimizer-selection', type=str, default='fused_adamw',
15+ choices=['fused_adamw', 'fused_torch_adamw'],
16+ help='Select from the former fused AdamW optimizer and Torch fused AdamW optimizer')
17+ 
11 def pre_register_patches(self, patch_manager, args):18 def pre_register_patches(self, patch_manager, args):
12 self.te_adaptation(patch_manager, args)19 self.te_adaptation(patch_manager, args)
13 self.apex_adaptation(patch_manager, args)20 self.apex_adaptation(patch_manager, args)
@@ -44,12 +51,10 @@ class RequirementsBasicFeature(MindSpeedFeature):
44 pm.register_patch('apex.normalization.fused_layer_norm.FusedLayerNormAffineFunction',51 pm.register_patch('apex.normalization.fused_layer_norm.FusedLayerNormAffineFunction',
45 FusedLayerNormAffineFunction, create_dummy=True)52 FusedLayerNormAffineFunction, create_dummy=True)
46 53 
47- if args.optimizer_selection == "fused_torch_adamw":54+ if args.optimizer_selection == 'fused_torch_adamw':
48- pm.register_patch(55+ pm.register_patch('apex.optimizers.FusedAdam', FusedTorchAdamW, create_dummy=True)
49- "apex.optimizers.FusedAdam", FusedTorchAdamW, create_dummy=True56+ elif args.optimizer_selection == 'fused_adamw':
50- )57+ pm.register_patch('apex.optimizers.FusedAdam', AdamW, create_dummy=True)
51- elif args.optimizer_selection == "fused_adamw":
52- pm.register_patch("apex.optimizers.FusedAdam", AdamW, create_dummy=True)
53 pm.register_patch('apex.optimizers.FusedSGD', torch.optim.SGD, create_dummy=True)58 pm.register_patch('apex.optimizers.FusedSGD', torch.optim.SGD, create_dummy=True)
54 59 
55 def torch_adaptation(self, pm, args):60 def torch_adaptation(self, pm, args):
Mmindspeed/functional/npu_deterministic/npu_deterministic.py+3-3
@@ -24,8 +24,8 @@ def extend_seed_all(seed=1234):
24 24 
25def npu_deterministic_wrapper(fn):25def npu_deterministic_wrapper(fn):
26 @wraps(fn)26 @wraps(fn)
27- def wrapper(seed_, data_parallel_random_init=False):27+ def wrapper(seed, *args, **kwargs):
28- fn(seed_, data_parallel_random_init)28+ fn(seed, *args, **kwargs)
29- extend_seed_all(seed_)29+ extend_seed_all(seed)
30 LOG.info("Deterministic computing is applied for npu.")30 LOG.info("Deterministic computing is applied for npu.")
31 return wrapper31 return wrapper
Mmindspeed/megatron_adaptor.py+4-11
@@ -11,13 +11,9 @@ import torch
11from torch.distributed import all_gather_into_tensor, reduce_scatter_tensor11from torch.distributed import all_gather_into_tensor, reduce_scatter_tensor
12from torch_npu.contrib import transfer_to_npu12from torch_npu.contrib import transfer_to_npu
13from mindspeed.features_manager import FEATURES_LIST13from mindspeed.features_manager import FEATURES_LIST
14-from .arguments import process_args14+from mindspeed.deprecate import AutoExecuteFunction, Deprecated, MEGATRON_ADAPTOR_DEPRECATED_TIME
15-from .deprecate import (
16- DisableExecution,
17- Deprecated,
18- MEGATRON_ADAPTOR_DEPRECATED_TIME,
19-)
20 15 
16+from .arguments import process_args
21 17 
22_ARGS = None18_ARGS = None
23LOG = getLogger(__name__)19LOG = getLogger(__name__)
@@ -1079,7 +1075,7 @@ def delete_lock_file(directory, lock):
1079 break1075 break
1080 1076 
1081 1077 
1082-@DisableExecution1078+@AutoExecuteFunction
1083def exe_adaptation():1079def exe_adaptation():
1084 mindspeed_args = get_mindspeed_args()1080 mindspeed_args = get_mindspeed_args()
1085 1081 
@@ -1115,10 +1111,7 @@ def exe_adaptation():
1115 1111 
1116 # New features structure1112 # New features structure
1117 for feature in FEATURES_LIST:1113 for feature in FEATURES_LIST:
1118- if (1114+ if getattr(mindspeed_args, feature.feature_name, None) or feature.default_patches:
1119- getattr(mindspeed_args, feature.feature_name, None)
1120- or feature.default_patches
1121- ):
1122 feature.register_patches(aspm, mindspeed_args)1115 feature.register_patches(aspm, mindspeed_args)
1123 1116 
1124 aspm.apply_patches()1117 aspm.apply_patches()
Mmindspeed/megatron_adaptor_v2.py+11-22
@@ -1,7 +1,5 @@
1"""Adaptor for all megatron functions by feature granularity."""1"""Adaptor for all megatron functions by feature granularity."""
2 2 
3-from argparse import Namespace
4-from typing import Type
5import os3import os
6import sys4import sys
7import shutil5import shutil
@@ -12,23 +10,19 @@ from pathlib import Path
12from torch.utils.cpp_extension import _get_build_directory10from torch.utils.cpp_extension import _get_build_directory
13from torch_npu.contrib import transfer_to_npu11from torch_npu.contrib import transfer_to_npu
14 12 
15-from .deprecate import DisableExecution
16 13 
17-# just in case execution of exe_adaptation function
18-# in megatron_adaptor when import it.
19- 
20-from mindspeed.features_manager import FEATURES_LIST_V2
21-DisableExecution.DISABLE = True
22-from mindspeed.patch_utils import MindSpeedPatchesManager
23-from mindspeed.arguments_v2 import get_mindspeed_args
24- 
25-DisableExecution.DISABLE = False
26from mindspeed.log_config import set_log_config14from mindspeed.log_config import set_log_config
15+from mindspeed.deprecate import AutoExecuteFunction, NoExecuteFunction
16+from mindspeed.features_manager.features_manager import MindSpeedFeaturesManager
17+ 
18+with NoExecuteFunction():
19+ from mindspeed.megatron_adaptor import get_mindspeed_args
27 20 
28LOG = getLogger(__name__)21LOG = getLogger(__name__)
29_IS_FEATURES_PATCHED = False22_IS_FEATURES_PATCHED = False
30 23 
31 24 
25+@AutoExecuteFunction
32def patch_features():26def patch_features():
33 """Patch all mindspeed related features."""27 """Patch all mindspeed related features."""
34 global _IS_FEATURES_PATCHED28 global _IS_FEATURES_PATCHED
@@ -44,19 +38,14 @@ def patch_features():
44 delete_lock_file()38 delete_lock_file()
45 39 
46 # apply patches before import megatron40 # apply patches before import megatron
47- for feature in FEATURES_LIST_V2:41+ MindSpeedFeaturesManager.apply_features_pre_patches(mindspeed_args)
48- if feature.is_need_apply(mindspeed_args):
49- feature.pre_register_patches(MindSpeedPatchesManager, mindspeed_args)
50- MindSpeedPatchesManager.apply_patches()
51 42 
52 # apply megatron patches43 # apply megatron patches
53- for feature in FEATURES_LIST_V2:44+ MindSpeedFeaturesManager.apply_features_patches(mindspeed_args)
54- if feature.is_need_apply(mindspeed_args):
55- feature.register_patches(MindSpeedPatchesManager, mindspeed_args)
56- MindSpeedPatchesManager.apply_patches()
57 45 
58- # accelerate package will check TE on sys.modules，so we need remove this patch46+ # accelerate package will check TE on sys.modules, so we need remove this patch
59- del sys.modules["transformer_engine"]47+ if 'transformer_engine' in sys.modules:
48+ del sys.modules["transformer_engine"]
60 49 
61 50 
62def delete_lock_file():51def delete_lock_file():
Mmindspeed/patch_utils.py+43-2
@@ -1,6 +1,7 @@
1import importlib1import importlib
2import sys2import sys
3import types3import types
4+from typing import List, Dict, Union
4 5 
5 6 
6def get_func_name(func):7def get_func_name(func):
@@ -44,13 +45,29 @@ class Patch:
44 45 
45 def set_patch_func(self, new_func, force_patch=False):46 def set_patch_func(self, new_func, force_patch=False):
46 if hasattr(new_func, '__name__') and new_func.__name__.endswith(('wrapper', 'decorator')):47 if hasattr(new_func, '__name__') and new_func.__name__.endswith(('wrapper', 'decorator')):
47- self.wrappers.append(new_func)48+ if new_func not in self.wrappers:
49+ self.wrappers.append(new_func)
48 else:50 else:
49 if self.patch_func and not force_patch:51 if self.patch_func and not force_patch:
50 raise RuntimeError('the patch of {} exist !'.format(self.orig_func_name))52 raise RuntimeError('the patch of {} exist !'.format(self.orig_func_name))
51 self.patch_func = new_func53 self.patch_func = new_func
52 self.is_applied = False54 self.is_applied = False
53 55 
56+ def remove_wrappers(self, wrapper_names: Union[str, List[str]] = None):
57+ if wrapper_names is None:
58+ self.wrappers.clear()
59+ return
60+ 
61+ if isinstance(wrapper_names, str):
62+ wrapper_names = [wrapper_names]
63+ for name in wrapper_names:
64+ i = 0
65+ while i < len(self.wrappers):
66+ if self.wrappers[i].__name__ == name:
67+ self.wrappers.pop(i)
68+ else:
69+ i += 1
70+ 
54 def apply_patch(self):71 def apply_patch(self):
55 if self.is_applied:72 if self.is_applied:
56 return73 return
@@ -105,16 +122,40 @@ class Patch:
105 122 
106 123 
107class MindSpeedPatchesManager:124class MindSpeedPatchesManager:
108- patches_info = {}125+ patches_info: Dict[str, Patch] = {}
109 126 
110 @staticmethod127 @staticmethod
111 def register_patch(orig_func_name, new_func=None, force_patch=False, create_dummy=False):128 def register_patch(orig_func_name, new_func=None, force_patch=False, create_dummy=False):
129+ """Patch registration method. When this method is executed, the patch does not take effect in real time.
130+ It takes effect only after the apply_patches method is invoked. Other details are as follows:
131+ 
132+ 1. If `orig_func_name` does not exist and create_dummy is set to True, a dummy function is created to ensure
133+ that the import is normal.
134+ 2. If `orig_func_name` is not None, `orig_func_name` is replaced with `new_func`.
135+ 3. If the `new_func` function name ends with `wrapper` or `decorator`, then `new_func` is decorated on
136+ `orig_func_name` as a decorator, and the decorator can be superimposed repeatedly.
137+ 4. When force_patch=False, a function cannot be replaced repeatedly (but can be decorated repeatedly),
138+ otherwise the replacement is overwritten.
139+ """
112 if orig_func_name not in MindSpeedPatchesManager.patches_info:140 if orig_func_name not in MindSpeedPatchesManager.patches_info:
113 MindSpeedPatchesManager.patches_info[orig_func_name] = Patch(orig_func_name, new_func, create_dummy)141 MindSpeedPatchesManager.patches_info[orig_func_name] = Patch(orig_func_name, new_func, create_dummy)
114 else:142 else:
115 MindSpeedPatchesManager.patches_info.get(orig_func_name).set_patch_func(new_func, force_patch)143 MindSpeedPatchesManager.patches_info.get(orig_func_name).set_patch_func(new_func, force_patch)
116 144 
145+ @staticmethod
146+ def remove_wrappers(orig_func_name, wrappers_name, remove_check=True):
147+ """Remove wrapper registered in orig_func_name."""
148+ if orig_func_name not in MindSpeedPatchesManager.patches_info:
149+ raise ValueError('The function <{}> not exist.'.format(orig_func_name))
150+ 
151+ patch = MindSpeedPatchesManager.patches_info.get(orig_func_name)
152+ wrappers_len = len(patch.wrappers)
153+ patch.remove_wrappers(wrappers_name)
154+ if remove_check and wrappers_len == len(patch.wrappers):
155+ raise RuntimeError('Remove wrappers has not remove anything.')
156+ 
117 @staticmethod157 @staticmethod
118 def apply_patches():158 def apply_patches():
159+ """Apply all patches registered in MindSpeedPatchesManager."""
119 for patch in MindSpeedPatchesManager.patches_info.values():160 for patch in MindSpeedPatchesManager.patches_info.values():
120 patch.apply_patch()161 patch.apply_patch()
Mtests_extend/unit_tests/mindspeed/test_register_patches.py+18-0
@@ -113,3 +113,21 @@ class TestRegisterPatchesResetEnv(DistributedTest):
113 from unit_tests.mindspeed.test_register_patches import function1113 from unit_tests.mindspeed.test_register_patches import function1
114 114 
115 assert function1() == 'this is function3'115 assert function1() == 'this is function3'
116+ 
117+ def test_double_wrapper(self):
118+ aspm.register_patch('unit_tests.mindspeed.test_register_patches.function1', function_wrapper)
119+ aspm.register_patch('unit_tests.mindspeed.test_register_patches.function1', function_wrapper)
120+ aspm.apply_patches()
121+ 
122+ from unit_tests.mindspeed.test_register_patches import function1
123+ 
124+ assert function1() == 'this is function1 wrapper'
125+ 
126+ def test_remove_wrapper(self):
127+ aspm.register_patch('unit_tests.mindspeed.test_register_patches.function1', function_wrapper)
128+ aspm.remove_wrappers('unit_tests.mindspeed.test_register_patches.function1', 'function_wrapper')
129+ aspm.apply_patches()
130+ 
131+ from unit_tests.mindspeed.test_register_patches import function1
132+ 
133+ assert function1() == 'this is function1'
Mtests_extend_v2/unit_tests/features/flash_attention/test_alibi.py+4-2
@@ -2,11 +2,11 @@ import pytest
2import torch2import torch
3import torch_npu3import torch_npu
4 4 
5+from mindspeed import megatron_adaptor_v2
5from megatron.training.global_vars import set_args6from megatron.training.global_vars import set_args
6from megatron.training.arguments import parse_args7from megatron.training.arguments import parse_args
7from megatron.core.transformer.transformer_config import TransformerConfig8from megatron.core.transformer.transformer_config import TransformerConfig
8 9 
9-from mindspeed import megatron_adaptor_v2
10from mindspeed.core.transformer.flash_attention.alibi.adaptor import MindSpeedDotProductAttention10from mindspeed.core.transformer.flash_attention.alibi.adaptor import MindSpeedDotProductAttention
11 11 
12DEVICE_NAME = torch_npu.npu.get_device_name(0)[:10]12DEVICE_NAME = torch_npu.npu.get_device_name(0)[:10]
@@ -33,6 +33,8 @@ def run_fusion_attn_with_pse_alibi(bs, seq_len, dtype):
33 config.sparse_mode = 233 config.sparse_mode = 2
34 config.seq_length = seq_len34 config.seq_length = seq_len
35 config.alibi_diagonal_opposite = False35 config.alibi_diagonal_opposite = False
36+ config.pre_tockens = 65536
37+ config.next_tockens = 0
36 38 
37 attn = MindSpeedDotProductAttention(39 attn = MindSpeedDotProductAttention(
38 config=config, 40 config=config,
@@ -55,7 +57,7 @@ def run_fusion_attn_with_pse_alibi(bs, seq_len, dtype):
55 assert isinstance(out, torch.Tensor)57 assert isinstance(out, torch.Tensor)
56 58 
57 59 
58-class TestAlibi():60+class TestAlibi:
59 61 
60 @pytest.mark.skipif(DEVICE_NAME != 'Ascend910B', reason='device type is not supported, skip this UT!')62 @pytest.mark.skipif(DEVICE_NAME != 'Ascend910B', reason='device type is not supported, skip this UT!')
61 def test_alibi(self, mocker):63 def test_alibi(self, mocker):