已合并
update:2.0.0升级mindspeed版本至2.0.0_core_r0.8.0 #2509
wx_4e191bb7af创建于 2025年4月3日
update:2.0.0升级mindspeed版本至2.0.0_core_r0.8.0 #2509
已合并
从refs/pull/2509/head合入到2.0.0
共 5 个文件变更+6-59
| @@ -49,7 +49,7 @@ source /usr/local/Ascend/nnal/atb/set_env.sh | |||
| 49 | # 安装MindSpeed加速库 | 49 | # 安装MindSpeed加速库 |
| 50 | git clone https://gitee.com/ascend/MindSpeed.git | 50 | git clone https://gitee.com/ascend/MindSpeed.git |
| 51 | cd MindSpeed | 51 | cd MindSpeed |
| 52 | -git checkout 0dfa0035ec54d9a74b2f6ee2867367df897299df # checkout commit from MindSpeed core_r0.8.0 in 2025.02.26 | 52 | +git checkout 2.0.0_core_r0.8.0 # checkout commit from MindSpeed 2.0.0_core_r0.8.0 |
| 53 | pip install -r requirements.txt | 53 | pip install -r requirements.txt |
| 54 | pip3 install -e . | 54 | pip3 install -e . |
| 55 | cd .. | 55 | cd .. |
| @@ -3,7 +3,6 @@ | |||
| 3 | 3 | ||
| 4 | import math | 4 | import math |
| 5 | from functools import wraps | 5 | from functools import wraps |
| 6 | -from typing import Any | ||
| 7 | 6 | ||
| 8 | import torch | 7 | import torch |
| 9 | import torch_npu | 8 | import torch_npu |
| @@ -19,7 +18,6 @@ from megatron.core.utils import divide | |||
| 19 | from megatron.core.packed_seq_params import PackedSeqParams | 18 | from megatron.core.packed_seq_params import PackedSeqParams |
| 20 | from mindspeed.core.context_parallel.ring_context_parallel import ringattn_context_parallel | 19 | from mindspeed.core.context_parallel.ring_context_parallel import ringattn_context_parallel |
| 21 | from mindspeed.core.context_parallel.ulysses_context_parallel import ulyssesattn_context_parallel | 20 | from mindspeed.core.context_parallel.ulysses_context_parallel import ulyssesattn_context_parallel |
| 22 | -from mindspeed.core.context_parallel.ulysses_context_parallel import _SeqAllToAll | ||
| 23 | from mindspeed.core.parallel_state import (get_context_parallel_group_for_hybrid_ring, | 21 | from mindspeed.core.parallel_state import (get_context_parallel_group_for_hybrid_ring, |
| 24 | get_context_parallel_for_hybrid_ring_world_size, | 22 | get_context_parallel_for_hybrid_ring_world_size, |
| 25 | get_context_parallel_for_hybrid_ring_rank, | 23 | get_context_parallel_for_hybrid_ring_rank, |
| @@ -246,58 +244,6 @@ def ulysses_context_parallel_forward_wrapper(fn): | |||
| 246 | return wrapper | 244 | return wrapper |
| 247 | 245 | ||
| 248 | 246 | ||
| 249 | -# bigfix ulysses, needs to be removed when mindspeed is upgraded to 2025.03.04. | ||
| 250 | -def ulysses_context_parallel_forward(self, query: Tensor, key: Tensor, value: Tensor, *args: Any, **kwargs: Any) -> Tensor: | ||
| 251 | - """ forward | ||
| 252 | - | ||
| 253 | - Arguments: | ||
| 254 | - query (Tensor): query input to the layer | ||
| 255 | - key (Tensor): key input to the layer | ||
| 256 | - value (Tensor): value input to the layer | ||
| 257 | - args: other args | ||
| 258 | - | ||
| 259 | - Returns: | ||
| 260 | - * output (Tensor): context output | ||
| 261 | - """ | ||
| 262 | - heads_per_gqa_group = self.local_attn.num_attention_heads_per_partition // self.local_attn.num_query_groups_per_partition | ||
| 263 | - global_args = get_args() | ||
| 264 | - should_kv_repeat_before_uly = global_args.use_flash_attn and global_args.kv_head_repeat_before_uly_alltoall | ||
| 265 | - | ||
| 266 | - if heads_per_gqa_group > 1 and should_kv_repeat_before_uly: | ||
| 267 | - key = key.repeat_interleave(heads_per_gqa_group, dim=2) | ||
| 268 | - value = value.repeat_interleave(heads_per_gqa_group, dim=2) | ||
| 269 | - | ||
| 270 | - | ||
| 271 | - use_custom_ulysses_backward = ( | ||
| 272 | - global_args.context_parallel_size > 1 and | ||
| 273 | - global_args.context_parallel_algo == "ulysses_cp_algo" and | ||
| 274 | - not global_args.use_legacy_models and | ||
| 275 | - global_args.context_parallel_kv_cache_policy | ||
| 276 | - ) | ||
| 277 | - if use_custom_ulysses_backward: | ||
| 278 | - output = self.local_attn(query, key, value, *args, **kwargs) | ||
| 279 | - else: | ||
| 280 | - spg = self.local_attn.ulysses_comm_para.get('spg') | ||
| 281 | - scatter_idx = self.local_attn.ulysses_comm_para.get('scatter_idx') | ||
| 282 | - gather_idx = self.local_attn.ulysses_comm_para.get('gather_idx') | ||
| 283 | - seq_world_size = torch.distributed.get_world_size(spg) | ||
| 284 | - if seq_world_size > key.shape[scatter_idx] and query.shape[scatter_idx] % key.shape[scatter_idx] == 0: | ||
| 285 | - key = key.repeat_interleave(query.shape[scatter_idx] // key.shape[scatter_idx], dim=scatter_idx) | ||
| 286 | - value = value.repeat_interleave(query.shape[scatter_idx] // value.shape[scatter_idx], dim=scatter_idx) | ||
| 287 | - | ||
| 288 | - # in shape : e.g., [s/p:h:] | ||
| 289 | - query_layer = _SeqAllToAll.apply(spg, query, scatter_idx, gather_idx) | ||
| 290 | - key_layer = _SeqAllToAll.apply(spg, key, scatter_idx, gather_idx) | ||
| 291 | - value_layer = _SeqAllToAll.apply(spg, value, scatter_idx, gather_idx) | ||
| 292 | - | ||
| 293 | - # out shape : e.g., [s:h/p:] | ||
| 294 | - context_layer = self.local_attn(query_layer, key_layer, value_layer, *args, **kwargs) | ||
| 295 | - | ||
| 296 | - output = _SeqAllToAll.apply(spg, context_layer, gather_idx, scatter_idx) | ||
| 297 | - | ||
| 298 | - return output | ||
| 299 | - | ||
| 300 | - | ||
| 301 | def dot_product_attention_forward_wrapper(fn): | 247 | def dot_product_attention_forward_wrapper(fn): |
| 302 | 248 | ||
| 303 | def wrapper(self, query, key, value, attention_mask, attn_mask_type, packed_seq_params): | 249 | def wrapper(self, query, key, value, attention_mask, attn_mask_type, packed_seq_params): |
| @@ -253,7 +253,7 @@ class CoreAdaptation(MegatronAdaptationABC): | |||
| 253 | from ..core import rotary_embedding_forward, apply_rotary_pos_emb_bshd | 253 | from ..core import rotary_embedding_forward, apply_rotary_pos_emb_bshd |
| 254 | from ..core.models.gpt.gpt_layer_specs import get_gpt_layer_local_spec_wrapper | 254 | from ..core.models.gpt.gpt_layer_specs import get_gpt_layer_local_spec_wrapper |
| 255 | from ..core.transformer.dot_product_attention import dot_product_attention_init, \ | 255 | from ..core.transformer.dot_product_attention import dot_product_attention_init, \ |
| 256 | - dot_product_attention_forward_wrapper, ulysses_context_parallel_forward | 256 | + dot_product_attention_forward_wrapper, ulysses_context_parallel_forward_wrapper |
| 257 | from ..core.models.gpt.gpt_model import gpt_model_init_wrapper, shared_embedding_weight | 257 | from ..core.models.gpt.gpt_model import gpt_model_init_wrapper, shared_embedding_weight |
| 258 | from ..core import rotary_embedding_init_wrapper, gpt_model_forward | 258 | from ..core import rotary_embedding_init_wrapper, gpt_model_forward |
| 259 | 259 | ||
| @@ -309,7 +309,7 @@ class CoreAdaptation(MegatronAdaptationABC): | |||
| 309 | # For GQA in ulysses and hybrid | 309 | # For GQA in ulysses and hybrid |
| 310 | MegatronAdaptation.register( | 310 | MegatronAdaptation.register( |
| 311 | 'mindspeed.core.context_parallel.ulysses_context_parallel.UlyssesContextAttention.forward', | 311 | 'mindspeed.core.context_parallel.ulysses_context_parallel.UlyssesContextAttention.forward', |
| 312 | - ulysses_context_parallel_forward) | 312 | + ulysses_context_parallel_forward_wrapper) |
| 313 | 313 | ||
| 314 | # Layer Definition | 314 | # Layer Definition |
| 315 | # For NPU, we use local-mcore-structrue in te layer. | 315 | # For NPU, we use local-mcore-structrue in te layer. |
| @@ -1347,6 +1347,7 @@ def _add_dummy_args(args): | |||
| 1347 | args.megatron_cp_in_bnsd = False | 1347 | args.megatron_cp_in_bnsd = False |
| 1348 | args.use_fusion_attn_v2 = False | 1348 | args.use_fusion_attn_v2 = False |
| 1349 | args.npu_deterministic = False | 1349 | args.npu_deterministic = False |
| 1350 | + args.schedules_method = None | ||
| 1350 | 1351 | ||
| 1351 | 1352 | ||
| 1352 | def _validate_noop_layer(args): | 1353 | def _validate_noop_layer(args): |
| @@ -9,8 +9,8 @@ source /usr/local/Ascend/nnal/atb/set_env.sh | |||
| 9 | git clone -b core_r0.8.0 https://gitee.com/ascend/MindSpeed.git | 9 | git clone -b core_r0.8.0 https://gitee.com/ascend/MindSpeed.git |
| 10 | cd MindSpeed | 10 | cd MindSpeed |
| 11 | 11 | ||
| 12 | -# checkout commit from MindSpeed core_r0.8.0 in 2025.02.26 | 12 | +# checkout commit from MindSpeed 2.0.0_core_r0.8.0 |
| 13 | -git checkout 0dfa0035ec54d9a74b2f6ee2867367df897299df | 13 | +git checkout 2.0.0_core_r0.8.0 |
| 14 | pip install -r requirements.txt | 14 | pip install -r requirements.txt |
| 15 | pip3 install -e . | 15 | pip3 install -e . |
| 16 | cd .. | 16 | cd .. |