已合并
【r1.9.0-beta1】【bugfix】增加softmax_out排布转换，以支持dsa一阶段tnd排布使用FA的tnd排布 #8288
wei_zhuoyi创建于 5月22日
【r1.9.0-beta1】【bugfix】增加softmax_out排布转换，以支持dsa一阶段tnd排布使用FA的tnd排布 #8288
已合并
共 2 个文件变更+135-37
| @@ -30,8 +30,8 @@ from mindformers.parallel_core.transformer_config import MLATransformerConfig | |||
| 30 | from mindformers.parallel_core.training_graph.transformer.dsa.dsa_indexer_loss import DSAIndexerLoss | 30 | from mindformers.parallel_core.training_graph.transformer.dsa.dsa_indexer_loss import DSAIndexerLoss |
| 31 | from mindformers.parallel_core.training_graph.device_matrix import layout | 31 | from mindformers.parallel_core.training_graph.device_matrix import layout |
| 32 | from mindformers.parallel_core.training_graph.communication import get_dp_cp_tp_id | 32 | from mindformers.parallel_core.training_graph.communication import get_dp_cp_tp_id |
| 33 | -from mindformers.parallel_core.training_graph.transformer.mask_generate import CausalEODMaskGenerate | ||
| 34 | from mindformers.parallel_core.training_graph.transformer.dsa.utils import adjust_bsnd_input, adjust_tnd_input | 33 | from mindformers.parallel_core.training_graph.transformer.dsa.utils import adjust_bsnd_input, adjust_tnd_input |
| 34 | +from mindformers.parallel_core.training_graph.transformer.softmax_convert import SoftmaxConverter | ||
| 35 | 35 | ||
| 36 | 36 | ||
| 37 | 37 | ||
| @@ -130,12 +130,11 @@ class DSAttention(nn.Cell): | |||
| 130 | head_num=self.head_num, | 130 | head_num=self.head_num, |
| 131 | scale_value=self.softmax_scale, | 131 | scale_value=self.softmax_scale, |
| 132 | inner_precise=0, | 132 | inner_precise=0, |
| 133 | - input_layout="BSND", | 133 | + input_layout=self.input_layout, |
| 134 | - sparse_mode=config.sparse_mode if not self.is_tnd else 0 | 134 | + sparse_mode=config.sparse_mode |
| 135 | ) | 135 | ) |
| 136 | - self.eod_mask_generator = CausalEODMaskGenerate(config) | 136 | + self.softmax_converter = SoftmaxConverter(config) |
| 137 | - self.softmax_transpose = aclnn_ops.Transpose() | 137 | + self.slice = aclnn_ops.StridedSlice().add_prim_attr("self_define_shard", True) |
| 138 | - self.slice = aclnn_ops.StridedSlice().add_prim_attr("self_define_shard", True) | ||
| 139 | if _get_parallel_mode() in (ParallelMode.SEMI_AUTO_PARALLEL,): | 138 | if _get_parallel_mode() in (ParallelMode.SEMI_AUTO_PARALLEL,): |
| 140 | self.shard() | 139 | self.shard() |
| 141 | 140 | ||
| @@ -196,17 +195,6 @@ class DSAttention(nn.Cell): | |||
| 196 | max_logits = self.reduce_max(softmax_max, (0, 1, 2)) | 195 | max_logits = self.reduce_max(softmax_max, (0, 1, 2)) |
| 197 | self.assign(self.max_logits_val, self.maximum(self.max_logits_val, max_logits)) | 196 | self.assign(self.max_logits_val, self.maximum(self.max_logits_val, max_logits)) |
| 198 | else: | 197 | else: |
| 199 | - # During dense warmup stage, common FlashAttention will be applied instead of SparseFlashAttention | ||
| 200 | - if self.is_tnd: | ||
| 201 | - t, n, _ = query.shape | ||
| 202 | - b, s = t // self.seq_length, self.seq_length | ||
| 203 | - attention_mask = self.eod_mask_generator(self.reshape(actual_seq_qlen, (b, -1))) | ||
| 204 | - query = self.reshape(query, (b, s, n, -1)) | ||
| 205 | - key = self.reshape(key, (b, s, n, -1)) | ||
| 206 | - value = self.reshape(value, (b, s, n, -1)) | ||
| 207 | - actual_seq_qlen = None | ||
| 208 | - actual_seq_kvlen = None | ||
| 209 | - | ||
| 210 | softmax_max, softmax_sum, _, attention_output = self.dense_flash_attention( | 198 | softmax_max, softmax_sum, _, attention_output = self.dense_flash_attention( |
| 211 | query, key, value, | 199 | query, key, value, |
| 212 | None, None, None, | 200 | None, None, None, |
| @@ -219,15 +207,14 @@ class DSAttention(nn.Cell): | |||
| 219 | step = (1,) * len(softmax_max.shape) | 207 | step = (1,) * len(softmax_max.shape) |
| 220 | softmax_max = self.slice(softmax_max, begin, (*softmax_shape, 1), step) | 208 | softmax_max = self.slice(softmax_max, begin, (*softmax_shape, 1), step) |
| 221 | softmax_sum = self.slice(softmax_sum, begin, (*softmax_shape, 1), step) | 209 | softmax_sum = self.slice(softmax_sum, begin, (*softmax_shape, 1), step) |
| 222 | - if self.track_max_attention_logit: | ||
| 223 | - max_logits = self.reduce_max(softmax_max, (0, 2, 3)) | ||
| 224 | - self.assign(self.max_logits_val, self.maximum(self.max_logits_val, max_logits)) | ||
| 225 | if self.is_tnd: | 210 | if self.is_tnd: |
| 226 | - # [b, s, n, d] -> [t, n, d] | 211 | + softmax_max, softmax_sum = self.softmax_converter(softmax_max, softmax_sum, actual_seq_qlen) |
| 227 | - attention_output = self.reshape(attention_output, (t, n, -1)) | 212 | + if self.track_max_attention_logit: |
| 228 | - # [b, n, s, 1] -> [n, b, s, 1] -> [n, t, 1] | 213 | + if self.is_tnd: |
| 229 | - softmax_max = self.reshape(self.softmax_transpose(softmax_max, (1, 0, 2, 3)), (n, t, 1)) | 214 | + max_logits = self.reduce_max(softmax_max, (1, 2)) |
| 230 | - softmax_sum = self.reshape(self.softmax_transpose(softmax_sum, (1, 0, 2, 3)), (n, t, 1)) | 215 | + else: |
| 216 | + max_logits = self.reduce_max(softmax_max, (0, 2, 3)) | ||
| 217 | + self.assign(self.max_logits_val, self.maximum(self.max_logits_val, max_logits)) | ||
| 231 | 218 | ||
| 232 | return attention_output, softmax_max, softmax_sum | 219 | return attention_output, softmax_max, softmax_sum |
| 233 | 220 | ||
| @@ -238,35 +225,42 @@ class DSAttention(nn.Cell): | |||
| 238 | q_shard = layout("dp_cp_tp", "None", "None") | 225 | q_shard = layout("dp_cp_tp", "None", "None") |
| 239 | kv_shard = layout("dp", "None", "None") | 226 | kv_shard = layout("dp", "None", "None") |
| 240 | idx_shard = layout("dp_cp_tp", "None", "None") | 227 | idx_shard = layout("dp_cp_tp", "None", "None") |
| 241 | - attn_shard = layout("dp_cp_tp", "None", "None") | ||
| 242 | - softmax_shard = layout("None", "dp_cp_tp", "None") | ||
| 243 | sfa_shard = (q_shard, kv_shard, kv_shard, idx_shard, q_shard, kv_shard, layout("dp"), layout("dp")) | 228 | sfa_shard = (q_shard, kv_shard, kv_shard, idx_shard, q_shard, kv_shard, layout("dp"), layout("dp")) |
| 229 | + fa_shard = ( | ||
| 230 | + layout("dp_cp", "tp", "None"), | ||
| 231 | + layout("dp", "tp", "None"), | ||
| 232 | + layout("dp", "tp", "None"), | ||
| 233 | + layout("None", "None"), | ||
| 234 | + layout("dp"), layout("dp") | ||
| 235 | + ) | ||
| 244 | else: | 236 | else: |
| 245 | q_shard = layout("dp", "cp_tp", "None", "None") | 237 | q_shard = layout("dp", "cp_tp", "None", "None") |
| 246 | kv_shard = layout("dp", "None", "None", "None") | 238 | kv_shard = layout("dp", "None", "None", "None") |
| 247 | idx_shard = layout("dp", "cp_tp", "None", "None") | 239 | idx_shard = layout("dp", "cp_tp", "None", "None") |
| 248 | - attn_shard = layout("dp", "cp_tp", "None", "None") | ||
| 249 | - softmax_shard = layout("dp", "None", "cp_tp", "None") | ||
| 250 | sfa_shard = (q_shard, kv_shard, kv_shard, idx_shard, q_shard, kv_shard) | 240 | sfa_shard = (q_shard, kv_shard, kv_shard, idx_shard, q_shard, kv_shard) |
| 251 | - self.split_q.shard((q_shard,)) | ||
| 252 | - self.split_k.shard((kv_shard,)) | ||
| 253 | - if self.sparse_loss: | ||
| 254 | - self.sparse_flash_attention.shard(sfa_shard, (attn_shard, softmax_shard, softmax_shard)) | ||
| 255 | - else: | ||
| 256 | fa_shard = ( | 241 | fa_shard = ( |
| 257 | layout("dp", "cp", "tp", "None"), | 242 | layout("dp", "cp", "tp", "None"), |
| 258 | layout("dp", "None", "tp", "None"), | 243 | layout("dp", "None", "tp", "None"), |
| 259 | layout("dp", "None", "tp", "None"), | 244 | layout("dp", "None", "tp", "None"), |
| 260 | layout("None", "None") if self.mask_compression else layout("dp", "None", "cp", "None"), | 245 | layout("None", "None") if self.mask_compression else layout("dp", "None", "cp", "None"), |
| 261 | ) | 246 | ) |
| 262 | - softmax_shard = layout("dp", "tp", "cp", "None") | 247 | + self.split_q.shard((q_shard,)) |
| 263 | - attn_shard = layout("dp", "cp", "tp", "None") | 248 | + self.split_k.shard((kv_shard,)) |
| 249 | + if self.sparse_loss: | ||
| 250 | + attn_shard = layout("dp_cp_tp", "None", "None") if self.is_tnd else layout("dp", "cp_tp", "None", "None") | ||
| 251 | + softmax_shard = layout("None", "dp_cp_tp", "None") if self.is_tnd else layout("dp", "None", "cp_tp", "None") | ||
| 252 | + self.sparse_flash_attention.shard(sfa_shard, (attn_shard, softmax_shard, softmax_shard)) | ||
| 253 | + else: | ||
| 254 | + softmax_shard = layout("dp_cp", "tp", "None") if self.is_tnd else layout("dp", "tp", "cp", "None") | ||
| 255 | + attn_shard = layout("dp_cp", "tp", "None") if self.is_tnd else layout("dp", "cp", "tp", "None") | ||
| 264 | self.dense_flash_attention.shard(fa_shard, (softmax_shard, softmax_shard, layout("None"), attn_shard)) | 256 | self.dense_flash_attention.shard(fa_shard, (softmax_shard, softmax_shard, layout("None"), attn_shard)) |
| 265 | self.slice.shard((softmax_shard,), (softmax_shard,)) | 257 | self.slice.shard((softmax_shard,), (softmax_shard,)) |
| 266 | - self.softmax_transpose.shard((layout("dp", "tp", "cp", "None"),)) | 258 | + |
| 267 | if self.track_max_attention_logit: | 259 | if self.track_max_attention_logit: |
| 268 | if self.sparse_loss: | 260 | if self.sparse_loss: |
| 269 | softmax_shard = layout("None", "dp_cp", "tp") if self.is_tnd else layout("dp", "None", "cp", "tp") | 261 | softmax_shard = layout("None", "dp_cp", "tp") if self.is_tnd else layout("dp", "None", "cp", "tp") |
| 262 | + else: | ||
| 263 | + softmax_shard = layout("tp", "dp_cp", "None") if self.is_tnd else layout("dp", "tp", "cp", "None") | ||
| 270 | self.assign.shard((layout("tp"), layout("tp")), (layout("tp"),)) | 264 | self.assign.shard((layout("tp"), layout("tp")), (layout("tp"),)) |
| 271 | self.maximum.shard((layout("tp"), layout("tp")), (layout("tp"),)) | 265 | self.maximum.shard((layout("tp"), layout("tp")), (layout("tp"),)) |
| 272 | self.reduce_max.shard((softmax_shard,), (layout("tp"),)) | 266 | self.reduce_max.shard((softmax_shard,), (layout("tp"),)) |
| @@ -0,0 +1,104 @@ | |||
| 1 | +# Copyright 2026 Huawei Technologies Co., Ltd | ||
| 2 | +# | ||
| 3 | +# Licensed under the Apache License, Version 2.0 (the "License"); | ||
| 4 | +# you may not use this file except in compliance with the License. | ||
| 5 | +# You may obtain a copy of the License at | ||
| 6 | +# | ||
| 7 | +# http://www.apache.org/licenses/LICENSE-2.0 | ||
| 8 | +# | ||
| 9 | +# Unless required by applicable law or agreed to in writing, software | ||
| 10 | +# distributed under the License is distributed on an "AS IS" BASIS, | ||
| 11 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| 12 | +# See the License for the specific language governing permissions and | ||
| 13 | +# limitations under the License. | ||
| 14 | +# ============================================================================ | ||
| 15 | +"""Convert Softmax_out Layout""" | ||
| 16 | +from mindspore import nn, ops | ||
| 17 | +from mindspore.ops import auto_generate as aclnn_ops | ||
| 18 | +from mindspore import dtype as mstype | ||
| 19 | + | ||
| 20 | +from mindformers.parallel_core.training_graph.communication import get_dp_cp_tp_id | ||
| 21 | +from mindformers.parallel_core.training_graph.device_matrix import layout | ||
| 22 | +from mindformers.parallel_core.transformer_config import TransformerConfig | ||
| 23 | + | ||
| 24 | + | ||
| 25 | +class SoftmaxConverter(nn.Cell): | ||
| 26 | + """convert softmax_max and softmax_sum layout of FlashAttentionScore in TND layout""" | ||
| 27 | + def __init__(self, config: TransformerConfig): | ||
| 28 | + super().__init__() | ||
| 29 | + if not config.input_layout != "TND": | ||
| 30 | + raise ValueError("Only softmax_out in TND layout need to be converted!") | ||
| 31 | + self.seq_length = config.seq_length | ||
| 32 | + self.head_num = config.num_attention_heads | ||
| 33 | + self.is_dryrun = config.is_dryrun | ||
| 34 | + self.tp = config.tensor_model_parallel_size | ||
| 35 | + self.cp = config.context_parallel_size | ||
| 36 | + self.cp_id = get_dp_cp_tp_id(config)[1] | ||
| 37 | + | ||
| 38 | + self.cast = aclnn_ops.Cast() | ||
| 39 | + self.reshape = aclnn_ops.Reshape() | ||
| 40 | + self.slice = aclnn_ops.StridedSlice().add_prim_attr("self_define_shard", True) | ||
| 41 | + self.clamp = aclnn_ops.ClampScalar() | ||
| 42 | + self.roll = aclnn_ops.Roll(1) | ||
| 43 | + self.repeat_scalar = aclnn_ops.RepeatInterleaveInt() | ||
| 44 | + self.arange = aclnn_ops.Arange() | ||
| 45 | + self.mod = aclnn_ops.FmodScalar() | ||
| 46 | + self.repeat_interleave = aclnn_ops.RepeatInterleaveTensor() | ||
| 47 | + self.equal = aclnn_ops.Equal() | ||
| 48 | + self.select = aclnn_ops.MaskedSelect() | ||
| 49 | + self.stack = aclnn_ops.StackExt() | ||
| 50 | + self.softmax_transpose = aclnn_ops.Transpose() | ||
| 51 | + infer_dtype = lambda *args: (args[0], args[1]) | ||
| 52 | + self.convert_softmax = ops.Morph( | ||
| 53 | + self._convert_softmax, | ||
| 54 | + self.infer_shape, infer_dtype | ||
| 55 | + ).add_prim_attr("self_define_shard", True) | ||
| 56 | + self.shard() | ||
| 57 | + | ||
| 58 | + def infer_shape(self, *args): | ||
| 59 | + t, n, _ = args[0] | ||
| 60 | + out_shape = [n, t, 1] | ||
| 61 | + return out_shape, out_shape | ||
| 62 | + | ||
| 63 | + def _convert_softmax(self, softmax_max, softmax_sum, actual_seq_len): | ||
| 64 | + """convert softmax layout""" | ||
| 65 | + t, n, _ = softmax_max.shape | ||
| 66 | + softmax_max = self.reshape(softmax_max, (-1,)) | ||
| 67 | + softmax_sum = self.reshape(softmax_sum, (-1,)) | ||
| 68 | + partial_num_head = self.head_num // self.tp | ||
| 69 | + offset_q = t * self.cp_id | ||
| 70 | + actual_seq_qlen = self.cast(self.clamp(actual_seq_len - offset_q, 0, t), mstype.int32) | ||
| 71 | + prev_seq_qlen = self.roll(actual_seq_qlen) | ||
| 72 | + prev_seq_qlen[0] = 0 | ||
| 73 | + interleave_seq_qlen = actual_seq_qlen - prev_seq_qlen | ||
| 74 | + interleave_seq_qlen = self.repeat_scalar(interleave_seq_qlen, partial_num_head) | ||
| 75 | + base = self.mod(self.arange(0, interleave_seq_qlen.shape[0], 1), partial_num_head) | ||
| 76 | + interleave_seq_qlen = self.repeat_interleave(base, interleave_seq_qlen) | ||
| 77 | + softmax_maxs = [] | ||
| 78 | + softmax_sums = [] | ||
| 79 | + for i in range(partial_num_head): | ||
| 80 | + head_mask = self.equal(interleave_seq_qlen, i) | ||
| 81 | + softmax_maxs.append(self.select(softmax_max, head_mask)) | ||
| 82 | + softmax_sums.append(self.select(softmax_sum, head_mask)) | ||
| 83 | + softmax_max = self.reshape(self.stack(softmax_maxs), (n, t, 1)) | ||
| 84 | + softmax_sum = self.reshape(self.stack(softmax_sums), (n, t, 1)) | ||
| 85 | + return softmax_max, softmax_sum | ||
| 86 | + | ||
| 87 | + def construct(self, softmax_max, softmax_sum, actual_seq_len): | ||
| 88 | + if self.is_dryrun: | ||
| 89 | + softmax_max = self.softmax_transpose(softmax_max, (1, 0, 2)) | ||
| 90 | + softmax_sum = self.softmax_transpose(softmax_sum, (1, 0, 2)) | ||
| 91 | + return softmax_max, softmax_sum | ||
| 92 | + if softmax_max.shape[-1] != 1: | ||
| 93 | + softmax_shape = softmax_max.shape[:-1] | ||
| 94 | + softmax_max = self.slice(softmax_max, (0, 0, 0), (*softmax_shape, 1), (1, 1, 1)) | ||
| 95 | + softmax_sum = self.slice(softmax_sum, (0, 0, 0), (*softmax_shape, 1), (1, 1, 1)) | ||
| 96 | + return self.convert_softmax(softmax_max, softmax_sum, actual_seq_len) | ||
| 97 | + | ||
| 98 | + def shard(self): | ||
| 99 | + self.slice.shard((layout("dp_cp", "tp", "None"),), (layout("dp_cp", "tp", "None"),)) | ||
| 100 | + self.convert_softmax.shard( | ||
| 101 | + (layout("dp_cp", "tp", "None"), layout("dp_cp", "tp", "None"), layout("dp")), | ||
| 102 | + (layout("tp", "dp_cp", "None"), layout("tp", "dp_cp", "None")) | ||
| 103 | + ) | ||
| 104 | + self.softmax_transpose.shard((layout("dp_cp", "tp", "None"),)) | ||