Star207
313
代码
代码
Issues71
Pull Requests91
流水线
项目讨论
Wiki
分析
Star207
313
已合并
[pytorch][feature] update cp args #3129
wx_4e191bb7af创建于 2025年8月12日
[pytorch][feature] update cp args #3129
已合并
wx_4e191bb7af创建于 2025年8月12日
从refs/pull/3129/head合入到master
共 22 个文件变更+43-308
Mdocs/pytorch/features/fine-tuning-with-context-parallel.md+14-22
@@ -1,10 +1,11 @@
1-### 长序列微调1+# 长序列微调
2 2 
3-## 数据集处理方法3+## 使用方法
4-数据集处理方法同[**多样本pack微调**](../solutions/finetune/multi_sample_pack_finetune.md)。
5 4 
5+### 数据预处理
6+数据预处理方法同[**多样本pack微调**](../solutions/finetune/multi_sample_pack_finetune.md)。
6 7 
7-## 启动长序列微调脚本8+### 微调参数
8 9 
9【--is-instruction-dataset】10【--is-instruction-dataset】
10 11 
@@ -22,16 +23,10 @@
22 23 
23CP切分的并行数目，要可以被序列长度整除。24CP切分的并行数目，要可以被序列长度整除。
24 25 
25-【--cp-attention-mask-type】26+【--attention-mask-type】
26 27 
27设置mask类型，微调开启CP的场景下只能为general。28设置mask类型，微调开启CP的场景下只能为general。
28 29 
29-【--adaptive-cp-manually-set-mask-list】
30- 
31-长序列场景完整mask会占用大量显存（约`seq-length * seq_length * 2`），并且影响端到端性能；当前pack模式长序列微调场景下，因mask具有一定规律性，
32-当`--context-parallel-algo`设置为 `adaptive_cp_algo` 或 `hybrid_adaptive_cp_algo`时，使能`--adaptive-cp-manually-set-mask-list`可以不生成完整的mask，
33-使能每个rank生成各自独立的mask_list，进而节省显存，提高端到端性能。
34- 
35【--context-parallel-algo】30【--context-parallel-algo】
36 31 
37通过传入指定参数，选择不同的cp算法，具体包含如下几种：32通过传入指定参数，选择不同的cp算法，具体包含如下几种：
@@ -40,23 +35,20 @@ CP切分的并行数目，要可以被序列长度整除。
402. [**ulysses_cp_algo**](https://gitee.com/ascend/MindSpeed/blob/master/docs/features/ulysses-context-parallel.md)352. [**ulysses_cp_algo**](https://gitee.com/ascend/MindSpeed/blob/master/docs/features/ulysses-context-parallel.md)
413. [**hybrid_cp_algo**](https://gitee.com/ascend/MindSpeed/blob/master/docs/features/hybrid-context-parallel.md)363. [**hybrid_cp_algo**](https://gitee.com/ascend/MindSpeed/blob/master/docs/features/hybrid-context-parallel.md)
42 37 
43-由于在微调场景，`--cp-attention-mask-type`只能设置为`general`，当训练样本为多条短样本拼接，CP较大时，建议优先选择`adaptive_cp_algo`，38+由于在微调场景，`--attention-mask-type`只能设置为`general`，
44理论上样本越短，拼接序列包含的样本数目越多，`--context-parallel-size`设置越大，性能收益越明显，但是要注意 seq-length / context-parallel-size > 8k时可以39理论上样本越短，拼接序列包含的样本数目越多，`--context-parallel-size`设置越大，性能收益越明显，但是要注意 seq-length / context-parallel-size > 8k时可以
45尽可能弥补CP带来的通信损失，针对该种场景参考配置如下，参数相关介绍参考上述对应算法的链接。CP较小时（一般<=4），`ulysses_cp_algo`是性能不错的选择。40尽可能弥补CP带来的通信损失，针对该种场景参考配置如下，参数相关介绍参考上述对应算法的链接。CP较小时（一般<=4），`ulysses_cp_algo`是性能不错的选择。
46 41 
47```bash42```bash
48 --seq-length 13107243 --seq-length 131072
49 --context-parallel-size 844 --context-parallel-size 8
50- --context-parallel-algo adaptive_cp_algo45+ --context-parallel-algo megatron_cp_algo
51- --cp-attention-mask-type general46+ --attention-mask-type general
52- --adaptive-cp-manually-set-mask-list
53- --adaptive-cp-dynamic-attn-mask
54- --adaptive-cp-only-reschedule
55```47```
56 48 
57## 使用效果49## 使用效果
58-| 模型 | 序列长度 | 分布式策略（TP/PP/CP） | gbs | CP类型 | cp-attention-mask-type | reset-position-ids | 显存 | 吞吐 TFLOP/s/GPU |50+| 模型 | 序列长度 | 分布式策略（TP/PP/CP） | gbs | CP类型 | attention-mask-type | reset-position-ids | 显存 | 吞吐 TFLOP/s/GPU |
59-|:---------:|:----:|:---------------:|:---:|:----------------:|:----------------------:|:------------------:|:-----:|:--------------:|51+|:---------:|:----:|:---------------:|:---:|:----------------:|:-------------------:|:------------------:|:-----:|:--------------:|
60-| Llama2-7B | 32k | 2/1/4 | 16 | megatron_cp_algo | general | True | 52777 | 102.7 |52+| Llama2-7B | 32k | 2/1/4 | 16 | megatron_cp_algo | general | True | 52777 | 102.7 |
61-| Llama2-7B | 32k | 2/1/4 | 16 | ulysses_cp_algo | general | True | 53681 | 192.3 |53+| Llama2-7B | 32k | 2/1/4 | 16 | ulysses_cp_algo | general | True | 53681 | 192.3 |
62-| Llama2-7B | 32k | 2/1/4 | 16 | adaptive_cp_algo | general | True | 52461 | 139.1 |54+ 
Mexamples/mcore/llama2/tune_llama2_7b_full_pack_16k.sh+1-1
@@ -38,7 +38,7 @@ GPT_ARGS="
38 --load ${CKPT_LOAD_DIR} \38 --load ${CKPT_LOAD_DIR} \
39 --context-parallel-algo megatron_cp_algo \39 --context-parallel-algo megatron_cp_algo \
40 --sequence-parallel \40 --sequence-parallel \
41- --cp-attention-mask-type general \41+ --attention-mask-type general \
42 --swap-attention \42 --swap-attention \
43 --num-workers 16 \43 --num-workers 16 \
44 --cp-window-size 1 \44 --cp-window-size 1 \
Mexamples/mcore/llama2/tune_llama2_7b_full_pack_32k.sh+1-1
@@ -39,7 +39,7 @@ GPT_ARGS="
39 --context-parallel-algo ulysses_cp_algo \39 --context-parallel-algo ulysses_cp_algo \
40 --ulysses-degree-in-cp 2 \40 --ulysses-degree-in-cp 2 \
41 --sequence-parallel \41 --sequence-parallel \
42- --cp-attention-mask-type causal \42+ --attention-mask-type causal \
43 --num-workers 16 \43 --num-workers 16 \
44 --cp-window-size 1 \44 --cp-window-size 1 \
45 --use-fused-rotary-pos-emb \45 --use-fused-rotary-pos-emb \
Mexamples/mcore/llama3/dpo_llama3_8b_full_ptd.sh+1-1
@@ -25,7 +25,7 @@ DISTRIBUTED_ARGS="
25 25 
26GPT_ARGS="26GPT_ARGS="
27 --variable-seq-lengths \27 --variable-seq-lengths \
28- --cp-attention-mask-type general \28+ --attention-mask-type general \
29 --sequence-parallel \29 --sequence-parallel \
30 --use-fused-rotary-pos-emb \30 --use-fused-rotary-pos-emb \
31 --use-deter-comp \31 --use-deter-comp \
Mexamples/mcore/mixtral/train_orm_mixtral_8x7B.sh+1-1
@@ -84,7 +84,7 @@ GPT_ARGS="
84 --sequence-parallel \84 --sequence-parallel \
85 --context-parallel-size ${CP} \85 --context-parallel-size ${CP} \
86 --context-parallel-algo ${CP_TYPE} \86 --context-parallel-algo ${CP_TYPE} \
87- --cp-attention-mask-type general \87+ --attention-mask-type general \
88 --tokenizer-type PretrainedFromHF \88 --tokenizer-type PretrainedFromHF \
89 --tokenizer-name-or-path ${TOKENIZER_PATH} \89 --tokenizer-name-or-path ${TOKENIZER_PATH} \
90 --load ${CKPT_LOAD_DIR} \90 --load ${CKPT_LOAD_DIR} \
Mexamples/mindspore/llama2/tune_llama2_7b_full_pack_16k_ms.sh+1-1
@@ -39,7 +39,7 @@ GPT_ARGS="
39 --load ${CKPT_LOAD_DIR} \39 --load ${CKPT_LOAD_DIR} \
40 --context-parallel-algo megatron_cp_algo \40 --context-parallel-algo megatron_cp_algo \
41 --sequence-parallel \41 --sequence-parallel \
42- --cp-attention-mask-type general \42+ --attention-mask-type general \
43 --swap-attention \43 --swap-attention \
44 --num-workers 16 \44 --num-workers 16 \
45 --cp-window-size 1 \45 --cp-window-size 1 \
Mexamples/mindspore/llama2/tune_llama2_7b_full_pack_32k_ms.sh+1-1
@@ -40,7 +40,7 @@ GPT_ARGS="
40 --context-parallel-algo ulysses_cp_algo \40 --context-parallel-algo ulysses_cp_algo \
41 --ulysses-degree-in-cp 2 \41 --ulysses-degree-in-cp 2 \
42 --sequence-parallel \42 --sequence-parallel \
43- --cp-attention-mask-type causal \43+ --attention-mask-type causal \
44 --num-workers 16 \44 --num-workers 16 \
45 --cp-window-size 1 \45 --cp-window-size 1 \
46 --use-fused-rotary-pos-emb \46 --use-fused-rotary-pos-emb \
Mmindspeed_llm/core/transformer/dot_product_attention.py+2-3
@@ -76,7 +76,7 @@ def do_ring_context_parallel(q, k, v, head_num, softmax_scale, attn_mask, dropou
76 76 
77 cp_para = dict()77 cp_para = dict()
78 78 
79- cp_para['causal'] = args.cp_attention_mask_type == 'causal'79+ cp_para['causal'] = args.attention_mask_type == 'causal'
80 cp_para['cp_group'] = cp_group80 cp_para['cp_group'] = cp_group
81 cp_para['cp_size'] = cp_size81 cp_para['cp_size'] = cp_size
82 cp_para['rank'] = rank82 cp_para['rank'] = rank
@@ -418,8 +418,7 @@ def flash_attention_forward(
418 if actual_seq_len is not None and args.mtp_num_layers:418 if actual_seq_len is not None and args.mtp_num_layers:
419 actual_seq_len = actual_seq_len[self.mtp_idx]419 actual_seq_len = actual_seq_len[self.mtp_idx]
420 420 
421- if args.context_parallel_size > 1 and args.context_parallel_algo in ['megatron_cp_algo', 'hybrid_cp_algo',421+ if args.context_parallel_size > 1 and args.context_parallel_algo in ['megatron_cp_algo', 'hybrid_cp_algo']:
422- 'adaptive_cp_algo', 'hybrid_adaptive_cp_algo']:
423 query, key, value = [rearrange(x, 's b h d -> s b (h d)') for x in [query, key, value]]422 query, key, value = [rearrange(x, 's b h d -> s b (h d)') for x in [query, key, value]]
424 return do_ring_context_parallel(423 return do_ring_context_parallel(
425 query, key, value, head_num=n_head, softmax_scale=scale, attn_mask=attention_mask, pse=self.pse,424 query, key, value, head_num=n_head, softmax_scale=scale, attn_mask=attention_mask, pse=self.pse,
Mmindspeed_llm/mindspore/core/transformer/dot_product_attention.py+1-2
@@ -47,8 +47,7 @@ def flash_attention_forward(
47 if actual_seq_len is not None and args.mtp_num_layers:47 if actual_seq_len is not None and args.mtp_num_layers:
48 actual_seq_len = actual_seq_len[self.mtp_idx]48 actual_seq_len = actual_seq_len[self.mtp_idx]
49 49 
50- if args.context_parallel_size > 1 and args.context_parallel_algo in ['megatron_cp_algo', 'hybrid_cp_algo',50+ if args.context_parallel_size > 1 and args.context_parallel_algo in ['megatron_cp_algo', 'hybrid_cp_algo']:
51- 'adaptive_cp_algo', 'hybrid_adaptive_cp_algo']:
52 query, key, value = [rearrange(x, 's b h d -> s b (h d)') for x in [query, key, value]]51 query, key, value = [rearrange(x, 's b h d -> s b (h d)') for x in [query, key, value]]
53 return do_ring_context_parallel(52 return do_ring_context_parallel(
54 query, key, value, head_num=n_head, softmax_scale=scale, attn_mask=attention_mask, pse=self.pse,53 query, key, value, head_num=n_head, softmax_scale=scale, attn_mask=attention_mask, pse=self.pse,
Mmindspeed_llm/tasks/megatron_adaptor.py+0-5
@@ -638,11 +638,6 @@ class CoreAdaptation(MegatronAdaptationABC):
638 MegatronAdaptation.register('megatron.training.checkpointing.unwrap_model', unwrap_model_wrapper)638 MegatronAdaptation.register('megatron.training.checkpointing.unwrap_model', unwrap_model_wrapper)
639 MegatronAdaptation.register('megatron.training.training.unwrap_model', unwrap_model_wrapper)639 MegatronAdaptation.register('megatron.training.training.unwrap_model', unwrap_model_wrapper)
640 640 
641- from mindspeed_llm.training.utils import generate_adaptive_cp_mask_list_by_user, generate_adaptive_cp_grid_mask_by_user
642- MegatronAdaptation.register('mindspeed.core.context_parallel.utils.generate_adaptive_cp_mask_list_by_user',
643- generate_adaptive_cp_mask_list_by_user)
644- MegatronAdaptation.register('mindspeed.core.context_parallel.utils.generate_adaptive_cp_grid_mask_by_user',
645- generate_adaptive_cp_grid_mask_by_user)
646 from mindspeed.training import get_device_arch_version641 from mindspeed.training import get_device_arch_version
647 MegatronAdaptation.register('megatron.training.utils.get_device_arch_version', get_device_arch_version)642 MegatronAdaptation.register('megatron.training.utils.get_device_arch_version', get_device_arch_version)
648 # Currently, it is not supported to Cast shard fp32 main params to fp8 model params643 # Currently, it is not supported to Cast shard fp32 main params to fp8 model params
Mmindspeed_llm/training/arguments.py+6-23
@@ -206,12 +206,11 @@ def _add_cp_args(parser):
206 group.add_argument('--context-parallel-size', type=int, default=1,206 group.add_argument('--context-parallel-size', type=int, default=1,
207 help='Degree of context parallelism.')207 help='Degree of context parallelism.')
208 group.add_argument('--context-parallel-algo', type=str, default='ulysses_cp_algo',208 group.add_argument('--context-parallel-algo', type=str, default='ulysses_cp_algo',
209- choices=['ulysses_cp_algo', 'megatron_cp_algo', 'hybrid_cp_algo', 'adaptive_cp_algo',209+ choices=['ulysses_cp_algo', 'megatron_cp_algo', 'hybrid_cp_algo', 'mamba_cp_algo'], help='context parallel algorithm')
210- 'hybrid_adaptive_cp_algo', 'mamba_cp_algo'], help='context parallel algorithm')
211 group.add_argument('--ulysses-degree-in-cp', type=int, default=None)210 group.add_argument('--ulysses-degree-in-cp', type=int, default=None)
212 group.add_argument('--attention-mask-type', type=str, default='causal',211 group.add_argument('--attention-mask-type', type=str, default='causal',
213 choices=['causal', 'general'], help='context parallel attention mask type')212 choices=['causal', 'general'], help='context parallel attention mask type')
214- group.add_argument('--cp-attention-mask-type', type=str, default='causal',213+ group.add_argument('--attention-mask-type', type=str, default='causal',
215 choices=['causal', 'general'], help='context parallel attention mask type')214 choices=['causal', 'general'], help='context parallel attention mask type')
216 group.add_argument('--use-cp-send-recv-overlap', action='store_true',215 group.add_argument('--use-cp-send-recv-overlap', action='store_true',
217 help='use it to enable cp send-recv-overlap.')216 help='use it to enable cp send-recv-overlap.')
@@ -288,8 +287,8 @@ def _validate_cp_args(args):
288 assert args.cp_window_size >= 1 and args.cp_window_size < args.context_parallel_size, f'cp_window_size should in range [1, context_parallel_size) when using double_ring_attention.'287 assert args.cp_window_size >= 1 and args.cp_window_size < args.context_parallel_size, f'cp_window_size should in range [1, context_parallel_size) when using double_ring_attention.'
289 n_window, remainder = divmod(args.context_parallel_size, args.cp_window_size)288 n_window, remainder = divmod(args.context_parallel_size, args.cp_window_size)
290 assert n_window >= 1 and remainder == 0, f'context parallel size must be divisible by cp_window_size when using double ring attention.'289 assert n_window >= 1 and remainder == 0, f'context parallel size must be divisible by cp_window_size when using double ring attention.'
291- if args.cp_attention_mask_type == 'general':290+ if args.attention_mask_type == 'general':
292- assert args.micro_batch_size == 1, f'When cp_attention_mask_type is set to general, the value of mbs can only be 1.'291+ assert args.micro_batch_size == 1, f'When attention_mask_type is set to general, the value of mbs can only be 1.'
293 292 
294 if args.context_parallel_algo == 'hybrid_cp_algo':293 if args.context_parallel_algo == 'hybrid_cp_algo':
295 assert args.ulysses_degree_in_cp is not None, "--ulysses-degree-in-cp must be specified in hybrid_cp_algo"294 assert args.ulysses_degree_in_cp is not None, "--ulysses-degree-in-cp must be specified in hybrid_cp_algo"
@@ -301,23 +300,8 @@ def _validate_cp_args(args):
301 n_window, remainder = divmod(ring_degree, args.cp_window_size)300 n_window, remainder = divmod(ring_degree, args.cp_window_size)
302 assert n_window >= 1 and remainder == 0, f'ring_degree should be divisible by cp_window_size when using double ring with hybrid context parallelism.'301 assert n_window >= 1 and remainder == 0, f'ring_degree should be divisible by cp_window_size when using double ring with hybrid context parallelism.'
303 _check_attention_head(args, args.ulysses_degree_in_cp)302 _check_attention_head(args, args.ulysses_degree_in_cp)
304- if args.cp_attention_mask_type == 'general':303+ if args.attention_mask_type == 'general':
305- assert args.micro_batch_size == 1, f'When cp_attention_mask_type is set to general, the value of mbs can only be 1.'304+ assert args.micro_batch_size == 1, f'When attention_mask_type is set to general, the value of mbs can only be 1.'
306- 
307- if args.context_parallel_size > 1 and args.context_parallel_algo == 'adaptive_cp_algo':
308- assert args.seq_length % args.context_parallel_size == 0, f"sequence length must be divisible by context_parallel_size"
309- if args.cp_attention_mask_type == 'general':
310- assert args.micro_batch_size == 1, f'When cp_attention_mask_type is set to general, the value of mbs can only be 1.'
311- 
312- if args.context_parallel_size > 1 and args.context_parallel_algo == 'hybrid_adaptive_cp_algo':
313- assert args.ulysses_degree_in_cp is not None, "--ulysses-degree-in-cp must be specified in hybrid_adaptive_cp_algo"
314- ring_degree, remainder = divmod(args.context_parallel_size, args.ulysses_degree_in_cp)
315- assert ring_degree > 1 and remainder == 0, "--ulysses-degree-in-cp must be devisible by --context-parallel-size"
316- head, remainder = divmod(args.num_attention_heads, args.ulysses_degree_in_cp * args.tensor_model_parallel_size)
317- assert head >= 1 and remainder == 0, f"num_attention_heads must be divisible by ulysse-degree-in-cp * tensor_model_parallel_size in hybrid cp"
318- assert args.seq_length % args.context_parallel_size == 0, f"sequence length must be divisible by context_parallel_size in hybrid cp"
319- if args.cp_attention_mask_type == 'general':
320- assert args.micro_batch_size == 1, f'When cp_attention_mask_type is set to general, the value of mbs can only be 1.'
321 305 
322 if args.sliding_window:306 if args.sliding_window:
323 raise AssertionError("sliding window is not supported in context parallel.")307 raise AssertionError("sliding window is not supported in context parallel.")
@@ -1271,7 +1255,6 @@ def _add_dummy_args(args):
1271 args.use_nanopipe = False1255 args.use_nanopipe = False
1272 args.moe_without_activation = False1256 args.moe_without_activation = False
1273 args.ampipe_degree = 01257 args.ampipe_degree = 0
1274- args.attention_mask_type = args.cp_attention_mask_type
1275 args.hccl_group_buffer_adaptive = False1258 args.hccl_group_buffer_adaptive = False
1276 args.moe_bmm_mc2 = False1259 args.moe_bmm_mc2 = False
1277 args.moe_hierarchical_alltoallv = False1260 args.moe_hierarchical_alltoallv = False
Mmindspeed_llm/training/utils.py+3-92
@@ -38,8 +38,6 @@ from mindspeed.utils import (set_actual_seq_len, set_position_ids,
38 _get_batch_on_this_cp_rank_in_ulysses_cp,38 _get_batch_on_this_cp_rank_in_ulysses_cp,
39 _get_batch_on_this_cp_rank_in_hybrid_cp_general,39 _get_batch_on_this_cp_rank_in_hybrid_cp_general,
40 _get_batch_on_this_cp_rank_in_hybrid_cp,40 _get_batch_on_this_cp_rank_in_hybrid_cp,
41- _get_batch_on_this_cp_rank_in_adaptive_cp,
42- _get_batch_on_this_cp_rank_in_hybrid_adaptive_cp,
43 broadcast_dynamic, _broadcast, get_ring_degree)41 broadcast_dynamic, _broadcast, get_ring_degree)
44from mindspeed.core.tensor_parallel_y_union_cp import TensorParallelYUnionCP42from mindspeed.core.tensor_parallel_y_union_cp import TensorParallelYUnionCP
45from mindspeed.model.transformer import set_attention_mask43from mindspeed.model.transformer import set_attention_mask
@@ -726,12 +724,12 @@ def get_batch_on_this_cp_rank(batch):
726 if not tp_y_cp_size > 1:724 if not tp_y_cp_size > 1:
727 return batch725 return batch
728 726 
729- if args.cp_attention_mask_type == 'general' and batch.get("attention_mask", None) is not None:727+ if args.attention_mask_type == 'general' and batch.get("attention_mask", None) is not None:
730 set_attention_mask(batch['attention_mask'].squeeze())728 set_attention_mask(batch['attention_mask'].squeeze())
731 729 
732 cp_expanded_by_2d_tp = args.tp_y > 1730 cp_expanded_by_2d_tp = args.tp_y > 1
733 if args.context_parallel_algo == 'megatron_cp_algo':731 if args.context_parallel_algo == 'megatron_cp_algo':
734- if args.cp_attention_mask_type == 'general':732+ if args.attention_mask_type == 'general':
735 batch = _get_batch_on_this_cp_rank_in_megatron_cp_general(batch)733 batch = _get_batch_on_this_cp_rank_in_megatron_cp_general(batch)
736 elif cp_expanded_by_2d_tp:734 elif cp_expanded_by_2d_tp:
737 batch = _get_batch_on_this_tp_y_cp_rank_in_megatron_cp(batch)735 batch = _get_batch_on_this_tp_y_cp_rank_in_megatron_cp(batch)
@@ -740,100 +738,13 @@ def get_batch_on_this_cp_rank(batch):
740 elif args.context_parallel_algo == 'ulysses_cp_algo' or args.context_parallel_algo == 'mamba_cp_algo':738 elif args.context_parallel_algo == 'ulysses_cp_algo' or args.context_parallel_algo == 'mamba_cp_algo':
741 batch = _get_batch_on_this_cp_rank_in_ulysses_cp(batch)739 batch = _get_batch_on_this_cp_rank_in_ulysses_cp(batch)
742 elif args.context_parallel_algo == 'hybrid_cp_algo':740 elif args.context_parallel_algo == 'hybrid_cp_algo':
743- if args.cp_attention_mask_type == 'general':741+ if args.attention_mask_type == 'general':
744 batch = _get_batch_on_this_cp_rank_in_hybrid_cp_general(batch)742 batch = _get_batch_on_this_cp_rank_in_hybrid_cp_general(batch)
745 else:743 else:
746 batch = _get_batch_on_this_cp_rank_in_hybrid_cp(batch)744 batch = _get_batch_on_this_cp_rank_in_hybrid_cp(batch)
747- elif args.context_parallel_algo == 'adaptive_cp_algo':
748- batch = _get_batch_on_this_cp_rank_in_adaptive_cp(batch)
749- elif args.context_parallel_algo == 'hybrid_adaptive_cp_algo':
750- batch = _get_batch_on_this_cp_rank_in_hybrid_adaptive_cp(batch)
751 return batch745 return batch
752 746 
753 747 
754-def generate_adaptive_cp_grid_mask_by_user(cp_size):
755- from mindspeed.utils import get_actual_seq_len
756- from mindspeed.core.context_parallel.utils import set_adaptive_cp_grid_mask_by_user
757- args = get_args()
758- actual_seq_len = get_actual_seq_len()
759- seq_length = args.seq_length
760- grid_mask = torch.zeros(cp_size, cp_size)
761- sub_seq_length = seq_length // cp_size
762- 
763- grid_actual_seq_len_dict = {}
764- for seq_len in actual_seq_len:
765- grid_actual_seq_len_dict[seq_len // sub_seq_length + 1] = seq_len % sub_seq_length == 0
766- grid_actual_seq_len = list(grid_actual_seq_len_dict.items())
767- start_index = 0
768- for i, _ in enumerate(grid_actual_seq_len):
769- end_index = grid_actual_seq_len[i][0]
770- grid_mask[start_index:end_index, start_index:end_index] = 1
771- 
772- if i != 0:
773- if grid_actual_seq_len[i - 1][1]:
774- start_index = grid_actual_seq_len[i - 1][0] - 1
775- else:
776- start_index = grid_actual_seq_len[i - 1][0]
777- grid_mask = torch.tril(grid_mask)
778- set_adaptive_cp_grid_mask_by_user(grid_mask)
779- 
780- 
781-def generate_adaptive_cp_mask_list_by_user(opt_seq, opt_scheduling, cp_rank, cp_size):
782- from mindspeed.utils import get_actual_seq_len
783- from mindspeed.core.context_parallel.utils import set_adaptive_cp_mask_list_by_user
784- actual_seq_len = get_actual_seq_len()
785- round_num = len(opt_scheduling)
786- grid_size = (opt_seq[-1] + 1) // cp_size
787- mask_list = []
788- for rnd_idx in range(round_num):
789- task_id = opt_scheduling[rnd_idx][cp_rank]
790- if task_id == -1:
791- mask_list.append(None)
792- continue
793- rank_x = task_id // cp_size
794- rank_y = task_id % cp_size
795- if rank_x == rank_y:
796- mask = torch.tril(torch.ones((grid_size, grid_size), device=torch.npu.current_device()))
797- for i in actual_seq_len:
798- if i - 1 < grid_size * rank_y:
799- continue
800- elif i - 1 >= grid_size * (rank_y + 1):
801- break
802- else:
803- mask[(i - 1 - grid_size * cp_rank + 1):, : (i - 1 - grid_size * cp_rank + 1)] = 0
804- elif cp_rank > rank_y:
805- mask = torch.zeros((grid_size, grid_size), device=torch.npu.current_device())
806- start_index = 0
807- end_index = grid_size
808- for i in actual_seq_len:
809- if i - 1 < grid_size * rank_y:
810- start_index = i - 1
811- continue
812- elif i - 1 >= grid_size * (rank_y + 1):
813- end_index = i - 1
814- break
815- else:
816- start_index = i - 1
817- start_index -= rank_y * grid_size
818- if start_index < 0:
819- start_index = 0
820- elif start_index > grid_size:
821- start_index = grid_size
822- end_index -= cp_rank * grid_size
823- if end_index < 0:
824- end_index = 0
825- elif end_index > grid_size:
826- end_index = grid_size
827- mask[: (end_index + 1), (start_index + 1):] = 1
828- else:
829- mask = torch.zeros((grid_size, grid_size), device=torch.npu.current_device())
830- if mask is not None:
831- # Convert attention mask to binary:
832- mask = mask < 0.5
833- mask_list.append(mask)
834- set_adaptive_cp_mask_list_by_user(mask_list)
835- 
836- 
837def _get_batch_on_this_cp_rank_in_megatron_cp_general(batch):748def _get_batch_on_this_cp_rank_in_megatron_cp_general(batch):
838 cp_rank = mpu.get_context_parallel_rank()749 cp_rank = mpu.get_context_parallel_rank()
839 cp_size = mpu.get_context_parallel_world_size()750 cp_size = mpu.get_context_parallel_world_size()
Mtests/README.md+1-9
@@ -14,7 +14,7 @@
14 <th>Mem.</th>14 <th>Mem.</th>
15 </tr>15 </tr>
16 <tr>16 <tr>
17- <td rowspan="16">ST</td>17+ <td rowspan="15">ST</td>
18 <td rowspan="14">Pretrain</td>18 <td rowspan="14">Pretrain</td>
19 <td>TP，PP，VPP，distributed_optimizer，o2_gradient，o2_optimizer，重计算，enable_recompute_layers_per_pp_rank，FA_TND，use_fused_rotary_pos_emb</td>19 <td>TP，PP，VPP，distributed_optimizer，o2_gradient，o2_optimizer，重计算，enable_recompute_layers_per_pp_rank，FA_TND，use_fused_rotary_pos_emb</td>
20 <td><a href="st/shell_scripts/llama2_tp2_pp4_vpp2_ptd.sh">llama2_tp2_pp4_vpp2_ptd.sh</a></td>20 <td><a href="st/shell_scripts/llama2_tp2_pp4_vpp2_ptd.sh">llama2_tp2_pp4_vpp2_ptd.sh</a></td>
@@ -113,14 +113,6 @@
113 <td>Y</td>113 <td>Y</td>
114 <td>Y</td>114 <td>Y</td>
115 </tr>115 </tr>
116- <tr>
117- <td rowspan="1">FullSFT</td>
118- <td>自适应cp，general_cp，SFT_pack_cp</td>
119- <td><a href="st/shell_scripts/tune_llama2_tp2_cp4_adaptive_cp.sh">tune_llama2_tp2_cp4_adaptive_cp.sh</a></td>
120- <td>Y</td>
121- <td>Y</td>
122- <td>Y</td>
123- </tr>
124 <tr>116 <tr>
125 <td rowspan="1">LoRA</td>117 <td rowspan="1">LoRA</td>
126 <td>CCLoRA, QLoRA</td>118 <td>CCLoRA, QLoRA</td>
Mtests/pipeline/context_parallel/test_adaptive_context_parallel.py+3-6
@@ -141,8 +141,7 @@ def run_adaptive_cp(cp_size, bs, seq_len, dtype, cp_args):
141 x.requires_grad = True141 x.requires_grad = True
142 142 
143 cp_para = dict()143 cp_para = dict()
144- cp_para['causal'] = args.cp_attention_mask_type == 'causal'144+ cp_para['causal'] = args.attention_mask_type == 'causal'
145- args.attention_mask_type = args.cp_attention_mask_type
146 cp_para['cp_group'] = cp_group145 cp_para['cp_group'] = cp_group
147 cp_para['cp_size'] = cp_size146 cp_para['cp_size'] = cp_size
148 cp_para['rank'] = rank147 cp_para['rank'] = rank
@@ -182,8 +181,7 @@ def run_hybrid_adaptive_cp(cp_size, bs, seq_len, dtype, cp_args):
182 from mindspeed.core.context_parallel.ulysses_context_parallel.ulysses_context_parallel import _SeqAllToAll181 from mindspeed.core.context_parallel.ulysses_context_parallel.ulysses_context_parallel import _SeqAllToAll
183 args = parse_args(None, True)182 args = parse_args(None, True)
184 args.seq_length = seq_len183 args.seq_length = seq_len
185- args.cp_attention_mask_type = 'general'184+ args.attention_mask_type = 'general'
186- args.attention_mask_type = args.cp_attention_mask_type
187 args.context_parallel_algo = 'hybrid_adaptive_cp_algo'185 args.context_parallel_algo = 'hybrid_adaptive_cp_algo'
188 args.ulysses_degree_in_cp = 2186 args.ulysses_degree_in_cp = 2
189 args.use_flash_attn = True187 args.use_flash_attn = True
@@ -245,8 +243,7 @@ def run_hybrid_adaptive_cp(cp_size, bs, seq_len, dtype, cp_args):
245 x.requires_grad = True243 x.requires_grad = True
246 244 
247 cp_para = dict()245 cp_para = dict()
248- cp_para['causal'] = args.cp_attention_mask_type == 'causal'246+ cp_para['causal'] = args.attention_mask_type == 'causal'
249- args.attention_mask_type = args.cp_attention_mask_type
250 cp_para['cp_group'] = get_context_parallel_group_for_hybrid_ring()247 cp_para['cp_group'] = get_context_parallel_group_for_hybrid_ring()
251 cp_para['cp_size'] = get_context_parallel_for_hybrid_ring_world_size()248 cp_para['cp_size'] = get_context_parallel_for_hybrid_ring_world_size()
252 cp_para['rank'] = get_context_parallel_for_hybrid_ring_rank()249 cp_para['rank'] = get_context_parallel_for_hybrid_ring_rank()
Mtests/pipeline/context_parallel/test_hybrid_context_parallel.py+1-2
@@ -53,8 +53,7 @@ def run_hybridattn_cp(test_args, cp_size, u_size, cp_args):
53 r_size = cp_size // u_size53 r_size = cp_size // u_size
54 args = parse_args(None, True)54 args = parse_args(None, True)
55 args.use_cp_send_recv_overlap = send_recv_overlap55 args.use_cp_send_recv_overlap = send_recv_overlap
56- args.cp_attention_mask_type = 'causal' if causal else 'general'56+ args.attention_mask_type = 'causal' if causal else 'general'
57- args.attention_mask_type = args.cp_attention_mask_type
58 # currently we always use FA in context parallel.57 # currently we always use FA in context parallel.
59 args.use_flash_attn = True58 args.use_flash_attn = True
60 if u_size == 1:59 if u_size == 1:
Mtests/pipeline/context_parallel/test_ringattn_context_parallel.py+1-2
@@ -46,8 +46,7 @@ def run_ringattn_cp(cp_size, bs, seq_len, dtype, cp_args):
46 args = parse_args(None, True)46 args = parse_args(None, True)
47 args.context_parallel_algo = 'megatron_cp_algo'47 args.context_parallel_algo = 'megatron_cp_algo'
48 args.use_cp_send_recv_overlap = send_recv_overlap48 args.use_cp_send_recv_overlap = send_recv_overlap
49- args.cp_attention_mask_type = 'causal' if causal else 'general'49+ args.attention_mask_type = 'causal' if causal else 'general'
50- args.attention_mask_type = args.cp_attention_mask_type
51 args.seq_length = seq_len50 args.seq_length = seq_len
52 args.use_flash_attn = True51 args.use_flash_attn = True
53 args.tp_2d = None52 args.tp_2d = None
Mtests/pipeline/model_module/test_attention.py+1-2
@@ -52,8 +52,7 @@ def run_attention_module(test_args, use_mcore, use_cp, cp_size, u_size, use_alib
52 r_size = cp_size // u_size52 r_size = cp_size // u_size
53 args = parse_args(None, True)53 args = parse_args(None, True)
54 args.use_cp_send_recv_overlap = True54 args.use_cp_send_recv_overlap = True
55- args.cp_attention_mask_type = 'causal'55+ args.attention_mask_type = 'causal'
56- args.attention_mask_type = args.cp_attention_mask_type
57 args.tp_2d = None56 args.tp_2d = None
58 args.tp_x = 157 args.tp_x = 1
59 args.tp_y = 158 args.tp_y = 1
Mtests/poc/gpt_moe_2t/pretrain_gpt_moe_2t_128k_GQA_256die_A2.sh+1-1
@@ -53,7 +53,7 @@ GPT_ARGS="
53 --group-query-attention \53 --group-query-attention \
54 --num-query-groups 8 \54 --num-query-groups 8 \
55 --reset-position-ids \55 --reset-position-ids \
56- --cp-attention-mask-type general \56+ --attention-mask-type general \
57 --no-gradient-accumulation-fusion \57 --no-gradient-accumulation-fusion \
58 --tensor-model-parallel-size ${TP} \58 --tensor-model-parallel-size ${TP} \
59 --pipeline-model-parallel-size ${PP} \59 --pipeline-model-parallel-size ${PP} \
Mtests/poc/gpt_moe_2t/pretrain_gpt_moe_2t_128k_GQA_256die_A3.sh+1-1
@@ -53,7 +53,7 @@ GPT_ARGS="
53 --group-query-attention \53 --group-query-attention \
54 --num-query-groups 8 \54 --num-query-groups 8 \
55 --reset-position-ids \55 --reset-position-ids \
56- --cp-attention-mask-type general \56+ --attention-mask-type general \
57 --no-gradient-accumulation-fusion \57 --no-gradient-accumulation-fusion \
58 --tensor-model-parallel-size ${TP} \58 --tensor-model-parallel-size ${TP} \
59 --pipeline-model-parallel-size ${PP} \59 --pipeline-model-parallel-size ${PP} \
Mtests/rlxf/train_orm_mixtral_tp2_pp2_cp2_ep2.sh+1-1
@@ -77,7 +77,7 @@ GPT_ARGS="
77 --sequence-parallel \77 --sequence-parallel \
78 --context-parallel-size ${CP} \78 --context-parallel-size ${CP} \
79 --context-parallel-algo ${CP_TYPE} \79 --context-parallel-algo ${CP_TYPE} \
80- --cp-attention-mask-type general \80+ --attention-mask-type general \
81 81 
82 --tokenizer-type PretrainedFromHF \82 --tokenizer-type PretrainedFromHF \
83 --tokenizer-name-or-path ${TOKENIZER_MODEL} \83 --tokenizer-name-or-path ${TOKENIZER_MODEL} \
Mtests/st/shell_scripts/llama2_tp2_cp4_general_double_ring.sh+1-1
@@ -35,7 +35,7 @@ DIST_ALGO=(
35 --cp-window-size 2 \35 --cp-window-size 2 \
36 --use-fused-ring-attention-update \36 --use-fused-ring-attention-update \
37 --context-parallel-algo megatron_cp_algo \37 --context-parallel-algo megatron_cp_algo \
38- --cp-attention-mask-type general \38+ --attention-mask-type general \
39 --sequence-parallel39 --sequence-parallel
40)40)
41 41 
Dtests/st/shell_scripts/tune_llama2_tp2_cp2_adaptive_cp.sh+0-130
@@ -1,130 +0,0 @@
1-#!/bin/bash
2- 
3-export CUDA_DEVICE_MAX_CONNECTIONS=1
4-export HCCL_DETERMINISITIC=True
5-export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True
6- 
7-NPUS_PER_NODE=4
8-MASTER_ADDR=localhost
9-MASTER_PORT=6015
10-NNODES=1
11-NODE_RANK=0
12-WORLD_SIZE=$(($NPUS_PER_NODE*$NNODES))
13- 
14-basepath=$(cd `dirname $0`; cd ../../../; pwd)
15- 
16-CKPT_LOAD_DIR="/data/ci/Llama2-mcore-tp2/"
17-DATA_PATH="/data/ci/llama2_pack/alpaca"
18-TOKENIZER_MODEL="/data/hf/llama-2-7b-hf"
19- 
20-TP=2
21-PP=1
22-CP=2
23- 
24-DISTRIBUTED_ARGS=(
25- --nproc_per_node $NPUS_PER_NODE
26- --nnodes $NNODES
27- --node_rank $NODE_RANK
28- --master_addr $MASTER_ADDR
29- --master_port $MASTER_PORT
30-)
31- 
32-DIST_ALGO=(
33- --tensor-model-parallel-size ${TP}
34- --pipeline-model-parallel-size ${PP}
35- --context-parallel-size ${CP}
36- --cp-window-size 1
37- --use-fused-ring-attention-update
38- --context-parallel-algo adaptive_cp_algo
39- --cp-attention-mask-type general
40- --adaptive-cp-manually-set-mask-list
41- --adaptive-cp-dynamic-attn-mask
42- --adaptive-cp-only-reschedule
43- --sequence-parallel
44-)
45- 
46-MODEL_ARGS=(
47- --use-mcore-models
48- --manual-gc
49- --manual-gc-interval 50
50- --num-layers 32
51- --hidden-size 4096
52- --ffn-hidden-size 11008
53- --num-attention-heads 32
54- --seq-length 4096
55- --max-position-embeddings 4096
56-)
57- 
58-ACCELERATE_ARGS=(
59- --reuse-fp32-param
60- --overlap-grad-reduce
61- --overlap-param-gather
62- --use-distributed-optimizer
63- --recompute-activation-function
64- --recompute-activation-function-num-layers 1
65-)
66- 
67-TRAINING_ARGS=(
68- --tokenizer-type PretrainedFromHF
69- --tokenizer-name-or-path ${TOKENIZER_MODEL}
70- --tokenizer-not-use-fast
71- --micro-batch-size 1
72- --global-batch-size 2
73- --make-vocab-size-divisible-by 1
74- --lr 1.25e-6
75- --train-iters 15
76- --lr-decay-style cosine
77- --untie-embeddings-and-output-weights
78- --disable-bias-linear
79- --attention-dropout 0.0
80- --init-method-std 0.01
81- --hidden-dropout 0.0
82- --position-embedding-type rope
83- --normalization RMSNorm
84- --use-fused-rmsnorm
85- --swiglu
86- --use-flash-attn
87- --reset-position-ids
88- --no-masked-softmax-fusion
89- --attention-softmax-in-fp32
90- --min-lr 1.25e-7
91- --weight-decay 1e-1
92- --lr-warmup-fraction 0.01
93- --clip-grad 1.0
94- --adam-beta1 0.9
95- --initial-loss-scale 65536
96- --adam-beta2 0.95
97- --no-gradient-accumulation-fusion
98- --no-load-optim
99- --no-load-rng
100- --use-fused-swiglu
101- --use-fused-rotary-pos-emb
102- --overlap-grad-reduce
103- --bf16
104- --finetune
105- --stage sft
106- --is-instruction-dataset
107-)
108- 
109-DATA_ARGS=(
110- --data-path $DATA_PATH
111- --split 949,50,1
112-)
113- 
114-OUTPUT_ARGS=(
115- --log-interval 1
116- --save-interval 10000
117- --eval-interval 1000
118- --eval-iters 0
119- --log-throughput
120-)
121- 
122-torchrun ${DISTRIBUTED_ARGS[@]} $basepath/posttrain_gpt.py \
123- ${DIST_ALGO[@]} \
124- ${MODEL_ARGS[@]} \
125- ${TRAINING_ARGS[@]} \
126- ${DATA_ARGS[@]} \
127- ${OUTPUT_ARGS[@]} \
128- ${ACCELERATE_ARGS[@]} \
129- --load ${CKPT_LOAD_DIR} \
130- --distributed-backend nccl