已合并
dualpipev dw detach #2154
zhao-yifan27创建于 2025年4月3日
dualpipev dw detach #2154
已合并
从refs/pull/2154/head合入到core_r0.8.0
共 4 个文件变更+11-4
| @@ -15,7 +15,8 @@ DualPipe流水不仅可以创造跨microbatch计算通信并行的条件，实 | |||
| 15 |  | 15 |  |
| 16 | 图中0~9和10~19代表同一个microbatch在同一张卡上的两个stage。绿色部分代表不同microbatch的前反向并行，开启跨micro batch前反向通信掩盖后，All2All通信和P2P通信都可以被没有依赖关系的计算掩盖，该特性的详细介绍参见[MoE跨microbatch前反向通信掩盖](megatron_moe/megatron-moe-fb-overlap.md)。 | 16 | 图中0~9和10~19代表同一个microbatch在同一张卡上的两个stage。绿色部分代表不同microbatch的前反向并行，开启跨micro batch前反向通信掩盖后，All2All通信和P2P通信都可以被没有依赖关系的计算掩盖，该特性的详细介绍参见[MoE跨microbatch前反向通信掩盖](megatron_moe/megatron-moe-fb-overlap.md)。 |
| 17 | 17 | ||
| 18 | -在warmup阶段，MindSpeed实现了尽可能多得P2P通信掩盖。在cooldown阶段，DualPipeV实现中PP尾卡会连续计算PP_size个反向stage，在对应的dw计算完成之前，激活值并不会完全释放，这会导致在重计算或其他内存特性场景下，cooldown阶段的峰值内存被大大拉高。因此MindSpeed实现中取消了cooldown阶段的dw分离，以较小的性能代价降低了重计算场景下内存峰值。 | 18 | +在warmup阶段，MindSpeed实现了尽可能多得P2P通信掩盖。在cooldown阶段，DualPipeV实现中PP尾卡会连续计算PP_size个反向stage，在对应的dw计算完成之前，激活值并不会完全释放，这会导致在某些重计算或其他内存特性场景下，cooldown阶段的峰值内存被大大拉高。因此MindSpeed实现中默认取消了cooldown阶段的dw分离，以较小的性能代价降低了重计算场景下内存峰值。同时也给出了dw分离的参数选项，开启dw分离时的流水如下图所示。 |
| 19 | +。 | ||
| 19 | 20 | ||
| 20 | 下图展示了在DeepseekV3 671B模型上采用PP8 TP2 EP32 DualPipeV策略采集的PP通信组profiling。 | 21 | 下图展示了在DeepseekV3 671B模型上采用PP8 TP2 EP32 DualPipeV策略采集的PP通信组profiling。 |
| 21 |  | 22 |  |
| @@ -50,6 +51,8 @@ DualPipe流水不仅可以创造跨microbatch计算通信并行的条件，实 | |||
| 50 | ## 使用方法 | 51 | ## 使用方法 |
| 51 | 在启动脚本中添加`--schedules-method dualpipev` 和`--moe-fb-overlap`。 | 52 | 在启动脚本中添加`--schedules-method dualpipev` 和`--moe-fb-overlap`。 |
| 52 | 53 | ||
| 54 | +在启动脚本中额外添加`--dualpipev-dw-detach` 来开启cooldown阶段的dw分离。 | ||
| 55 | + | ||
| 53 | 使用DualPipeV时，模型层数设置应为`PP*2`的倍数。同时每个PP组的micro batch数至少设置为`PP*2`。 | 56 | 使用DualPipeV时，模型层数设置应为`PP*2`的倍数。同时每个PP组的micro batch数至少设置为`PP*2`。 |
| 54 | 57 | ||
| 55 | 58 | ||
| @@ -708,7 +708,6 @@ def forward_backward_pipelining_with_cutinhalf( | |||
| 708 | ): | 708 | ): |
| 709 | args = get_args() | 709 | args = get_args() |
| 710 | args.moe_fb_overlap = True | 710 | args.moe_fb_overlap = True |
| 711 | - args.dualpipe_no_dw_detach = True | ||
| 712 | 711 | ||
| 713 | set_shared_embedding_from_dual_chunk(model[0], model[1]) | 712 | set_shared_embedding_from_dual_chunk(model[0], model[1]) |
| 714 | assert ( | 713 | assert ( |
| @@ -1435,7 +1434,7 @@ def forward_backward_pipelining_with_cutinhalf( | |||
| 1435 | input_tensor_bwd = merged_input_tensors.pop(0)[1] | 1434 | input_tensor_bwd = merged_input_tensors.pop(0)[1] |
| 1436 | output_tensor_bwd, bwd_model_chunk_id = merged_output_tensors.pop(0) | 1435 | output_tensor_bwd, bwd_model_chunk_id = merged_output_tensors.pop(0) |
| 1437 | 1436 | ||
| 1438 | - if not args.dualpipe_no_dw_detach: | 1437 | + if args.dualpipev_dw_detach: |
| 1439 | WeightGradStore.start_decouple() | 1438 | WeightGradStore.start_decouple() |
| 1440 | 1439 | ||
| 1441 | if args.moe_fb_overlap: | 1440 | if args.moe_fb_overlap: |
| @@ -1449,7 +1448,7 @@ def forward_backward_pipelining_with_cutinhalf( | |||
| 1449 | input_tensor_bwd, output_tensor_bwd, output_tensor_grad_bwd, model_type, config | 1448 | input_tensor_bwd, output_tensor_bwd, output_tensor_grad_bwd, model_type, config |
| 1450 | ) | 1449 | ) |
| 1451 | 1450 | ||
| 1452 | - if not args.dualpipe_no_dw_detach: | 1451 | + if args.dualpipev_dw_detach: |
| 1453 | WeightGradStore.end_decouple() | 1452 | WeightGradStore.end_decouple() |
| 1454 | 1453 | ||
| 1455 | if i == pp_size - 1: | 1454 | if i == pp_size - 1: |
| @@ -14,12 +14,17 @@ class DualpipeVFeature(MindSpeedFeature): | |||
| 14 | group = parser.add_argument_group(title=self.feature_name) | 14 | group = parser.add_argument_group(title=self.feature_name) |
| 15 | group.add_argument('--schedules-method', type=str, | 15 | group.add_argument('--schedules-method', type=str, |
| 16 | default=None, choices=['dualpipev']) | 16 | default=None, choices=['dualpipev']) |
| 17 | + group.add_argument('--dualpipev-dw-detach', action='store_true', | ||
| 18 | + help='detach dw in cooldown to reduce bubble') | ||
| 17 | 19 | ||
| 18 | def validate_args(self, args): | 20 | def validate_args(self, args): |
| 19 | if args.schedules_method == "dualpipev": | 21 | if args.schedules_method == "dualpipev": |
| 20 | if args.num_layers_per_virtual_pipeline_stage is not None: | 22 | if args.num_layers_per_virtual_pipeline_stage is not None: |
| 21 | raise AssertionError( | 23 | raise AssertionError( |
| 22 | "The dualpipev and virtual_pipeline are incompatible.") | 24 | "The dualpipev and virtual_pipeline are incompatible.") |
| 25 | + if args.pipeline_model_parallel_size == 1: | ||
| 26 | + raise AssertionError( | ||
| 27 | + "pipeline_model_parallel_size should be larger than 1 with dualpipev schedules") | ||
| 23 | if args.num_layers < args.pipeline_model_parallel_size * 2: | 28 | if args.num_layers < args.pipeline_model_parallel_size * 2: |
| 24 | raise AssertionError( | 29 | raise AssertionError( |
| 25 | 'number of layers must be at least 2*pipeline_model_parallel_size in dualpipe') | 30 | 'number of layers must be at least 2*pipeline_model_parallel_size in dualpipe') |