已合并
【r1.9.0-beta1】【bugfix】dsa消重排 #8196
wei_zhuoyi创建于 4月15日
【r1.9.0-beta1】【bugfix】dsa消重排 #8196
已合并
共 8 个文件变更+50-61
| @@ -737,21 +737,11 @@ class GPTModel(nn.Cell): | |||
| 737 | 737 | ||
| 738 | def shard(self, config: TransformerConfig): | 738 | def shard(self, config: TransformerConfig): |
| 739 | """parallel shard.""" | 739 | """parallel shard.""" |
| 740 | - dp = config.data_parallel_size | 740 | + self.concat_prefix.shard((layout("dp", "None", "cp", "None"), layout("dp", "None", "cp", "None"))) |
| 741 | - tp = config.tensor_model_parallel_size | 741 | + self.transpose.shard((layout("cp", "dp", "tp"),)) |
| 742 | - cp = 1 if config is None else config.context_parallel_size | ||
| 743 | - | ||
| 744 | - slice_in_strategy = ((dp, 1),) | ||
| 745 | - self.slice.shard(in_strategy=slice_in_strategy) | ||
| 746 | - not_equal_in_strategy = ((dp, 1), ()) | ||
| 747 | - self.not_equal.shard(in_strategy=not_equal_in_strategy) | ||
| 748 | - mul_in_strategy = ((dp, 1), (dp, 1)) | ||
| 749 | - self.mul.shard(in_strategy=mul_in_strategy) | ||
| 750 | - self.concat_prefix.shard(((dp, 1, cp, 1), (dp, 1, cp, 1))) | ||
| 751 | - self.transpose.shard(((cp, dp, tp),)) | ||
| 752 | 742 | ||
| 753 | if config.moe_router_load_balancing_type == "gbs_aux_loss": | 743 | if config.moe_router_load_balancing_type == "gbs_aux_loss": |
| 754 | - self.assign_aux.shard(((1,), (1,))) | 744 | + self.assign_aux.shard((layout("None"), layout("None"))) |
| 755 | 745 | ||
| 756 | pipeline_stage = config.pipeline_model_parallel_size | 746 | pipeline_stage = config.pipeline_model_parallel_size |
| 757 | if pipeline_stage > 1: | 747 | if pipeline_stage > 1: |
| @@ -127,6 +127,7 @@ class DSAttention(nn.Cell): | |||
| 127 | input_layout=self.input_layout, | 127 | input_layout=self.input_layout, |
| 128 | sparse_mode=config.sparse_mode | 128 | sparse_mode=config.sparse_mode |
| 129 | ) | 129 | ) |
| 130 | + self.slice = aclnn_ops.StridedSlice().add_prim_attr("self_define_shard", True) | ||
| 130 | if _get_parallel_mode() in (ParallelMode.SEMI_AUTO_PARALLEL,): | 131 | if _get_parallel_mode() in (ParallelMode.SEMI_AUTO_PARALLEL,): |
| 131 | self.shard() | 132 | self.shard() |
| 132 | 133 | ||
| @@ -195,8 +196,11 @@ class DSAttention(nn.Cell): | |||
| 195 | actual_seq_qlen=actual_seq_qlen, | 196 | actual_seq_qlen=actual_seq_qlen, |
| 196 | actual_seq_kvlen=actual_seq_kvlen | 197 | actual_seq_kvlen=actual_seq_kvlen |
| 197 | ) | 198 | ) |
| 198 | - softmax_max = softmax_max[..., :1] | 199 | + softmax_shape = softmax_max.shape[:-1] |
| 199 | - softmax_sum = softmax_sum[..., :1] | 200 | + begin = (0,) * len(softmax_max.shape) |
| 201 | + step = (1,) * len(softmax_max.shape) | ||
| 202 | + softmax_max = self.slice(softmax_max, begin, (*softmax_shape, 1), step) | ||
| 203 | + softmax_sum = self.slice(softmax_sum, begin, (*softmax_shape, 1), step) | ||
| 200 | if self.track_max_attention_logit: | 204 | if self.track_max_attention_logit: |
| 201 | if self.is_tnd: | 205 | if self.is_tnd: |
| 202 | max_logits = self.reduce_max(softmax_max, (0, 2)) | 206 | max_logits = self.reduce_max(softmax_max, (0, 2)) |
| @@ -216,7 +220,7 @@ class DSAttention(nn.Cell): | |||
| 216 | idx_shard = layout("dp_cp", "None", "None") | 220 | idx_shard = layout("dp_cp", "None", "None") |
| 217 | mask_shard = layout("None", "None") | 221 | mask_shard = layout("None", "None") |
| 218 | attn_shard = layout("dp_cp", "tp", "None") | 222 | attn_shard = layout("dp_cp", "tp", "None") |
| 219 | - softmax_shard = layout("None", "dp_cp", "tp") if self.sparse_loss else layout("dp_cp", "None", "tp") | 223 | + softmax_shard = layout("None", "dp_cp", "tp") if self.sparse_loss else layout("dp_cp", "tp", "None") |
| 220 | sfa_shard = (q_shard, kv_shard, kv_shard, idx_shard, q_shard, kv_shard, layout("dp"), layout("dp")) | 224 | sfa_shard = (q_shard, kv_shard, kv_shard, idx_shard, q_shard, kv_shard, layout("dp"), layout("dp")) |
| 221 | fa_shard = (q_shard, kv_shard, kv_shard, mask_shard, layout("dp"), layout("dp")) | 225 | fa_shard = (q_shard, kv_shard, kv_shard, mask_shard, layout("dp"), layout("dp")) |
| 222 | else: | 226 | else: |
| @@ -234,6 +238,7 @@ class DSAttention(nn.Cell): | |||
| 234 | self.sparse_flash_attention.shard(sfa_shard, (attn_shard, softmax_shard, softmax_shard)) | 238 | self.sparse_flash_attention.shard(sfa_shard, (attn_shard, softmax_shard, softmax_shard)) |
| 235 | else: | 239 | else: |
| 236 | self.dense_flash_attention.shard(fa_shard, (softmax_shard, softmax_shard, layout("None"), attn_shard)) | 240 | self.dense_flash_attention.shard(fa_shard, (softmax_shard, softmax_shard, layout("None"), attn_shard)) |
| 241 | + self.slice.shard((softmax_shard,), (softmax_shard,)) | ||
| 237 | if self.track_max_attention_logit: | 242 | if self.track_max_attention_logit: |
| 238 | self.assign.shard((layout("tp"), layout("tp")), (layout("tp"),)) | 243 | self.assign.shard((layout("tp"), layout("tp")), (layout("tp"),)) |
| 239 | self.maximum.shard((layout("tp"), layout("tp")), (layout("tp"),)) | 244 | self.maximum.shard((layout("tp"), layout("tp")), (layout("tp"),)) |
| @@ -124,6 +124,7 @@ class DSAIndexer(nn.Cell): | |||
| 124 | self.index_topk = config.dsa_indexer_topk | 124 | self.index_topk = config.dsa_indexer_topk |
| 125 | self.sparse_loss = config.dsa_indexer_use_sparse_loss | 125 | self.sparse_loss = config.dsa_indexer_use_sparse_loss |
| 126 | self.softmax_scale: float = self.index_head_dim ** -0.5 | 126 | self.softmax_scale: float = self.index_head_dim ** -0.5 |
| 127 | + self.weight_scale = Tensor([(self.index_n_heads ** -0.5) * self.softmax_scale], mstype.float32) | ||
| 127 | 128 | ||
| 128 | self.apply_rotary_emb_q = ApplyRotaryPosEmb(config) | 129 | self.apply_rotary_emb_q = ApplyRotaryPosEmb(config) |
| 129 | self.apply_rotary_emb_k = ApplyRotaryPosEmb(config, for_k_pos_emb=True) | 130 | self.apply_rotary_emb_k = ApplyRotaryPosEmb(config, for_k_pos_emb=True) |
| @@ -179,6 +180,7 @@ class DSAIndexer(nn.Cell): | |||
| 179 | self.cat_k = aclnn_ops.Concat(-1).add_prim_attr("self_define_shard", True) | 180 | self.cat_k = aclnn_ops.Concat(-1).add_prim_attr("self_define_shard", True) |
| 180 | self.cast = aclnn_ops.Cast() | 181 | self.cast = aclnn_ops.Cast() |
| 181 | self.reshape = aclnn_ops.Reshape() | 182 | self.reshape = aclnn_ops.Reshape() |
| 183 | + self.mul_scalar = aclnn_ops.Mul() | ||
| 182 | self.transpose_q = aclnn_ops.Transpose() | 184 | self.transpose_q = aclnn_ops.Transpose() |
| 183 | self.transpose_k = aclnn_ops.Transpose() | 185 | self.transpose_k = aclnn_ops.Transpose() |
| 184 | self.transpose_w = aclnn_ops.Transpose() | 186 | self.transpose_w = aclnn_ops.Transpose() |
| @@ -247,7 +249,7 @@ class DSAIndexer(nn.Cell): | |||
| 247 | k, _ = self.linear_wk(x) | 249 | k, _ = self.linear_wk(x) |
| 248 | k = self.k_norm(k) | 250 | k = self.k_norm(k) |
| 249 | # [seqlen, batch, index_head_dim] -> [seqlen, batch, 1, index_head_dim] | 251 | # [seqlen, batch, index_head_dim] -> [seqlen, batch, 1, index_head_dim] |
| 250 | - k = k.reshape(seqlen, bsz, 1, self.index_head_dim) | 252 | + k = self.reshape(k, (seqlen, bsz, 1, self.index_head_dim)) |
| 251 | # [seqlen, batch, 1, index_head_dim] -> [seqlen, batch, index_head_dim] | 253 | # [seqlen, batch, 1, index_head_dim] -> [seqlen, batch, index_head_dim] |
| 252 | k_nope, k_pe = self.split_k( | 254 | k_nope, k_pe = self.split_k( |
| 253 | k, [self.index_head_dim - self.qk_pos_emb_head_dim, self.qk_pos_emb_head_dim], dim=-1 | 255 | k, [self.index_head_dim - self.qk_pos_emb_head_dim, self.qk_pos_emb_head_dim], dim=-1 |
| @@ -272,7 +274,7 @@ class DSAIndexer(nn.Cell): | |||
| 272 | # ========================================= | 274 | # ========================================= |
| 273 | # [seqlen, batch, hidden_size] -> [seqlen, batch, index_n_heads] | 275 | # [seqlen, batch, hidden_size] -> [seqlen, batch, index_n_heads] |
| 274 | weights, _ = self.linear_weights_proj(x) | 276 | weights, _ = self.linear_weights_proj(x) |
| 275 | - weights = weights * (self.index_n_heads ** -0.5) * self.softmax_scale | 277 | + weights = self.mul_scalar(weights, self.weight_scale) |
| 276 | 278 | ||
| 277 | # convert q, k, weights from SB* layout to BS* or T* layout | 279 | # convert q, k, weights from SB* layout to BS* or T* layout |
| 278 | weights = self.transpose_w(weights, (1, 0, 2)) | 280 | weights = self.transpose_w(weights, (1, 0, 2)) |
| @@ -301,6 +303,9 @@ class DSAIndexer(nn.Cell): | |||
| 301 | return_value=True | 303 | return_value=True |
| 302 | ) | 304 | ) |
| 303 | if not self.sparse_loss: | 305 | if not self.sparse_loss: |
| 306 | + if self.is_tnd: | ||
| 307 | + actual_seq_qlen = self.cast(actual_seq_qlen, mstype.int64) | ||
| 308 | + actual_seq_klen = self.cast(actual_seq_klen, mstype.int64) | ||
| 304 | softmax_max_index, softmax_sum_index = self.dense_softmax_lse( | 309 | softmax_max_index, softmax_sum_index = self.dense_softmax_lse( |
| 305 | q, k, weights, | 310 | q, k, weights, |
| 306 | actual_seq_qlen=actual_seq_qlen, | 311 | actual_seq_qlen=actual_seq_qlen, |
| @@ -343,6 +348,7 @@ class DSAIndexer(nn.Cell): | |||
| 343 | w_shard = layout("cp", "dp", "None") | 348 | w_shard = layout("cp", "dp", "None") |
| 344 | self.split_q.shard((q_shard,)) | 349 | self.split_q.shard((q_shard,)) |
| 345 | self.split_k.shard((k_shard,)) | 350 | self.split_k.shard((k_shard,)) |
| 351 | + self.mul_scalar.shard((w_shard, layout("None"))) | ||
| 346 | self.cat_q.shard(((q_shard, q_shard),), (q_shard,)) | 352 | self.cat_q.shard(((q_shard, q_shard),), (q_shard,)) |
| 347 | self.cat_k.shard(((k_shard, k_shard),), (k_shard,)) | 353 | self.cat_k.shard(((k_shard, k_shard),), (k_shard,)) |
| 348 | # merge tp in q-head and cp in k-seq | 354 | # merge tp in q-head and cp in k-seq |
| @@ -28,7 +28,7 @@ def adjust_tnd_input(offset_id, q, k, actual_seq_qlen, actual_seq_klen): | |||
| 28 | prev_seq_klen[0] = 0 | 28 | prev_seq_klen[0] = 0 |
| 29 | new_actual_seq_klen = F.cumsum(F.relu(new_actual_seq_klen - prev_seq_klen), 0) | 29 | new_actual_seq_klen = F.cumsum(F.relu(new_actual_seq_klen - prev_seq_klen), 0) |
| 30 | new_actual_seq_klen[-1] = slice_tk | 30 | new_actual_seq_klen[-1] = slice_tk |
| 31 | - return new_actual_seq_klen, new_actual_seq_klen | 31 | + return new_actual_seq_qlen, new_actual_seq_klen |
| 32 | 32 | ||
| 33 | 33 | ||
| 34 | def adjust_bsnd_input(offset_id, q, x): | 34 | def adjust_bsnd_input(offset_id, q, x): |
| @@ -183,21 +183,18 @@ class FusedLayerNorm(nn.Cell): | |||
| 183 | output = self.cast(output, original_type) | 183 | output = self.cast(output, original_type) |
| 184 | return output | 184 | return output |
| 185 | 185 | ||
| 186 | + # pylint: disable=W0613 | ||
| 186 | def shard(self, config: TransformerConfig, in_strategy=None): | 187 | def shard(self, config: TransformerConfig, in_strategy=None): |
| 187 | """shard method""" | 188 | """shard method""" |
| 188 | if in_strategy: | 189 | if in_strategy: |
| 189 | - if isinstance(in_strategy[0], Layout): | 190 | + strategy = in_strategy[0] |
| 190 | - strategy = layout.to_tuple_strategy(in_strategy[0]) | ||
| 191 | - else: | ||
| 192 | - strategy = in_strategy | ||
| 193 | else: | 191 | else: |
| 194 | - dp, _, cp = get_strategy(config) | 192 | + strategy = layout("cp", "dp", "None") |
| 195 | - strategy = (cp, dp, 1) | ||
| 196 | 193 | ||
| 197 | - if strategy[-1] != 1: | 194 | + if layout.to_tuple_strategy(strategy)[-1] != 1: |
| 198 | raise TypeError(f'The last dim in FusedLayerNorm can not equal to 1! Strategy {strategy} not supported!') | 195 | raise TypeError(f'The last dim in FusedLayerNorm can not equal to 1! Strategy {strategy} not supported!') |
| 199 | 196 | ||
| 200 | - self.layer_norm.shard((strategy, (strategy[-1],), (strategy[-1],))) | 197 | + self.layer_norm.shard((strategy, layout("None"), layout("None"))) |
| 201 | 198 | ||
| 202 | def sharding_propagation(self, config: TransformerConfig): | 199 | def sharding_propagation(self, config: TransformerConfig): |
| 203 | pass | 200 | pass |
| @@ -350,6 +347,7 @@ class FusedRMSNorm(nn.Cell): | |||
| 350 | output = self.cast(output, original_type) | 347 | output = self.cast(output, original_type) |
| 351 | return output | 348 | return output |
| 352 | 349 | ||
| 350 | + # pylint: disable=W0613 | ||
| 353 | def shard(self, config: TransformerConfig, in_strategy: tuple = None): | 351 | def shard(self, config: TransformerConfig, in_strategy: tuple = None): |
| 354 | """shard method""" | 352 | """shard method""" |
| 355 | if in_strategy: | 353 | if in_strategy: |