已合并
feat:支持moe-allgather-overlap-comm特性 #2282
wx_4e191bb7af创建于 2025年2月23日
feat:支持moe-allgather-overlap-comm特性 #2282
已合并
从refs/pull/2282/head合入到master
共 6 个文件变更+49-28
| @@ -14,6 +14,7 @@ from megatron.core.transformer.moe.experts import GroupedMLP, SequentialMLP | |||
| 14 | from megatron.core.transformer.moe.moe_utils import save_to_aux_losses_tracker | 14 | from megatron.core.transformer.moe.moe_utils import save_to_aux_losses_tracker |
| 15 | from megatron.training import get_args | 15 | from megatron.training import get_args |
| 16 | from mindspeed.core.transformer.moe.moe_layer_overlap_all2all import MoELayerOverlapAll2All | 16 | from mindspeed.core.transformer.moe.moe_layer_overlap_all2all import MoELayerOverlapAll2All |
| 17 | +from mindspeed.core.transformer.moe.moe_layer_overlap_allgather import MoELayerOverlapAllGather | ||
| 17 | 18 | ||
| 18 | 19 | ||
| 19 | def moe_layer_init_wrapper(init_func): | 20 | def moe_layer_init_wrapper(init_func): |
| @@ -61,6 +62,8 @@ def moe_layer_forward(self, hidden_states: torch.Tensor): | |||
| 61 | global_args = get_args() | 62 | global_args = get_args() |
| 62 | if global_args.moe_token_dispatcher_type == 'alltoall' and global_args.moe_alltoall_overlap_comm: | 63 | if global_args.moe_token_dispatcher_type == 'alltoall' and global_args.moe_alltoall_overlap_comm: |
| 63 | return MoELayerOverlapAll2All.apply(hidden_states, self) | 64 | return MoELayerOverlapAll2All.apply(hidden_states, self) |
| 65 | + if global_args.moe_token_dispatcher_type == 'allgather' and global_args.moe_allgather_overlap_comm: | ||
| 66 | + return MoELayerOverlapAllGather.apply(hidden_states, self) | ||
| 64 | 67 | ||
| 65 | # process MoE | 68 | # process MoE |
| 66 | scores, indices = self.router(hidden_states) | 69 | scores, indices = self.router(hidden_states) |
| @@ -411,14 +411,25 @@ class CoreAdaptation(MegatronAdaptationABC): | |||
| 411 | if args.moe_permutation_async_comm: | 411 | if args.moe_permutation_async_comm: |
| 412 | if args.moe_token_dispatcher_type == 'allgather': | 412 | if args.moe_token_dispatcher_type == 'allgather': |
| 413 | from mindspeed.core.transformer.moe.router import aux_loss_load_balancing | 413 | from mindspeed.core.transformer.moe.router import aux_loss_load_balancing |
| 414 | - MegatronAdaptation.register( | ||
| 415 | - 'megatron.core.transformer.moe.token_dispatcher.MoEAllGatherTokenDispatcher.token_permutation', | ||
| 416 | - allgather_token_permutation) | ||
| 417 | - MegatronAdaptation.register( | ||
| 418 | - 'megatron.core.transformer.moe.token_dispatcher.MoEAllGatherTokenDispatcher.token_unpermutation', | ||
| 419 | - allgather_token_unpermutation) | ||
| 420 | MegatronAdaptation.register('megatron.core.transformer.moe.router.TopKRouter.aux_loss_load_balancing', | 414 | MegatronAdaptation.register('megatron.core.transformer.moe.router.TopKRouter.aux_loss_load_balancing', |
| 421 | aux_loss_load_balancing) | 415 | aux_loss_load_balancing) |
| 416 | + if args.moe_allgather_overlap_comm: | ||
| 417 | + from mindspeed.core.transformer.moe.token_dispatcher import (allgather_token_permutation_new, | ||
| 418 | + allgather_token_unpermutation_new) | ||
| 419 | + MegatronAdaptation.register('megatron.core.transformer.moe.experts.GroupedMLP.forward', group_mlp_forward) | ||
| 420 | + MegatronAdaptation.register( | ||
| 421 | + 'megatron.core.transformer.moe.token_dispatcher.MoEAllGatherTokenDispatcher.token_permutation', | ||
| 422 | + allgather_token_permutation_new) | ||
| 423 | + MegatronAdaptation.register( | ||
| 424 | + 'megatron.core.transformer.moe.token_dispatcher.MoEAllGatherTokenDispatcher.token_unpermutation', | ||
| 425 | + allgather_token_unpermutation_new) | ||
| 426 | + else: | ||
| 427 | + MegatronAdaptation.register( | ||
| 428 | + 'megatron.core.transformer.moe.token_dispatcher.MoEAllGatherTokenDispatcher.token_permutation', | ||
| 429 | + allgather_token_permutation) | ||
| 430 | + MegatronAdaptation.register( | ||
| 431 | + 'megatron.core.transformer.moe.token_dispatcher.MoEAllGatherTokenDispatcher.token_unpermutation', | ||
| 432 | + allgather_token_unpermutation) | ||
| 422 | elif args.moe_token_dispatcher_type == 'alltoall': | 433 | elif args.moe_token_dispatcher_type == 'alltoall': |
| 423 | from mindspeed.core.transformer.moe.token_dispatcher import preprocess, alltoall_token_permutation | 434 | from mindspeed.core.transformer.moe.token_dispatcher import preprocess, alltoall_token_permutation |
| 424 | from mindspeed.core.transformer.moe.moe_utils import permute, unpermute | 435 | from mindspeed.core.transformer.moe.moe_utils import permute, unpermute |
| @@ -484,7 +495,7 @@ class CoreAdaptation(MegatronAdaptationABC): | |||
| 484 | MegatronAdaptation.register('megatron.core.transformer.moe.moe_utils.permute', permute_wrapper) | 495 | MegatronAdaptation.register('megatron.core.transformer.moe.moe_utils.permute', permute_wrapper) |
| 485 | MegatronAdaptation.register('megatron.core.transformer.moe.moe_utils.unpermute', unpermute_wrapper) | 496 | MegatronAdaptation.register('megatron.core.transformer.moe.moe_utils.unpermute', unpermute_wrapper) |
| 486 | 497 | ||
| 487 | - if not args.moe_alltoall_overlap_comm: | 498 | + if not args.moe_alltoall_overlap_comm and not args.moe_allgather_overlap_comm: |
| 488 | MegatronAdaptation.register('megatron.core.transformer.moe.experts.GroupedMLP.forward', | 499 | MegatronAdaptation.register('megatron.core.transformer.moe.experts.GroupedMLP.forward', |
| 489 | groupedmlp_forward) | 500 | groupedmlp_forward) |
| 490 | 501 | ||
| @@ -233,7 +233,7 @@ def _add_coc_args(parser): | |||
| 233 | group.add_argument('--disable-gloo-group', action='store_true', | 233 | group.add_argument('--disable-gloo-group', action='store_true', |
| 234 | help='Replace the communication method of the DP group in the distributed optimizer from gloo to hccl.') | 234 | help='Replace the communication method of the DP group in the distributed optimizer from gloo to hccl.') |
| 235 | group.add_argument('--hccl-slice-size', type=int, default=10 * 1024 * 1024, | 235 | group.add_argument('--hccl-slice-size', type=int, default=10 * 1024 * 1024, |
| 236 | - help='data slice size on each dp rank in distributed optimizer') | 236 | + help='data slice size on each dp rank in distributed optimizer') |
| 237 | return parser | 237 | return parser |
| 238 | 238 | ||
| 239 | 239 | ||
| @@ -423,6 +423,8 @@ def _add_moe_args(parser): | |||
| 423 | group.add_argument('--moe-zero-memory-num-layers', type=int, default=None, | 423 | group.add_argument('--moe-zero-memory-num-layers', type=int, default=None, |
| 424 | help='the number of layers using moe-zero-memory level1' | 424 | help='the number of layers using moe-zero-memory level1' |
| 425 | 'in each pp stage.') | 425 | 'in each pp stage.') |
| 426 | + group.add_argument('--moe-allgather-overlap-comm', action='store_true', default=False, | ||
| 427 | + help='moe_allgather_overlap_comm') | ||
| 426 | return parser | 428 | return parser |
| 427 | 429 | ||
| 428 | 430 | ||
| @@ -1099,7 +1101,8 @@ def _validate_moe_args(args): | |||
| 1099 | raise AssertionError('shared expert gate output dimension can only be configured with 1 or hidden_size') | 1101 | raise AssertionError('shared expert gate output dimension can only be configured with 1 or hidden_size') |
| 1100 | if hasattr(args, 'use_fused_moe_token_permute_and_unpermute') and args.use_fused_moe_token_permute_and_unpermute: | 1102 | if hasattr(args, 'use_fused_moe_token_permute_and_unpermute') and args.use_fused_moe_token_permute_and_unpermute: |
| 1101 | raise AssertionError('moe_expert_capacity_factor mode does not support use_fused_moe_token_permute_and_unpermute') | 1103 | raise AssertionError('moe_expert_capacity_factor mode does not support use_fused_moe_token_permute_and_unpermute') |
| 1102 | - if args.moe_alltoall_overlap_comm and (not args.moe_permutation_async_comm or not args.moe_grouped_gemm): | 1104 | + if args.moe_alltoall_overlap_comm or args.moe_allgather_overlap_comm: |
| 1105 | + if not args.moe_permutation_async_comm or not args.moe_grouped_gemm: | ||
| 1103 | raise AssertionError( | 1106 | raise AssertionError( |
| 1104 | '`--moe-alltoall-overlap-comm` or `--moe-allgather-overlap-comm` only support with `--moe-permutation-async-comm` and `--moe-grouped-gemm`.') | 1107 | '`--moe-alltoall-overlap-comm` or `--moe-allgather-overlap-comm` only support with `--moe-permutation-async-comm` and `--moe-grouped-gemm`.') |
| 1105 | if args.moe_alltoall_overlap_comm and not args.moe_token_dispatcher_type == 'alltoall': | 1108 | if args.moe_alltoall_overlap_comm and not args.moe_token_dispatcher_type == 'alltoall': |
| @@ -1117,6 +1120,10 @@ def _validate_moe_args(args): | |||
| 1117 | raise AssertionError('`--moe-zero-memory` only supports `--moe-alltoall-overlap-comm` for now.') | 1120 | raise AssertionError('`--moe-zero-memory` only supports `--moe-alltoall-overlap-comm` for now.') |
| 1118 | if args.moe_zero_memory != "disable" and args.recompute_method is not None: | 1121 | if args.moe_zero_memory != "disable" and args.recompute_method is not None: |
| 1119 | raise AssertionError('`--moe-zero-memory` does not support full recomputation for now.') | 1122 | raise AssertionError('`--moe-zero-memory` does not support full recomputation for now.') |
| 1123 | + if args.moe_allgather_overlap_comm and not args.moe_token_dispatcher_type == 'allgather': | ||
| 1124 | + raise AssertionError('`--moe-allgather-overlap-comm` only support with `--moe-token-dispatcher-type allgather`.') | ||
| 1125 | + if args.moe_allgather_overlap_comm and not args.tensor_model_parallel_size > 1 and not args.expert_model_parallel_size > 1: | ||
| 1126 | + raise AssertionError('`--moe_allgather_overlap_comm` requires enabling tp or ep.') | ||
| 1120 | if args.shared_expert_gate and args.gradient_accumulation_fusion: | 1127 | if args.shared_expert_gate and args.gradient_accumulation_fusion: |
| 1121 | raise AssertionError('args.shared_expert_gate does not support gradient_accumulation_fusion.') | 1128 | raise AssertionError('args.shared_expert_gate does not support gradient_accumulation_fusion.') |
| 1122 | 1129 | ||
| @@ -1319,7 +1326,6 @@ def _add_dummy_args(args): | |||
| 1319 | args.adaptive_recompute_profiling_step = 10 | 1326 | args.adaptive_recompute_profiling_step = 10 |
| 1320 | args.recompute_in_bubble = False | 1327 | args.recompute_in_bubble = False |
| 1321 | args.use_nanopipe = False | 1328 | args.use_nanopipe = False |
| 1322 | - args.moe_allgather_overlap_comm = False | ||
| 1323 | args.moe_without_activation = False | 1329 | args.moe_without_activation = False |
| 1324 | args.disable_gloo_group = None | 1330 | args.disable_gloo_group = None |
| 1325 | args.ampipe_degree = 0 | 1331 | args.ampipe_degree = 0 |
| @@ -96,7 +96,7 @@ | |||
| 96 | </tr> | 96 | </tr> |
| 97 | <tr> | 97 | <tr> |
| 98 | <td>Mcore</td> | 98 | <td>Mcore</td> |
| 99 | - <td>mla_attention，moe_grouped_gemm，EP，allgather_dispatcher，use_fused_rotary_pos_emb_new，recompute_norm</td> | 99 | + <td>mla_attention，moe_grouped_gemm，EP，allgather_dispatcher，moe_allgather_overlap_comm，use_fused_rotary_pos_emb_new，recompute_norm</td> |
| 100 | <td><a href="st/shell_scripts/deepseek_v2_mcore_tp1_pp1_ep8.sh">deepseek_v2_mcore_tp1_pp1_ep8.sh</a></td> | 100 | <td><a href="st/shell_scripts/deepseek_v2_mcore_tp1_pp1_ep8.sh">deepseek_v2_mcore_tp1_pp1_ep8.sh</a></td> |
| 101 | <td>Y</td> | 101 | <td>Y</td> |
| 102 | <td>Y</td> | 102 | <td>Y</td> |
| @@ -17,27 +17,27 @@ | |||
| 17 | 12.80007 | 17 | 12.80007 |
| 18 | ], | 18 | ], |
| 19 | "throughput": [ | 19 | "throughput": [ |
| 20 | - 15.1, | 20 | + 12, |
| 21 | - 237.9, | 21 | + 256, |
| 22 | - 238.6, | 22 | + 256, |
| 23 | - 239.0, | 23 | + 257, |
| 24 | - 240.4, | 24 | + 257, |
| 25 | - 242.4, | 25 | + 258, |
| 26 | - 242.4, | 26 | + 258, |
| 27 | - 242.1, | 27 | + 258, |
| 28 | - 240.5, | 28 | + 259, |
| 29 | - 240.4, | 29 | + 258, |
| 30 | - 240.6, | 30 | + 259, |
| 31 | - 240.4, | 31 | + 259, |
| 32 | - 241.4, | 32 | + 259, |
| 33 | - 241.2, | 33 | + 259, |
| 34 | - 242.0 | 34 | + 259 |
| 35 | ], | 35 | ], |
| 36 | "memo info": [ | 36 | "memo info": [ |
| 37 | { | 37 | { |
| 38 | "rank": 0, | 38 | "rank": 0, |
| 39 | - "allocated memory": 19434.04541015625, | 39 | + "allocated memory": 19593.77734375, |
| 40 | - "max allocated memory": 28959.240234375 | 40 | + "max allocated memory": 28277.3251953125 |
| 41 | } | 41 | } |
| 42 | ] | 42 | ] |
| 43 | } | 43 | } |
| @@ -46,6 +46,7 @@ MOE_ARGS=" | |||
| 46 | --moe-permutation-async-comm \ | 46 | --moe-permutation-async-comm \ |
| 47 | --moe-grouped-gemm \ | 47 | --moe-grouped-gemm \ |
| 48 | --moe-token-dispatcher-type allgather \ | 48 | --moe-token-dispatcher-type allgather \ |
| 49 | + --moe-allgather-overlap-comm \ | ||
| 49 | --first-k-dense-replace 1 \ | 50 | --first-k-dense-replace 1 \ |
| 50 | --moe-layer-freq 1 \ | 51 | --moe-layer-freq 1 \ |
| 51 | --n-shared-experts 2 \ | 52 | --n-shared-experts 2 \ |