已合并
[pytorch][bugfix]update for support 32k-pack #2989
jzh6229创建于 2025年7月9日
[pytorch][bugfix]update for support 32k-pack #2989
已合并
从refs/pull/2989/head合入到2.0.0
共 2 个文件变更+27-9
| @@ -1,6 +1,7 @@ | |||
| 1 | # coding=utf-8 | 1 | # coding=utf-8 |
| 2 | # Copyright (c) 2024, Huawei Technologies Co., Ltd. All rights reserved. | 2 | # Copyright (c) 2024, Huawei Technologies Co., Ltd. All rights reserved. |
| 3 | 3 | ||
| 4 | +import logging | ||
| 4 | import math | 5 | import math |
| 5 | from functools import wraps | 6 | from functools import wraps |
| 6 | 7 | ||
| @@ -30,18 +31,22 @@ from mindspeed.core.parallel_state import (get_context_parallel_group_for_hybrid | |||
| 30 | from mindspeed.core.context_parallel.adaptive_context_parallel import adaptive_attn_context_parallel | 31 | from mindspeed.core.context_parallel.adaptive_context_parallel import adaptive_attn_context_parallel |
| 31 | from mindspeed.core.context_parallel.utils import get_scheduling_info | 32 | from mindspeed.core.context_parallel.utils import get_scheduling_info |
| 32 | from mindspeed.model.transformer import get_attention_mask | 33 | from mindspeed.model.transformer import get_attention_mask |
| 33 | -from mindspeed.utils import get_actual_seq_len | ||
| 34 | from mindspeed.core.context_parallel.context_parallel_kv_cache import get_cache_policy | 34 | from mindspeed.core.context_parallel.context_parallel_kv_cache import get_cache_policy |
| 35 | from mindspeed.utils import get_actual_seq_len, compute_qkv_index, get_position_ids | 35 | from mindspeed.utils import get_actual_seq_len, compute_qkv_index, get_position_ids |
| 36 | 36 | ||
| 37 | from mindspeed_llm.core.models.common.embeddings.rotary_pos_embedding import yarn_get_mscale | 37 | from mindspeed_llm.core.models.common.embeddings.rotary_pos_embedding import yarn_get_mscale |
| 38 | from mindspeed_llm.tasks.models.common.alibi import Alibi | 38 | from mindspeed_llm.tasks.models.common.alibi import Alibi |
| 39 | +from mindspeed_llm.training.utils import recompute_valid_actual_seq_len | ||
| 40 | + | ||
| 41 | +logger = logging.getLogger(__name__) | ||
| 39 | 42 | ||
| 40 | try: | 43 | try: |
| 41 | from einops import rearrange | 44 | from einops import rearrange |
| 42 | except ImportError: | 45 | except ImportError: |
| 43 | rearrange = None | 46 | rearrange = None |
| 44 | 47 | ||
| 48 | +ACTUAL_SEQ_LEN_THRESHOLD = 2048 | ||
| 49 | + | ||
| 45 | 50 | ||
| 46 | def do_ring_context_parallel(q, k, v, head_num, softmax_scale, attn_mask, dropout_p=0., pse=None, pse_type=None, packed_seq_params=None): | 51 | def do_ring_context_parallel(q, k, v, head_num, softmax_scale, attn_mask, dropout_p=0., pse=None, pse_type=None, packed_seq_params=None): |
| 47 | args = get_args() | 52 | args = get_args() |
| @@ -51,12 +56,12 @@ def do_ring_context_parallel(q, k, v, head_num, softmax_scale, attn_mask, dropou | |||
| 51 | cu_seqlens_q=torch.tensor(actual_seq_len, dtype=torch.int64, device=torch.cuda.current_device()), | 56 | cu_seqlens_q=torch.tensor(actual_seq_len, dtype=torch.int64, device=torch.cuda.current_device()), |
| 52 | cu_seqlens_kv=torch.tensor(actual_seq_len, dtype=torch.int64, device=torch.cuda.current_device()) | 57 | cu_seqlens_kv=torch.tensor(actual_seq_len, dtype=torch.int64, device=torch.cuda.current_device()) |
| 53 | ) | 58 | ) |
| 54 | - | 59 | + |
| 55 | q_index, kv_index = compute_qkv_index(torch.tensor(actual_seq_len, dtype=torch.int64, device=torch.cuda.current_device()).clone().tolist()) | 60 | q_index, kv_index = compute_qkv_index(torch.tensor(actual_seq_len, dtype=torch.int64, device=torch.cuda.current_device()).clone().tolist()) |
| 56 | packed_seq_params.q_index = q_index | 61 | packed_seq_params.q_index = q_index |
| 57 | packed_seq_params.kv_index = kv_index | 62 | packed_seq_params.kv_index = kv_index |
| 58 | packed_seq_params.position_ids = get_position_ids() | 63 | packed_seq_params.position_ids = get_position_ids() |
| 59 | - | 64 | + |
| 60 | in_hybrid_mode = get_context_parallel_group_for_hybrid_ring(check_initialized=False) is not None | 65 | in_hybrid_mode = get_context_parallel_group_for_hybrid_ring(check_initialized=False) is not None |
| 61 | if in_hybrid_mode: | 66 | if in_hybrid_mode: |
| 62 | cp_group = get_context_parallel_group_for_hybrid_ring() | 67 | cp_group = get_context_parallel_group_for_hybrid_ring() |
| @@ -105,14 +110,14 @@ def do_ulyssesattn_context_parallel(self, | |||
| 105 | attn_mask_type, | 110 | attn_mask_type, |
| 106 | packed_seq_params): | 111 | packed_seq_params): |
| 107 | args = get_args() | 112 | args = get_args() |
| 108 | - | 113 | + |
| 109 | sparse_mode = args.sparse_mode | 114 | sparse_mode = args.sparse_mode |
| 110 | if attn_mask_type == AttnMaskType.no_mask: | 115 | if attn_mask_type == AttnMaskType.no_mask: |
| 111 | sparse_mode = 0 # default mask | 116 | sparse_mode = 0 # default mask |
| 112 | - | 117 | + |
| 113 | scale = 1.0 / math.sqrt( | 118 | scale = 1.0 / math.sqrt( |
| 114 | self.hidden_size_per_attention_head) if self.scale_mask_softmax.scale is None else self.softmax_scale | 119 | self.hidden_size_per_attention_head) if self.scale_mask_softmax.scale is None else self.softmax_scale |
| 115 | - | 120 | + |
| 116 | self.ulysses_comm_para['cache_policy'] = get_cache_policy( | 121 | self.ulysses_comm_para['cache_policy'] = get_cache_policy( |
| 117 | self.layer_number, args.context_parallel_kv_cache_policy, args.context_parallel_cache_interval | 122 | self.layer_number, args.context_parallel_kv_cache_policy, args.context_parallel_cache_interval |
| 118 | ) | 123 | ) |
| @@ -126,10 +131,10 @@ def do_ulyssesattn_context_parallel(self, | |||
| 126 | attn_para['keep_prob'] = 1 - self.attention_dropout.p | 131 | attn_para['keep_prob'] = 1 - self.attention_dropout.p |
| 127 | attn_para['sparse_mode'] = sparse_mode | 132 | attn_para['sparse_mode'] = sparse_mode |
| 128 | output = ulyssesattn_context_parallel(query, key, value, attn_para, self.ulysses_comm_para) | 133 | output = ulyssesattn_context_parallel(query, key, value, attn_para, self.ulysses_comm_para) |
| 129 | - | 134 | + |
| 130 | return output | 135 | return output |
| 131 | - | 136 | + |
| 132 | - | 137 | + |
| 133 | def dot_product_attention_init( | 138 | def dot_product_attention_init( |
| 134 | self, | 139 | self, |
| 135 | config: TransformerConfig, | 140 | config: TransformerConfig, |
| @@ -489,6 +494,10 @@ def flash_attention_forward( | |||
| 489 | ) | 494 | ) |
| 490 | output = output.transpose(0, 1) | 495 | output = output.transpose(0, 1) |
| 491 | else: | 496 | else: |
| 497 | + if actual_seq_len is not None and len(actual_seq_len) > ACTUAL_SEQ_LEN_THRESHOLD: | ||
| 498 | + logger.warning("flash-attention get a long actual_seq_len, maybe create a coredump!") | ||
| 499 | + actual_seq_len = recompute_valid_actual_seq_len(get_position_ids(), actual_seq_len) | ||
| 500 | + | ||
| 492 | output = torch_npu.npu_fusion_attention( | 501 | output = torch_npu.npu_fusion_attention( |
| 493 | query, key, value, n_head, args.shape_order, | 502 | query, key, value, n_head, args.shape_order, |
| 494 | pse=pse, | 503 | pse=pse, |
| @@ -70,6 +70,15 @@ def compute_actual_seq_len(seq, stride=0): | |||
| 70 | return res | 70 | return res |
| 71 | 71 | ||
| 72 | 72 | ||
| 73 | +def recompute_valid_actual_seq_len(pos_ids, actual_seq_len): | ||
| 74 | + seq = pos_ids.view(-1) | ||
| 75 | + valid_seq = (seq != 0).nonzero()[-1] + 1 + 1 | ||
| 76 | + valid_actual_seq_len_clip = (torch.tensor(actual_seq_len).to(pos_ids.device) < valid_seq).nonzero()[-1] | ||
| 77 | + valid_actual_seq_len = actual_seq_len[:valid_actual_seq_len_clip + 1] | ||
| 78 | + valid_actual_seq_len.append(actual_seq_len[-1]) | ||
| 79 | + return valid_actual_seq_len | ||
| 80 | + | ||
| 81 | + | ||
| 73 | def generate_actual_seq_len(batch): | 82 | def generate_actual_seq_len(batch): |
| 74 | position_ids = batch.get('position_ids').transpose(0, 1).contiguous() | 83 | position_ids = batch.get('position_ids').transpose(0, 1).contiguous() |
| 75 | set_position_ids(position_ids) | 84 | set_position_ids(position_ids) |