From ba6d1d9418dbe2deaa5d16ef6059311d774a8a9f Mon Sep 17 00:00:00 2001 From: yiyison Date: Wed, 28 May 2025 19:44:39 +0800 Subject: [PATCH] =?UTF-8?q?=E5=A2=9E=E5=8A=A0=20MLA=20=E5=AE=9E=E7=8E=B0?= =?UTF-8?q?=E4=BB=A5=E5=8F=8A=E6=B5=8B=E8=AF=95=E7=94=A8=E4=BE=8B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../transformer/multi_latent_attention.py | 839 ++++++++++++++++++ .../transformer/rotary_pos_embedding.py | 110 ++- .../test_multi_latent_attention/__init__.py | 14 + .../data_gen_utils.py | 306 +++++++ .../run_multi_latent_attention.py | 166 ++++ .../test_multi_latent_attention.py | 309 +++++++ 6 files changed, 1729 insertions(+), 15 deletions(-) create mode 100644 mindformers/parallel_core/training_graph/transformer/multi_latent_attention.py create mode 100644 tests/st/test_ut/test_parallel_core/test_training_graph/test_transformer/test_multi_latent_attention/__init__.py create mode 100644 tests/st/test_ut/test_parallel_core/test_training_graph/test_transformer/test_multi_latent_attention/data_gen_utils.py create mode 100644 tests/st/test_ut/test_parallel_core/test_training_graph/test_transformer/test_multi_latent_attention/run_multi_latent_attention.py create mode 100644 tests/st/test_ut/test_parallel_core/test_training_graph/test_transformer/test_multi_latent_attention/test_multi_latent_attention.py diff --git a/mindformers/parallel_core/training_graph/transformer/multi_latent_attention.py b/mindformers/parallel_core/training_graph/transformer/multi_latent_attention.py new file mode 100644 index 000000000..d0893ebe8 --- /dev/null +++ b/mindformers/parallel_core/training_graph/transformer/multi_latent_attention.py @@ -0,0 +1,839 @@ +# Copyright 2025 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ +"""Multi-head Latent Attention (MLA) mechanism with low-rank compression.""" +import math +from dataclasses import dataclass +from typing import Union +import mindspore as ms +from mindspore import nn, Tensor +from mindspore.mint import unsqueeze +from mindspore.ops import cast +from mindspore.common.initializer import initializer +from mindspore.context import ParallelMode +from mindspore.parallel._utils import _get_parallel_mode, _is_sharding_propagation + +from mindformers.parallel_core.utils.spec_utils import ModuleSpec, build_module +from mindformers.parallel_core.transformer_config import MLATransformerConfig +from mindformers.parallel_core.training_graph.transformer.rotary_pos_embedding import ApplyRotaryPosEmb, ROPE_FUNCTIONS + + +@dataclass +class MLASelfAttentionSubmodules: + """Submodules for the MLA self-attention layer.""" + linear_qkv: Union[ModuleSpec, type] = None + linear_qb: Union[ModuleSpec, type] = None + linear_kvb: Union[ModuleSpec, type] = None + core_attention: Union[ModuleSpec, type] = None + linear_proj: Union[ModuleSpec, type] = None + q_layernorm: Union[ModuleSpec, type] = None + k_layernorm: Union[ModuleSpec, type] = None + + +@dataclass +class MLASelfAttentionSubmodulesMegatron: + """Submodules for the MLA self-attention layer.""" + + linear_q_proj: Union[ModuleSpec, type] = None + linear_q_down_proj: Union[ModuleSpec, type] = None + linear_q_up_proj: Union[ModuleSpec, type] = None + linear_kv_down_proj: Union[ModuleSpec, type] = None + linear_kv_up_proj: Union[ModuleSpec, type] = None + core_attention: Union[ModuleSpec, type] = None + linear_proj: Union[ModuleSpec, type] = None + q_layernorm: Union[ModuleSpec, type] = None + kv_layernorm: Union[ModuleSpec, type] = None + + +def yarn_get_mscale(scale: float = 1, mscale: float = 1) -> float: + if scale <= 1: + return 1.0 + return 0.1 * mscale * math.log(scale) + 1.0 + + +class MultiLatentAttention(nn.Cell): + """Multi-head Latent Attention (MLA) with KV compression and rotary position encoding.""" + def __init__( + self, + config: MLATransformerConfig, + submodules: Union[MLASelfAttentionSubmodules, MLASelfAttentionSubmodulesMegatron], + layer_number: int, + attention_type: str, + attn_mask_type: str = None, + cp_comm_type: str = None, + ) -> None: + super().__init__() + if attn_mask_type: + raise NotImplementedError("For Attention, 'attn_mask_type' is not supported for now.") + if cp_comm_type: + raise NotImplementedError("cp_comm_type is not supported for now.") + self.config = config + self.layer_number = layer_number + self.layer_index = max(1, layer_number) + self.attn_mask_type = attn_mask_type + self.attention_type = attention_type + + # model structure config + self.use_flash_attention = self.config.use_flash_attention + self.use_ring_attention = self.config.use_ring_attention + self.use_eod_attn_mask_compression = self.config.use_eod_attn_mask_compression + self.use_attn_mask_compression = self.config.use_attn_mask_compression + self.num_attention_heads = self.config.num_attention_heads + self.query_projection_size = self.config.v_head_dim * self.config.num_attention_heads + self.qk_head_dim = self.config.qk_head_dim + self.qk_pos_emb_head_dim = self.config.qk_pos_emb_head_dim + self.q_head_dim = self.qk_head_dim + self.qk_pos_emb_head_dim + self.kv_lora_rank = self.config.kv_lora_rank + self.v_head_dim = self.config.v_head_dim + self.input_layout = self.config.input_layout + self.compute_dtype = self.config.compute_dtype + + zero_pad_length = self.q_head_dim - self.v_head_dim + if zero_pad_length == 0: + self.use_zero_pad = False + elif zero_pad_length < 0: + raise ValueError("qk_head_dim + qk_pos_emb_head_dim should not less than v_head_dim") + else: + self.use_zero_pad = True + self.pad_zeros = initializer('zeros', shape=(1, 1, self.num_attention_heads, zero_pad_length), + dtype=self.compute_dtype) + + mscale = yarn_get_mscale(self.config.rotary_scaling_factor, self.config.mscale) + self.softmax_scale = mscale * mscale / math.sqrt(self.q_head_dim) + self.rotary_pos_emb = ROPE_FUNCTIONS[self.config.rope_type](self.config) + + self.core_attention = build_module( + submodules.core_attention, + config=self.config, + layer_number=self.layer_index + ) + + # Output. + self.linear_proj = build_module( + submodules.linear_proj, + self.query_projection_size, + self.config.hidden_size, + config=self.config, + init_method=self.config.output_layer_init_method, + bias=self.config.add_bias_linear, + skip_bias_add=True, + ) + + self.shape = ms.ops.auto_generate.Shape() + self.reshape = ms.ops.auto_generate.Reshape() + self.merge_head_transpose = ms.ops.auto_generate.Transpose() + self.tile = ms.ops.auto_generate.Tile() + self.value_concat = ms.ops.auto_generate.Concat(-1) + self.value_tnd_concat = ms.ops.auto_generate.Concat(-1) + self.dim_slice = ms.ops.auto_generate.StridedSlice() + + if _get_parallel_mode() in (ParallelMode.AUTO_PARALLEL,) and _is_sharding_propagation(): + self.sharding_propagation() + elif _get_parallel_mode() in (ParallelMode.SEMI_AUTO_PARALLEL,): + self.shard() + + def sharding_propagation(self): + pass + + def shard(self): + dp = self.config.data_parallel_size + tp = self.config.tensor_model_parallel_size + + self.value_concat.shard(((dp, tp, 1, 1), (dp, tp, 1, 1))) + self.value_tnd_concat.shard(((dp, 1, tp, 1), (dp, 1, tp, 1))) + self.dim_slice.shard(((dp, tp, 1, 1),)) + self.tile.shard(((1, 1, 1, 1),)) + + def construct(self, x: Tensor, attention_mask=None, rotary_pos_emb=None, rotary_pos_cos=None, + rotary_pos_sin=None, prefix_keys_values=None, pad_zeros=None, actual_seq_len=None): + """Forward process.""" + if rotary_pos_cos: + raise NotImplementedError("rotary_pos_cos is not supported for now.") + if rotary_pos_sin: + raise NotImplementedError("rotary_pos_sin is not supported for now.") + if prefix_keys_values: + raise NotImplementedError("prefix_keys_values is not supported for now.") + ori_dtype = x.dtype + seq_len, bs, _ = self.shape(x) + + if self.use_eod_attn_mask_compression: + query, key, value = self.get_tnd_query_key_value_tensors(x, rotary_pos_emb=rotary_pos_emb) + else: + query, key, value = self.get_query_key_value_tensors(x, rotary_pos_emb=rotary_pos_emb) + + query = cast(query, self.compute_dtype) + key = cast(key, self.compute_dtype) + value = cast(value, self.compute_dtype) + if self.use_flash_attention: + if self.use_zero_pad: + if self.input_layout == "TND": + pad_zeros = self.tile(self.pad_zeros, (seq_len, bs, 1, 1)) + pad_zeros = self.reshape(pad_zeros, (seq_len * bs, -1, pad_zeros.shape[-1])) + value = self.value_tnd_concat((value, pad_zeros)) + else: + pad_zeros = self.tile(self.pad_zeros, (seq_len, bs, 1, 1)) + value = self.value_concat((value, pad_zeros)) + if self.use_eod_attn_mask_compression: + context_layer = self.core_attention( + query, key, value, attention_mask, + actual_seq_qlen=actual_seq_len, actual_seq_kvlen=actual_seq_len + ) + if self.use_zero_pad: + context_layer = self.reshape(context_layer, (seq_len, bs, self.num_attention_heads, -1)) + context_layer = self.dim_slice(context_layer, (0, 0, 0, 0), + (seq_len, bs, self.num_attention_heads, self.v_head_dim), + (1, 1, 1, 1)) + attn_out = self.reshape(context_layer, (seq_len, bs, -1)) + else: + context_layer = self.core_attention( + query, key, value, attention_mask, + ) + if self.use_zero_pad: + context_layer = self.reshape(context_layer, (seq_len, bs, self.num_attention_heads, -1)) + context_layer = self.dim_slice(context_layer, (0, 0, 0, 0), + (seq_len, bs, self.num_attention_heads, self.v_head_dim), + (1, 1, 1, 1)) + attn_out = self.reshape(context_layer, (seq_len, bs, -1)) + else: + attn_out = self.core_attention(query, key, value, attention_mask) + + output = self.linear_proj(attn_out)[0] # dp, mp -> dp, 1 / dp * mp, 1 + output = cast(output, ori_dtype) + return output + + +class MLASelfAttention(MultiLatentAttention): + """MLA Self-attention layer class, the implementation follows the same structure as Mindspeed A2 + + Self-attention layer takes input with size [s, b, h] + and returns output of the same size. + """ + + def __init__( + self, + config: MLATransformerConfig, + submodules: MLASelfAttentionSubmodules, + layer_number: int, + attn_mask_type: str = None, + cp_comm_type: str = None, + ): + super().__init__( + config=config, + submodules=submodules, + layer_number=layer_number, + attn_mask_type=attn_mask_type, + cp_comm_type=cp_comm_type, + attention_type="self", + ) + + self.split = ms.ops.auto_generate.SplitWithSize().add_prim_attr("skip_redistribution", True) + self.tile_kv = ms.ops.auto_generate.Tile() + self.tile_tnd_kv = ms.ops.auto_generate.Tile() + self.pe_concat = ms.ops.auto_generate.Concat(axis=3) + self.pe_tnd_concat = ms.ops.auto_generate.Concat(axis=3) + self.apply_rotary_emb = ApplyRotaryPosEmb(config) + self.reshape = ms.ops.auto_generate.Reshape() + + if self.config.q_lora_rank is None: + self.q_rank = self.config.num_attention_heads * self.q_head_dim + self.q_layernorm = None + else: + self.q_rank = self.config.q_lora_rank + if submodules.q_layernorm is not None: + self.q_layernorm = build_module( + submodules.q_layernorm, + dim=self.config.q_lora_rank, + config=self.config, + eps=self.config.layernorm_epsilon, + ) + else: + self.q_layernorm = None + + self.linear_qb = build_module( + submodules.linear_qb, + self.config.q_lora_rank, + self.config.num_attention_heads * self.q_head_dim, + config=self.config, + init_method=self.config.init_method, + bias=self.config.add_bias_linear or self.config.add_qkv_bias, + skip_bias_add=False, + ) + + self.linear_qkv = build_module( + submodules.linear_qkv, + self.config.hidden_size, + self.q_rank + self.kv_lora_rank + self.qk_pos_emb_head_dim, + config=self.config, + init_method=self.config.init_method, + bias=self.config.add_bias_linear or self.config.add_qkv_bias, + skip_bias_add=False, + ) + + if submodules.k_layernorm is not None: + self.k_layernorm = build_module( + submodules.k_layernorm, + dim=self.kv_lora_rank, + config=self.config, + eps=self.config.layernorm_epsilon, + ) + else: + self.k_layernorm = None + + self.linear_kvb = build_module( + submodules.linear_kvb, + self.kv_lora_rank, + self.config.num_attention_heads * (self.q_head_dim - self.qk_pos_emb_head_dim + self.v_head_dim), + config=self.config, + init_method=self.config.init_method, + bias=self.config.add_bias_linear or self.config.add_qkv_bias, + skip_bias_add=False, + ) + + self.linear_proj = build_module( + submodules.linear_proj, + self.query_projection_size, + self.config.hidden_size, + config=self.config, + init_method=self.config.output_layer_init_method, + bias=self.config.add_bias_linear, + skip_bias_add=True, + ) + + if _get_parallel_mode() in (ParallelMode.AUTO_PARALLEL,) and _is_sharding_propagation(): + self.sharding_propagation() + elif _get_parallel_mode() in (ParallelMode.SEMI_AUTO_PARALLEL,): + self.shard_self_attn() + + def sharding_propagation(self): + dp = self.config.data_parallel_size + tp = self.config.tensor_model_parallel_size + + self.tile_kv.shard(((1, dp, tp, 1),)) + self.tile_tnd_kv.shard(((1, dp, tp, 1),)) + + def shard_self_attn(self): + dp = self.config.data_parallel_size + tp = self.config.tensor_model_parallel_size + + self.tile_kv.shard(((1, dp, tp, 1),)) + self.tile_tnd_kv.shard(((1, dp, tp, 1),)) + self.pe_concat.shard(((1, dp, tp, 1), (1, dp, tp, 1))) + self.pe_tnd_concat.shard(((1, dp, tp, 1), (1, dp, tp, 1))) + + + def get_query_key_value_tensors(self, + hidden_states, + key_value_states=None, + position_ids=None, + packed_seq_params=None, + inference_params=None, + rotary_pos_emb=None): + """get_query_key_value_tensors""" + if key_value_states: + raise NotImplementedError("`key_value_states` is not supported for now") + if position_ids: + raise NotImplementedError("`position_ids` is not supported for now") + if packed_seq_params: + raise NotImplementedError("`packed_seq_params` is not supported for now") + if inference_params: + raise NotImplementedError("`inference_params` is not supported for now") + + tmp = self.shape(hidden_states) + seq_len, bs = tmp[0], tmp[1] + + seq_len = self.shape(hidden_states)[0] + rotary_pos_emb = self.rotary_pos_emb(seq_len) + + qkv_combo = self.linear_qkv(hidden_states)[0] + + q_a, compressed_kv, k_pe = self.split( + qkv_combo, + [ + self.q_rank, + self.kv_lora_rank, + self.qk_pos_emb_head_dim, + ], + dim=-1, + ) + + if self.q_layernorm is not None: + q_a = self.q_layernorm(q_a) + q = self.linear_qb(q_a)[0] + q = self.reshape(q, (seq_len, bs, self.num_attention_heads, -1)) + q_nope, q_pe = self.split( + q, [self.qk_head_dim, self.qk_pos_emb_head_dim], dim=-1 + ) + + else: + q = self.reshape(q_a, (seq_len, bs, self.num_attention_heads, -1)) + q_nope, q_pe = self.split( + q, [self.qk_head_dim, self.qk_pos_emb_head_dim], dim=-1 + ) + + k_pe = self.reshape(k_pe, (seq_len, bs, 1, self.qk_pos_emb_head_dim)) + compressed_kv_norm = self.k_layernorm(compressed_kv) + + kv = self.linear_kvb(compressed_kv_norm)[0] + kv = self.reshape(kv, ( + seq_len, + bs, + self.num_attention_heads, + self.qk_head_dim + self.v_head_dim, + )) + + k_nope, value = self.split(kv, [self.qk_head_dim, self.v_head_dim], dim=-1) + + if rotary_pos_emb is not None: + q_pe = self.apply_rotary_emb( + q_pe, + rotary_pos_emb, + rotary_interleaved=self.config.rotary_interleaved, + multi_latent_attention=self.config.multi_latent_attention + ) + k_pe = self.apply_rotary_emb( + k_pe, + rotary_pos_emb, + rotary_interleaved=self.config.rotary_interleaved, + multi_latent_attention=self.config.multi_latent_attention, + input_is_parallel=True + ) + + query = self.pe_concat([q_nope, q_pe]) + k_pe = self.tile_kv(k_pe, (1, 1, self.num_attention_heads, 1)) + key = self.pe_concat([k_nope, k_pe]) + + return query, key, value + + def get_tnd_query_key_value_tensors(self, + hidden_states, + key_value_states=None, + position_ids=None, + packed_seq_params=None, + inference_params=None, + rotary_pos_emb=None): + """get_tnd_query_key_value_tensors""" + if key_value_states: + raise NotImplementedError("`position_ids` is not supported for now") + if position_ids: + raise NotImplementedError("`position_ids` is not supported for now") + if packed_seq_params: + raise NotImplementedError("`packed_seq_params` is not supported for now") + if inference_params: + raise NotImplementedError("`inference_params` is not supported for now") + + tmp = self.shape(hidden_states) + seq_len, bs = tmp[0], tmp[1] + + seq_len = self.shape(hidden_states)[0] + rotary_pos_emb = self.rotary_pos_emb(seq_len) + + qkv_combo = self.linear_qkv(hidden_states)[0] + + q_a, compressed_kv, k_pe = self.split( + qkv_combo, + [ + self.q_rank, + self.kv_lora_rank, + self.qk_pos_emb_head_dim, + ], + dim=-1, + ) + + if self.q_layernorm is not None: + q_a = self.q_layernorm(q_a) + q = self.linear_qb(q_a)[0] + q = self.reshape(q, (seq_len, bs, self.num_attention_heads, -1)) + + q_nope, q_pe = self.split( + q, [self.qk_head_dim, self.qk_pos_emb_head_dim], dim=-1 + ) + + else: + q = self.reshape(q_a, (seq_len, bs, self.num_attention_heads, -1)) + q_nope, q_pe = self.split( + q, [self.qk_head_dim, self.qk_pos_emb_head_dim], dim=-1 + ) + + k_pe = self.reshape(k_pe, (seq_len, bs, 1, self.qk_pos_emb_head_dim)) + compressed_kv_norm = self.k_layernorm(compressed_kv) + + kv = self.linear_kvb(compressed_kv_norm)[0] + kv = self.reshape(kv, ( + seq_len, + bs, + self.num_attention_heads, + self.qk_head_dim + self.v_head_dim, + )) + + k_nope, value = self.split(kv, [self.qk_head_dim, self.v_head_dim], dim=-1) + + if rotary_pos_emb is not None: + q_pe = self.apply_rotary_emb( + q_pe, + rotary_pos_emb, + rotary_interleaved=self.config.rotary_interleaved, + multi_latent_attention=self.config.multi_latent_attention + ) + k_pe = self.apply_rotary_emb( + k_pe, + rotary_pos_emb, + rotary_interleaved=self.config.rotary_interleaved, + multi_latent_attention=self.config.multi_latent_attention, + input_is_parallel=True + ) + + query = self.pe_tnd_concat([q_nope, q_pe]) + k_pe = self.tile_tnd_kv(k_pe, (1, 1, self.num_attention_heads, 1)) + key = self.pe_tnd_concat([k_nope, k_pe]) + + query = self.reshape(query, (bs * seq_len, self.num_attention_heads, -1)) + key = self.reshape(key, (bs * seq_len, self.num_attention_heads, -1)) + value = self.reshape(value, (bs * seq_len, self.num_attention_heads, -1)) + + return query, key, value + + +class MLASelfAttentionMegatron(MultiLatentAttention): + """MLA Self-attention layer class + Self-attention layer takes input with size [s, b, h] + and returns output of the same size. + """ + + def __init__( + self, + config: MLATransformerConfig, + submodules: MLASelfAttentionSubmodulesMegatron, + layer_number: int, + attn_mask_type: str = None, + cp_comm_type: str = None, + ) -> None: + super().__init__( + config=config, + submodules=submodules, + layer_number=layer_number, + attn_mask_type=attn_mask_type, + cp_comm_type=cp_comm_type, + attention_type="self", + ) + + self.transpose = ms.ops.auto_generate.Transpose() + self.split = ms.ops.auto_generate.SplitWithSize().add_prim_attr("skip_redistribution", True) + self.tile_kv = ms.ops.auto_generate.Tile() + self.tile_tnd_kv = ms.ops.auto_generate.Tile() + self.pe_concat = ms.ops.auto_generate.Concat(axis=3) + self.pe_tnd_concat = ms.ops.auto_generate.Concat(axis=3) + self.apply_rotary_emb = ApplyRotaryPosEmb(config) + self.reshape = ms.ops.auto_generate.Reshape() + + if self.config.q_lora_rank is None: + # Not projecting query + self.linear_q_proj = build_module( + submodules.linear_q_proj, + self.config.hidden_size, + self.config.num_attention_heads * self.q_head_dim, + config=self.config, + init_method=self.config.init_method, + gather_output=False, + bias=False, + skip_bias_add=False, + is_expert=False, + ) + + else: + self.linear_q_down_proj = build_module( + submodules.linear_q_down_proj, + self.config.hidden_size, + self.config.q_lora_rank, + config=self.config, + init_method=self.config.init_method, + bias=False, + skip_bias_add=False, + skip_weight_param_allocation=False, + ) + + self.linear_q_up_proj = build_module( + submodules.linear_q_up_proj, + self.config.q_lora_rank, + self.config.num_attention_heads * self.q_head_dim, + config=self.config, + init_method=self.config.init_method, + gather_output=False, + bias=False, + skip_bias_add=False, + skip_weight_param_allocation=False, + ) + + self.linear_kv_down_proj = build_module( + submodules.linear_kv_down_proj, + self.config.hidden_size, + self.kv_lora_rank + self.qk_pos_emb_head_dim, + config=self.config, + init_method=self.config.init_method, + bias=False, + skip_bias_add=False, + skip_weight_param_allocation=False, + ) + + self.linear_kv_up_proj = build_module( + submodules.linear_kv_up_proj, + self.config.kv_lora_rank, + self.config.num_attention_heads * (self.config.qk_head_dim + self.config.v_head_dim), + config=self.config, + init_method=self.config.init_method, + gather_output=False, + bias=False, + skip_bias_add=False, + is_expert=False, + ) + + if self.config.q_lora_rank is not None: + self.q_layernorm = build_module( + submodules.q_layernorm, + dim=self.config.q_lora_rank, + config=self.config, + param_init_type=self.config.layernorm_compute_dtype, + eps=self.config.layernorm_epsilon, + ) + + self.kv_layernorm = build_module( + submodules.kv_layernorm, + dim=self.config.kv_lora_rank, + config=self.config, + param_init_type=self.config.layernorm_compute_dtype, + eps=self.config.layernorm_epsilon, + ) + + if _get_parallel_mode() in (ParallelMode.AUTO_PARALLEL,) and _is_sharding_propagation(): + self.sharding_propagation() + elif _get_parallel_mode() in (ParallelMode.SEMI_AUTO_PARALLEL,): + self.shard_self_attn() + + def sharding_propagation_self_attn(self): + dp = self.config.data_parallel_size + tp = self.config.tensor_model_parallel_size + + self.tile_kv.shard(((1, dp, tp, 1),)) + self.tile_tnd_kv.shard(((1, dp, tp, 1),)) + + def shard_self_attn(self): + dp = self.config.data_parallel_size + tp = self.config.tensor_model_parallel_size + + self.tile_kv.shard(((1, dp, tp, 1),)) + self.tile_tnd_kv.shard(((1, dp, tp, 1),)) + self.pe_concat.shard(((1, dp, tp, 1), (1, dp, tp, 1))) + self.pe_tnd_concat.shard(((1, dp, tp, 1), (1, dp, tp, 1))) + + def get_query_key_value_tensors(self, hidden_states, + key_value_states=None, + position_ids=None, + packed_seq_params=None, + inference_params=None, + rotary_pos_emb=None): + """ + Derives `query`, `key` and `value` tensors from `hidden_states`. + """ + if key_value_states: + raise NotImplementedError("`key_value_states` is not supported for now") + if position_ids: + raise NotImplementedError("`position_ids` is not supported for now") + if packed_seq_params: + raise NotImplementedError("`packed_seq_params` is not supported for now") + if inference_params: + raise NotImplementedError("`inference_params` is not supported for now") + + if hidden_states.ndim != 3: + raise ValueError(f"hidden_states should be 3D, [s, b, n*h], got {hidden_states.ndim}D") + + # ========================================= + # Prepare RoPE and seqlen related params + # ========================================= + + seq_len = self.shape(hidden_states)[0] + rotary_pos_emb = self.rotary_pos_emb(seq_len) + + # ========================================= + # QKV down projection and layernorm + # ========================================= + if self.config.q_lora_rank is not None: + q_compressed, _ = self.linear_q_down_proj(hidden_states) + q_compressed = self.q_layernorm(q_compressed) + else: + q_compressed = hidden_states + + kv_combined, _ = self.linear_kv_down_proj(hidden_states) + + kv_compressed, k_pos_emb = self.split( + kv_combined, [self.config.kv_lora_rank, self.config.qk_pos_emb_head_dim], dim=-1 + ) + + kv_compressed = self.kv_layernorm(kv_compressed) + + # ========================================= + # QKV up projection and RoPE apply + # ========================================= + def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_pos_emb): + if self.config.q_lora_rank is not None: + q, _ = self.linear_q_up_proj(q_compressed) + else: + q, _ = self.linear_q_proj(q_compressed) + + q_len, bsz, _ = self.shape(q) + + q = self.reshape(q, (q_len, bsz, self.num_attention_heads, self.q_head_dim)) + kv, _ = self.linear_kv_up_proj(kv_compressed) + + kv = self.reshape(kv, + (q_len, bsz, self.num_attention_heads, self.config.qk_head_dim + self.config.v_head_dim)) + + k_pos_emb = unsqueeze(k_pos_emb, 2) + + q_no_pe, q_pos_emb = self.split( + q, [self.config.qk_head_dim, self.config.qk_pos_emb_head_dim], dim=-1 + ) + + k_no_pe, value = self.split( + kv, [self.config.qk_head_dim, self.config.v_head_dim], dim=-1 + ) + + if rotary_pos_emb: + q_pos_emb = self.apply_rotary_emb( + q_pos_emb, + rotary_pos_emb, + rotary_interleaved=self.config.rotary_interleaved, + multi_latent_attention=self.config.multi_latent_attention + ) + k_pos_emb = self.apply_rotary_emb( + k_pos_emb, + rotary_pos_emb, + rotary_interleaved=self.config.rotary_interleaved, + multi_latent_attention=self.config.multi_latent_attention, + input_is_parallel=True + ) + + query = self.pe_concat([q_no_pe, q_pos_emb]) + + k_pos_emb = self.tile_kv(k_pos_emb, (1, 1, self.num_attention_heads, 1)) + key = self.pe_concat([k_no_pe, k_pos_emb]) + + return query, key, value + + query, key, value = qkv_up_proj_and_rope_apply( + q_compressed, kv_compressed, k_pos_emb, rotary_pos_emb + ) + + return query, key, value + + def get_tnd_query_key_value_tensors(self, hidden_states, + key_value_states=None, + position_ids=None, + packed_seq_params=None, + inference_params=None, + rotary_pos_emb=None): + """ + Derives `query`, `key` and `value` tensors from `hidden_states`. + """ + if key_value_states: + raise NotImplementedError("`key_value_states` is not supported for now") + if position_ids: + raise NotImplementedError("`position_ids` is not supported for now") + if packed_seq_params: + raise NotImplementedError("`packed_seq_params` is not supported for now") + if inference_params: + raise NotImplementedError("`inference_params` is not supported for now") + + if hidden_states.ndim != 3: + raise ValueError(f"hidden_states should be 3D, [s, b, n*h], got {hidden_states.ndim}D") + + # ========================================= + # Prepare RoPE and seqlen related params + # ========================================= + + seq_len = self.shape(hidden_states)[0] + rotary_pos_emb = self.rotary_pos_emb(seq_len) + + # ========================================= + # QKV down projection and layernorm + # ========================================= + if self.config.q_lora_rank is not None: + q_compressed, _ = self.linear_q_down_proj(hidden_states) + q_compressed = self.q_layernorm(q_compressed) + else: + q_compressed = hidden_states + + kv_combined, _ = self.linear_kv_down_proj(hidden_states) + + kv_compressed, k_pos_emb = self.split( + kv_combined, [self.config.kv_lora_rank, self.config.qk_pos_emb_head_dim], dim=-1 + ) + + kv_compressed = self.kv_layernorm(kv_compressed) + + # ========================================= + # QKV up projection and RoPE apply + # ========================================= + def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_pos_emb): + if self.config.q_lora_rank is not None: + q, _ = self.linear_q_up_proj(q_compressed) + else: + q, _ = self.linear_q_proj(q_compressed) + + q_len, bsz, _ = self.shape(q) + + q = self.reshape(q, (q_len, bsz, self.num_attention_heads, self.q_head_dim)) + kv, _ = self.linear_kv_up_proj(kv_compressed) + + kv = self.reshape(kv, + (q_len, bsz, self.num_attention_heads, self.config.qk_head_dim + self.config.v_head_dim)) + + k_pos_emb = unsqueeze(k_pos_emb, 2) + + q_no_pe, q_pos_emb = self.split( + q, [self.config.qk_head_dim, self.config.qk_pos_emb_head_dim], dim=-1 + ) + + k_no_pe, value = self.split( + kv, [self.config.qk_head_dim, self.config.v_head_dim], dim=-1 + ) + + if rotary_pos_emb: + q_pos_emb = self.apply_rotary_emb( + q_pos_emb, + rotary_pos_emb, + rotary_interleaved=self.config.rotary_interleaved, + multi_latent_attention=self.config.multi_latent_attention + ) + k_pos_emb = self.apply_rotary_emb( + k_pos_emb, + rotary_pos_emb, + rotary_interleaved=self.config.rotary_interleaved, + multi_latent_attention=self.config.multi_latent_attention, + input_is_parallel=True + ) + + query = self.pe_tnd_concat([q_no_pe, q_pos_emb]) + + k_pos_emb = self.tile_tnd_kv(k_pos_emb, (1, 1, self.num_attention_heads, 1)) + key = self.pe_tnd_concat([k_no_pe, k_pos_emb]) + + query = self.reshape(query, (bsz * q_len, self.num_attention_heads, -1)) + key = self.reshape(key, (bsz * q_len, self.num_attention_heads, -1)) + value = self.reshape(value, (bsz * q_len, self.num_attention_heads, -1)) + return query, key, value + + query, key, value = qkv_up_proj_and_rope_apply( + q_compressed, kv_compressed, k_pos_emb, rotary_pos_emb + ) + + return query, key, value diff --git a/mindformers/parallel_core/training_graph/transformer/rotary_pos_embedding.py b/mindformers/parallel_core/training_graph/transformer/rotary_pos_embedding.py index f07efe4a3..6b5d43faf 100644 --- a/mindformers/parallel_core/training_graph/transformer/rotary_pos_embedding.py +++ b/mindformers/parallel_core/training_graph/transformer/rotary_pos_embedding.py @@ -29,9 +29,42 @@ from mindspore import nn, Tensor from mindspore import ops from mindspore.context import ParallelMode from mindspore.parallel._utils import _get_parallel_mode, _is_sharding_propagation -from mindspore.ops.auto_generate import (AddExt, Reshape, Mul, Cos, Sin, - Split, Neg, Concat, StackExt, StridedSlice, Outer) -from mindformers.parallel_core.transformer_config import TransformerConfig +from mindspore.ops.auto_generate import AddExt, Reshape, Mul, Cos, Sin, Split, Neg, Concat, StackExt, StridedSlice, \ + Outer +from mindformers.parallel_core.transformer_config import TransformerConfig, MLATransformerConfig + + +def _get_rope(config: MLATransformerConfig): + return RotaryEmbedding( + config.qk_pos_emb_head_dim, + rotary_percent=config.rotary_percent, + rotary_interleaved=config.rotary_interleaved, + seq_len_interpolation_factor=config.rotary_seq_len_interpolation_factor, + rotary_base=config.rotary_base, + rope_scaling=config.rope_scaling + ) + + +def _get_yarn(config: MLATransformerConfig): + return YarnRotaryEmbedding( + config.qk_pos_emb_head_dim, + rotary_percent=config.rotary_percent, + rotary_interleaved=config.rotary_interleaved, + seq_len_interpolation_factor=config.rotary_seq_len_interpolation_factor, + rotary_base=config.rotary_base, + scaling_factor=config.rotary_scaling_factor, + original_max_position_embeddings=config.max_position_embeddings, + beta_fast=config.beta_fast, + beta_slow=config.beta_slow, + mscale=config.mscale, + mscale_all_dim=config.mscale_all_dim, + ) + + +ROPE_FUNCTIONS = { + 'rope': _get_rope, + 'yarn': _get_yarn, +} class RotaryEmbedding(nn.Cell): @@ -270,17 +303,23 @@ class ApplyRotaryPosEmb(nn.Cell): """ def __init__(self, - config: TransformerConfig): + config: TransformerConfig + ): super(ApplyRotaryPosEmb, self).__init__() self.append_eod = config.use_eod_attn_mask_compression self.add = AddExt() + self.add_input_is_parallel = AddExt() self.mul = Mul() + self.mul_input_is_parallel = Mul() self.cos = Cos() self.sin = Sin() self.neg = Neg() + self.neg_input_is_parallel = Neg() self.split = Split(axis=-1, output_num=2) + self.split_input_is_parallel = Split(axis=-1, output_num=2) self.tnd_split = Split(axis=-1, output_num=2) self.cat = Concat(axis=-1) + self.cat_input_is_parallel = Concat(axis=-1) self.stack = StackExt(dim=-1) self.slice = StridedSlice() self.reshape = Reshape() @@ -295,15 +334,18 @@ class ApplyRotaryPosEmb(nn.Cell): t: Tensor, freqs: tuple, rotary_interleaved: bool = False, - multi_latent_attention: bool = False) -> Tensor: + multi_latent_attention: bool = False, + input_is_parallel: bool = False) -> Tensor: """Apply rotary position embedding to input tensor. Args: t (Tensor): Input tensor is of shape [seq_length, ... , dim] freqs (tuple): A tuple of frequencies and mscale. Rotary position embedding frequencies - of shape [seq_length, ... , dim], mscale is float. - rotary_interleaved (bool): Whether to use interleaved rotary position embedding - multi_latent_attention (bool): Whether to use multi latent attention + of shape [seq_length, ... , dim], mscale is float + rotary_interleaved (bool): Whether to use interleaved rotary position embedding. Default: False + multi_latent_attention (bool): Whether to use multi latent attention. Default: False + input_is_parallel (bool): Whether the input tensor is already sliced。The shard strategy will be different + according to this argument. Default: False Returns: Tensor: Output tensor after applying rotary position embedding @@ -317,28 +359,59 @@ class ApplyRotaryPosEmb(nn.Cell): t_not_rotary = self.slice(t, (0, 0, 0, rot_dim), (seq_len, bs, n_heads, head_dim), (1, 1, 1, 1)) t = self.slice(t, (0, 0, 0, 0), (seq_len, bs, n_heads, rot_dim), (1, 1, 1, 1)) + if input_is_parallel: + cat = self.cat_input_is_parallel + mul = self.mul_input_is_parallel + add = self.add_input_is_parallel + else: + cat = self.cat + mul = self.mul + add = self.add + if multi_latent_attention: x1 = t[..., 0::2] x2 = t[..., 1::2] - t = self.cat((x1, x2)) + t = cat((x1, x2)) cos_ = self.cos(freqs * m_scale).astype(t.dtype) sin_ = self.sin(freqs * m_scale).astype(t.dtype) - output = self.add(self.mul(t, cos_), self.mul(self._rotate_half(t, rotary_interleaved), sin_)) + t_rot = self._rotate_half(t, rotary_interleaved, input_is_parallel) + output = add(mul(t, cos_), mul(t_rot, sin_)) + if t_not_rotary is not None: - output = self.cat((output, t_not_rotary)) + output = cat((output, t_not_rotary)) return output - def _rotate_half(self, t: Tensor, rotary_interleaved: bool = False) -> Tensor: + def _rotate_half(self, t: Tensor, rotary_interleaved: bool = False, input_is_parallel: bool = False) -> Tensor: + """Rotates half of the input tensor for rotary position embeddings. + + Args: + t: Input tensor of shape (seq_len, bs, n_heads, head_dim) + rotary_interleaved: If True, processes interleaved features + input_is_parallel: If True, uses parallel-optimized operations + + Returns: + Rotated tensor with same shape as input + """ seq_len, bs, n_heads, head_dim = t.shape + + if input_is_parallel: + split = self.split_input_is_parallel + neg = self.neg_input_is_parallel + cat = self.cat_input_is_parallel + else: + split = self.split + neg = self.neg + cat = self.cat + if rotary_interleaved: t_1 = self.slice(t, (0, 0, 0, 0), (seq_len, bs, n_heads, head_dim), (1, 1, 1, 2)) t_2 = self.slice(t, (0, 0, 0, 1), (seq_len, bs, n_heads, head_dim), (1, 1, 1, 2)) - t_rot = self.reshape(self.stack((self.neg(t_2), t_1)), (seq_len, bs, n_heads, -1)) + t_rot = self.reshape(self.stack((neg(t_2), t_1)), (seq_len, bs, n_heads, -1)) else: - t_1, t_2 = self.split(t) - t_rot = self.cat((self.neg(t_2), t_1)) + t_1, t_2 = split(t) + t_rot = cat((neg(t_2), t_1)) return t_rot def shard(self, config: TransformerConfig): @@ -347,6 +420,7 @@ class ApplyRotaryPosEmb(nn.Cell): tp = config.tensor_model_parallel_size if config and config.tensor_model_parallel_size is not None else 1 strategy_in = (1, dp, tp, 1) + strategy_in_input_is_parallel = (1, dp, 1, 1) sin_in_strategy = ((1, 1, 1, 1),) cos_in_strategy = ((1, 1, 1, 1),) @@ -358,15 +432,21 @@ class ApplyRotaryPosEmb(nn.Cell): add_in_strategy = (strategy_in, strategy_in) self.add.shard(in_strategy=add_in_strategy) + self.add_input_is_parallel.shard(in_strategy=(strategy_in_input_is_parallel, strategy_in_input_is_parallel)) if self.append_eod: self.mul.shard(in_strategy=(strategy_in, (1, strategy_in[0], 1, 1))) + self.mul_input_is_parallel.shard(in_strategy=(strategy_in_input_is_parallel, (1, strategy_in[0], 1, 1))) else: self.mul.shard(in_strategy=(strategy_in, (1, 1, 1, 1))) + self.mul_input_is_parallel.shard(in_strategy=(strategy_in_input_is_parallel, (1, 1, 1, 1))) self.split.shard(in_strategy=split_in_strategy) + self.split_input_is_parallel.shard(in_strategy=(strategy_in_input_is_parallel,)) self.cat.shard(in_strategy=cat_in_strategy) + self.cat_input_is_parallel.shard(in_strategy=(strategy_in_input_is_parallel, strategy_in_input_is_parallel)) self.stack.shard(in_strategy=stack_in_strategy) self.slice.shard(in_strategy=slice_in_strategy) self.neg.shard(in_strategy=neg_in_strategy) + self.neg_input_is_parallel.shard(in_strategy=(strategy_in_input_is_parallel,)) self.sin.shard(in_strategy=sin_in_strategy) self.cos.shard(in_strategy=cos_in_strategy) diff --git a/tests/st/test_ut/test_parallel_core/test_training_graph/test_transformer/test_multi_latent_attention/__init__.py b/tests/st/test_ut/test_parallel_core/test_training_graph/test_transformer/test_multi_latent_attention/__init__.py new file mode 100644 index 000000000..2c703f37d --- /dev/null +++ b/tests/st/test_ut/test_parallel_core/test_training_graph/test_transformer/test_multi_latent_attention/__init__.py @@ -0,0 +1,14 @@ +# Copyright 2025 Huawei Technologies Co., Ltd +# # +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# # +# http://www.apache.org/licenses/LICENSE-2.0 +# # +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ diff --git a/tests/st/test_ut/test_parallel_core/test_training_graph/test_transformer/test_multi_latent_attention/data_gen_utils.py b/tests/st/test_ut/test_parallel_core/test_training_graph/test_transformer/test_multi_latent_attention/data_gen_utils.py new file mode 100644 index 000000000..6e5920532 --- /dev/null +++ b/tests/st/test_ut/test_parallel_core/test_training_graph/test_transformer/test_multi_latent_attention/data_gen_utils.py @@ -0,0 +1,306 @@ +# Copyright 2025 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ +"""Get reference data.""" + +import numpy as np + + +def get_init_params(loc, scale, config, seq_length, batch_size, hidden_size): + """ + Generates initial parameters for Multi-head Latent Attention (MLA). + Input shape is (seq_length, batch_size, hidden_size). + """ + rng = np.random.default_rng(42) + hidden_state = rng.normal(loc=loc, scale=scale, size=(seq_length, batch_size, hidden_size)) + megatron_state_dict = { + "linear_kv_down_proj.weight": rng.normal( + loc=loc, scale=scale, size=(config.kv_lora_rank + config.qk_pos_emb_head_dim, config.hidden_size)), + "linear_kv_up_proj.weight": rng.normal( + loc=loc, scale=scale, + size=(config.num_attention_heads * (config.qk_head_dim + config.v_head_dim), config.kv_lora_rank) + ), + "linear_proj.weight": rng.normal( + loc=loc, scale=scale, size=(config.hidden_size, config.v_head_dim * config.num_attention_heads)), + "kv_layernorm.weight": rng.normal(loc=loc, scale=scale, size=config.kv_lora_rank), + } + + if config.q_lora_rank: + megatron_state_dict["linear_q_down_proj.weight"] = rng.normal( + loc=loc, scale=scale, size=(config.q_lora_rank, config.hidden_size)) + megatron_state_dict["linear_q_up_proj.weight"] = rng.normal( + loc=loc, scale=scale, + size=(config.num_attention_heads * (config.qk_head_dim + config.qk_pos_emb_head_dim), config.q_lora_rank), + ) + megatron_state_dict["q_layernorm.weight"] = rng.normal(loc=loc, scale=scale, size=config.q_lora_rank) + else: + megatron_state_dict["linear_q_proj.weight"] = rng.normal( + loc=loc, scale=scale, + size=(config.num_attention_heads * (config.qk_head_dim + config.qk_pos_emb_head_dim), config.hidden_size), + ) + + mindspeed_state_dict = { + "linear_kvb.weight": megatron_state_dict["linear_kv_up_proj.weight"], + "linear_proj.weight": megatron_state_dict["linear_proj.weight"], + "k_layernorm.weight": megatron_state_dict["kv_layernorm.weight"], + } + + if config.q_lora_rank: + mindspeed_state_dict["linear_qkv.weight"] = np.concatenate(( + megatron_state_dict["linear_q_down_proj.weight"], + megatron_state_dict["linear_kv_down_proj.weight"] + )) + mindspeed_state_dict["linear_qb.weight"] = megatron_state_dict["linear_q_up_proj.weight"] + mindspeed_state_dict["q_layernorm.weight"] = megatron_state_dict["q_layernorm.weight"] + else: + mindspeed_state_dict["linear_qkv.weight"] = np.concatenate(( + megatron_state_dict["linear_q_proj.weight"], + megatron_state_dict["linear_kv_down_proj.weight"] + )) + + return hidden_state, megatron_state_dict, mindspeed_state_dict + + +def get_gpu_datas() -> dict[str, np.ndarray]: + """Generate gpu data for test.""" + q8_flash_ql_kl = [[[-0.00221252, -0.00332642, -0.00065231, -0.00078201, -0.00337219, + 0.00312805, 0.00157166, -0.00262451, -0.00154114, -0.00549316, + -0.00117493, 0.00473022, 0.00396729, 0.00215149, 0.00233459, + 0.00063705], + [0.00430298, 0.00549316, 0.00195312, 0.00224304, 0.02026367, + -0.01123047, -0.00744629, 0.01275635, 0.00769043, 0.01556396, + -0.00270081, -0.01525879, -0.01501465, -0.01220703, 0.01214600, + 0.00756836]], + + [[-0.00221252, -0.00332642, -0.00065231, -0.00078201, -0.00337219, + 0.00312805, 0.00157166, -0.00262451, -0.00154114, -0.00549316, + -0.00117493, 0.00473022, 0.00396729, 0.00215149, 0.00233459, + 0.00063705], + [0.00430298, 0.00549316, 0.00195312, 0.00224304, 0.02026367, + -0.01123047, -0.00744629, 0.01275635, 0.00769043, 0.01556396, + -0.00270081, -0.01525879, -0.01501465, -0.01220703, 0.01214600, + 0.00756836]], + + [[-0.00221252, -0.00332642, -0.00065231, -0.00078201, -0.00337219, + 0.00312805, 0.00157166, -0.00262451, -0.00154114, -0.00549316, + -0.00117493, 0.00473022, 0.00396729, 0.00215149, 0.00233459, + 0.00063705], + [0.00430298, 0.00549316, 0.00195312, 0.00224304, 0.02026367, + -0.01123047, -0.00744629, 0.01275635, 0.00769043, 0.01556396, + -0.00270081, -0.01525879, -0.01501465, -0.01220703, 0.01214600, + 0.00756836]], + + [[-0.00221252, -0.00332642, -0.00065231, -0.00078201, -0.00337219, + 0.00312805, 0.00157166, -0.00262451, -0.00154114, -0.00549316, + -0.00117493, 0.00473022, 0.00396729, 0.00215149, 0.00233459, + 0.00063705], + [0.00430298, 0.00549316, 0.00195312, 0.00224304, 0.02026367, + -0.01123047, -0.00744629, 0.01275635, 0.00769043, 0.01556396, + -0.00270081, -0.01525879, -0.01501465, -0.01220703, 0.01214600, + 0.00756836]]] + q0_flash_ql_kl = [[[-0.00221252, -0.00332642, -0.00065231, -0.00078201, -0.00337219, + 0.00312805, 0.00157166, -0.00262451, -0.00154114, -0.00549316, + -0.00117493, 0.00473022, 0.00396729, 0.00215149, 0.00233459, + 0.00063705], + [0.00430298, 0.00549316, 0.00195312, 0.00224304, 0.02026367, + -0.01123047, -0.00744629, 0.01275635, 0.00769043, 0.01556396, + -0.00270081, -0.01525879, -0.01501465, -0.01220703, 0.01214600, + 0.00756836]], + + [[-0.00221252, -0.00332642, -0.00065231, -0.00078201, -0.00337219, + 0.00312805, 0.00157166, -0.00262451, -0.00154114, -0.00549316, + -0.00117493, 0.00473022, 0.00396729, 0.00215149, 0.00233459, + 0.00063705], + [0.00430298, 0.00549316, 0.00195312, 0.00224304, 0.02026367, + -0.01123047, -0.00744629, 0.01275635, 0.00769043, 0.01556396, + -0.00270081, -0.01525879, -0.01501465, -0.01220703, 0.01214600, + 0.00756836]], + + [[-0.00221252, -0.00332642, -0.00065231, -0.00078201, -0.00337219, + 0.00312805, 0.00157166, -0.00262451, -0.00154114, -0.00549316, + -0.00117493, 0.00473022, 0.00396729, 0.00215149, 0.00233459, + 0.00063705], + [0.00430298, 0.00549316, 0.00195312, 0.00224304, 0.02026367, + -0.01123047, -0.00744629, 0.01275635, 0.00769043, 0.01556396, + -0.00270081, -0.01525879, -0.01501465, -0.01220703, 0.01214600, + 0.00756836]], + + [[-0.00221252, -0.00332642, -0.00065231, -0.00078201, -0.00337219, + 0.00312805, 0.00157166, -0.00262451, -0.00154114, -0.00549316, + -0.00117493, 0.00473022, 0.00396729, 0.00215149, 0.00233459, + 0.00063705], + [0.00430298, 0.00549316, 0.00195312, 0.00224304, 0.02026367, + -0.01123047, -0.00744629, 0.01275635, 0.00769043, 0.01556396, + -0.00270081, -0.01525879, -0.01501465, -0.01220703, 0.01214600, + 0.00756836]]] + q8_flash_kl = [[[-0.00221252, -0.00332642, -0.00065231, -0.00078201, -0.00337219, + 0.00312805, 0.00157166, -0.00262451, -0.00154114, -0.00549316, + -0.00117493, 0.00473022, 0.00396729, 0.00215149, 0.00233459, + 0.00063705], + [0.00430298, 0.00549316, 0.00195312, 0.00224304, 0.02026367, + -0.01123047, -0.00744629, 0.01275635, 0.00769043, 0.01556396, + -0.00270081, -0.01525879, -0.01501465, -0.01220703, 0.01214600, + 0.00756836]], + + [[-0.00221252, -0.00332642, -0.00065231, -0.00078201, -0.00337219, + 0.00312805, 0.00157166, -0.00262451, -0.00154114, -0.00549316, + -0.00117493, 0.00473022, 0.00396729, 0.00215149, 0.00233459, + 0.00063705], + [0.00430298, 0.00549316, 0.00195312, 0.00224304, 0.02026367, + -0.01123047, -0.00744629, 0.01275635, 0.00769043, 0.01556396, + -0.00270081, -0.01525879, -0.01501465, -0.01220703, 0.01214600, + 0.00756836]], + + [[-0.00221252, -0.00332642, -0.00065231, -0.00078201, -0.00337219, + 0.00312805, 0.00157166, -0.00262451, -0.00154114, -0.00549316, + -0.00117493, 0.00473022, 0.00396729, 0.00215149, 0.00233459, + 0.00063705], + [0.00430298, 0.00549316, 0.00195312, 0.00224304, 0.02026367, + -0.01123047, -0.00744629, 0.01275635, 0.00769043, 0.01556396, + -0.00270081, -0.01525879, -0.01501465, -0.01220703, 0.01214600, + 0.00756836]], + + [[-0.00221252, -0.00332642, -0.00065231, -0.00078201, -0.00337219, + 0.00312805, 0.00157166, -0.00262451, -0.00154114, -0.00549316, + -0.00117493, 0.00473022, 0.00396729, 0.00215149, 0.00233459, + 0.00063705], + [0.00430298, 0.00549316, 0.00195312, 0.00224304, 0.02026367, + -0.01123047, -0.00744629, 0.01275635, 0.00769043, 0.01556396, + -0.00270081, -0.01525879, -0.01501465, -0.01220703, 0.01214600, + 0.00756836]]] + + return { + "q8_flash_ql_kl": np.array(q8_flash_ql_kl), + "q0_flash_ql_kl": np.array(q0_flash_ql_kl), + "q8_flash_kl": np.array(q8_flash_kl), + } + + +def get_golden() -> dict[str, np.ndarray]: + """Generate golden data for test.""" + q8_flash_ql_kl = [[[-0.00218201, -0.00332642, -0.00063324, -0.00079727, -0.00341797, + 0.00314331, 0.00164032, -0.00270081, -0.00156403, -0.00549316, + -0.00119781, 0.00479126, 0.00402832, 0.00221252, 0.00230408, + 0.00059891], + [0.00433350, 0.00561523, 0.00198364, 0.00219727, 0.02026367, + -0.01123047, -0.00738525, 0.01281738, 0.00775146, 0.01562500, + -0.00265503, -0.01538086, -0.01501465, -0.01220703, 0.01214600, + 0.00756836]], + + [[-0.00218201, -0.00332642, -0.00063324, -0.00079727, -0.00341797, + 0.00314331, 0.00164032, -0.00270081, -0.00156403, -0.00549316, + -0.00119781, 0.00479126, 0.00402832, 0.00221252, 0.00230408, + 0.00059891], + [0.00433350, 0.00561523, 0.00198364, 0.00219727, 0.02026367, + -0.01123047, -0.00738525, 0.01281738, 0.00775146, 0.01562500, + -0.00265503, -0.01538086, -0.01501465, -0.01220703, 0.01214600, + 0.00756836]], + + [[-0.00218201, -0.00332642, -0.00063324, -0.00079727, -0.00341797, + 0.00314331, 0.00164032, -0.00270081, -0.00156403, -0.00549316, + -0.00119781, 0.00479126, 0.00402832, 0.00221252, 0.00230408, + 0.00059891], + [0.00433350, 0.00561523, 0.00198364, 0.00219727, 0.02026367, + -0.01123047, -0.00738525, 0.01281738, 0.00775146, 0.01562500, + -0.00265503, -0.01538086, -0.01501465, -0.01220703, 0.01214600, + 0.00756836]], + + [[-0.00218201, -0.00332642, -0.00063324, -0.00079727, -0.00341797, + 0.00314331, 0.00164032, -0.00270081, -0.00156403, -0.00549316, + -0.00119781, 0.00479126, 0.00402832, 0.00221252, 0.00230408, + 0.00059891], + [0.00433350, 0.00561523, 0.00198364, 0.00219727, 0.02026367, + -0.01123047, -0.00738525, 0.01281738, 0.00775146, 0.01562500, + -0.00265503, -0.01538086, -0.01501465, -0.01220703, 0.01214600, + 0.00756836]]] + q0_flash_ql_kl = [[[-0.00218201, -0.00332642, -0.00063324, -0.00079727, -0.00341797, + 0.00314331, 0.00164032, -0.00270081, -0.00156403, -0.00549316, + -0.00119781, 0.00479126, 0.00402832, 0.00221252, 0.00230408, + 0.00059891], + [0.00433350, 0.00561523, 0.00198364, 0.00219727, 0.02026367, + -0.01123047, -0.00738525, 0.01281738, 0.00775146, 0.01562500, + -0.00265503, -0.01538086, -0.01501465, -0.01220703, 0.01214600, + 0.00756836]], + + [[-0.00218201, -0.00332642, -0.00063324, -0.00079727, -0.00341797, + 0.00314331, 0.00164032, -0.00270081, -0.00156403, -0.00549316, + -0.00119781, 0.00479126, 0.00402832, 0.00221252, 0.00230408, + 0.00059891], + [0.00433350, 0.00561523, 0.00198364, 0.00219727, 0.02026367, + -0.01123047, -0.00738525, 0.01281738, 0.00775146, 0.01562500, + -0.00265503, -0.01538086, -0.01501465, -0.01220703, 0.01214600, + 0.00756836]], + + [[-0.00218201, -0.00332642, -0.00063324, -0.00079727, -0.00341797, + 0.00314331, 0.00164032, -0.00270081, -0.00156403, -0.00549316, + -0.00119781, 0.00479126, 0.00402832, 0.00221252, 0.00230408, + 0.00059891], + [0.00433350, 0.00561523, 0.00198364, 0.00219727, 0.02026367, + -0.01123047, -0.00738525, 0.01281738, 0.00775146, 0.01562500, + -0.00265503, -0.01538086, -0.01501465, -0.01220703, 0.01214600, + 0.00756836]], + + [[-0.00218201, -0.00332642, -0.00063324, -0.00079727, -0.00341797, + 0.00314331, 0.00164032, -0.00270081, -0.00156403, -0.00549316, + -0.00119781, 0.00479126, 0.00402832, 0.00221252, 0.00230408, + 0.00059891], + [0.00433350, 0.00561523, 0.00198364, 0.00219727, 0.02026367, + -0.01123047, -0.00738525, 0.01281738, 0.00775146, 0.01562500, + -0.00265503, -0.01538086, -0.01501465, -0.01220703, 0.01214600, + 0.00756836]]] + q8_flash_kl = [[[-0.00218201, -0.00332642, -0.00063324, -0.00079727, -0.00341797, + 0.00314331, 0.00164032, -0.00270081, -0.00156403, -0.00549316, + -0.00119781, 0.00479126, 0.00402832, 0.00221252, 0.00230408, + 0.00059891], + [0.00433350, 0.00561523, 0.00198364, 0.00219727, 0.02026367, + -0.01123047, -0.00738525, 0.01281738, 0.00775146, 0.01562500, + -0.00265503, -0.01538086, -0.01501465, -0.01220703, 0.01214600, + 0.00756836]], + + [[-0.00218201, -0.00332642, -0.00063324, -0.00079727, -0.00341797, + 0.00314331, 0.00164032, -0.00270081, -0.00156403, -0.00549316, + -0.00119781, 0.00479126, 0.00402832, 0.00221252, 0.00230408, + 0.00059891], + [0.00433350, 0.00561523, 0.00198364, 0.00219727, 0.02026367, + -0.01123047, -0.00738525, 0.01281738, 0.00775146, 0.01562500, + -0.00265503, -0.01538086, -0.01501465, -0.01220703, 0.01214600, + 0.00756836]], + + [[-0.00218201, -0.00332642, -0.00063324, -0.00079727, -0.00341797, + 0.00314331, 0.00164032, -0.00270081, -0.00156403, -0.00549316, + -0.00119781, 0.00479126, 0.00402832, 0.00221252, 0.00230408, + 0.00059891], + [0.00433350, 0.00561523, 0.00198364, 0.00219727, 0.02026367, + -0.01123047, -0.00738525, 0.01281738, 0.00775146, 0.01562500, + -0.00265503, -0.01538086, -0.01501465, -0.01220703, 0.01214600, + 0.00756836]], + + [[-0.00218201, -0.00332642, -0.00063324, -0.00079727, -0.00341797, + 0.00314331, 0.00164032, -0.00270081, -0.00156403, -0.00549316, + -0.00119781, 0.00479126, 0.00402832, 0.00221252, 0.00230408, + 0.00059891], + [0.00433350, 0.00561523, 0.00198364, 0.00219727, 0.02026367, + -0.01123047, -0.00738525, 0.01281738, 0.00775146, 0.01562500, + -0.00265503, -0.01538086, -0.01501465, -0.01220703, 0.01214600, + 0.00756836]]] + + return { + "q8_flash_ql_kl": np.array(q8_flash_ql_kl), + "q0_flash_ql_kl": np.array(q0_flash_ql_kl), + "q8_flash_kl": np.array(q8_flash_kl), + } + + +GOLDEN_DATA = get_golden() +GPU_DATA = get_gpu_datas() diff --git a/tests/st/test_ut/test_parallel_core/test_training_graph/test_transformer/test_multi_latent_attention/run_multi_latent_attention.py b/tests/st/test_ut/test_parallel_core/test_training_graph/test_transformer/test_multi_latent_attention/run_multi_latent_attention.py new file mode 100644 index 000000000..ae165ea64 --- /dev/null +++ b/tests/st/test_ut/test_parallel_core/test_training_graph/test_transformer/test_multi_latent_attention/run_multi_latent_attention.py @@ -0,0 +1,166 @@ +# Copyright 2025 Huawei Technologies Co., Ltd +# # +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# # +# http://www.apache.org/licenses/LICENSE-2.0 +# # +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ +"""Run Multi-head Latent Attention (MLA) accuracy test with configurable parameters via args.""" +import os +import argparse +from pathlib import Path +import numpy as np +import mindspore as ms +from mindspore.communication import init +from mindspore.mint import ones +from mindformers.parallel_core.training_graph.transformer.multi_latent_attention import MLASelfAttention, \ + MLASelfAttentionMegatron, MLASelfAttentionSubmodules, MLASelfAttentionSubmodulesMegatron +from mindformers.parallel_core.transformer_config import MLATransformerConfig +from mindformers.parallel_core.training_graph.transformer.flash_attention import FlashAttention +from mindformers.parallel_core.training_graph.transformer.identity_op import IdentityOp +from mindformers.parallel_core.training_graph.transformer.norm import RMSNorm +from mindformers.parallel_core.training_graph.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear +from mindformers.parallel_core.utils.spec_utils import ModuleSpec, build_module +from tests.st.test_ut.test_parallel_core.test_training_graph.test_transformer.test_multi_latent_attention.data_gen_utils import get_init_params + +SCRIPT_DIR = Path(__file__).parent.resolve() + + +class MLARunner: + """Class to manage Multi-head Latent Attention (MLA) execution.""" + + def __init__(self, args_from_parser): + self.args = args_from_parser + self.compute_dtype = "bfloat16" + + rank_id_str: str | None = os.environ.get("RANK_ID") + self.rank_id = int(rank_id_str) if rank_id_str is not None else None + self.worker_num = int(os.environ.get("MS_WORKER_NUM", "1")) + + if self.rank_id is not None: + ms.set_auto_parallel_context( + parallel_mode=ms.ParallelMode.SEMI_AUTO_PARALLEL, + full_batch=True + ) + init() + + self.config = MLATransformerConfig( + use_flash_attention=self.args.use_flash_attn, + hidden_size=self.args.hidden_size, + data_parallel_size=self.worker_num // self.args.tensor_parallel, + tensor_model_parallel_size=self.args.tensor_parallel, + compute_dtype=self.compute_dtype, + num_attention_heads=2, + num_layers=1, + max_position_embeddings=2, + q_lora_rank=self.args.q_lora_rank, + kv_lora_rank=4, + v_head_dim=8, + qk_head_dim=4, + qk_pos_emb_head_dim=4, + hidden_dropout=0.0, + attention_dropout=0.0, + output_layer_init_method=None, + add_bias_linear=False + ) + + def build_model(self): + """Build and initialize Multi-head Latent Attention model.""" + core_attention = FlashAttention + q_layernorm = IdentityOp if self.args.q_layernorm.lower() == 'none' else RMSNorm + k_layernorm = IdentityOp if self.args.k_layernorm.lower() == 'none' else RMSNorm + + if self.args.struct == 'a2': + spec = ModuleSpec( + module=MLASelfAttention, + submodules=MLASelfAttentionSubmodules( + linear_qkv=ColumnParallelLinear, + linear_qb=ColumnParallelLinear, + linear_kvb=ColumnParallelLinear, + core_attention=core_attention, + linear_proj=RowParallelLinear, + q_layernorm=q_layernorm, + k_layernorm=k_layernorm + ) + ) + else: + spec = ModuleSpec( + module=MLASelfAttentionMegatron, + submodules=MLASelfAttentionSubmodulesMegatron( + linear_q_proj=ColumnParallelLinear, + linear_q_down_proj=ColumnParallelLinear, + linear_q_up_proj=ColumnParallelLinear, + linear_kv_down_proj=ColumnParallelLinear, + linear_kv_up_proj=ColumnParallelLinear, + core_attention=core_attention, + linear_proj=RowParallelLinear, + q_layernorm=q_layernorm, + kv_layernorm=k_layernorm + ) + ) + + model = build_module(spec, config=self.config, layer_number=1) + return model + + def run(self): + """Run the model with given inputs""" + net = self.build_model() + hidden_state, mega_weight, mind_weight = get_init_params(0.0, 0.15, self.config, self.args.seq_length, + self.args.batch_size, self.args.hidden_size) + hidden_state = ms.Tensor(hidden_state, dtype=ms.float32) + weight = mind_weight if self.args.struct == 'a2' else mega_weight + for k in weight: + weight[k] = ms.Parameter(ms.Tensor(weight[k], dtype=ms.float32)) + ms.load_param_into_net(net, weight) + + output = net( + hidden_state, + attention_mask=ones((self.args.batch_size, 1, self.args.seq_length, self.args.seq_length), dtype=ms.bool_), + rotary_pos_emb=ones((4, 1, 1, 4)) + ) + output_ms = {"output": output} + + if self.rank_id is None or int(self.rank_id) == 0: + # Convert to float32 for saving, common practice for bf16/fp16 + output_np = {k: v.asnumpy().astype(np.float32) for k, v in output_ms.items() if v is not None} + output_path = self.args.output_path + np.savez(output_path, **output_np) + + +def main(): + """Main function.""" + parser = argparse.ArgumentParser(description="Run Multi-head Latent Attention test") + # Input shape parameters + parser.add_argument("--struct", type=str, default="a2") + parser.add_argument("--batch_size", type=int, default=2) + parser.add_argument("--seq_length", type=int, default=4) + parser.add_argument("--hidden_size", type=int, default=16) + # Model configuration parameters + parser.add_argument("--q_lora_rank", type=int, default=8) + parser.add_argument("--use_flash_attn", type=lambda x: x.lower() == "true", default=True) + parser.add_argument("--q_layernorm", type=str, default="RMSNorm") + parser.add_argument("--k_layernorm", type=str, default="RMSNorm") + # Output and parallelism + parser.add_argument("--output_path", type=str, default="output_mla_ms.npz") + parser.add_argument("--tensor_parallel", type=int, default=1) + parser.add_argument("--test_name", type=str, default="q16_flash_ql_kl") + args = parser.parse_args() + args.q_lora_rank = None if args.q_lora_rank == 0 else args.q_lora_rank + + ms.context.set_context(deterministic="ON") + ms.set_context(mode=ms.GRAPH_MODE) + ms.set_seed(42) + + runner = MLARunner(args) + runner.run() + + +if __name__ == "__main__": + main() diff --git a/tests/st/test_ut/test_parallel_core/test_training_graph/test_transformer/test_multi_latent_attention/test_multi_latent_attention.py b/tests/st/test_ut/test_parallel_core/test_training_graph/test_transformer/test_multi_latent_attention/test_multi_latent_attention.py new file mode 100644 index 000000000..d9b12847a --- /dev/null +++ b/tests/st/test_ut/test_parallel_core/test_training_graph/test_transformer/test_multi_latent_attention/test_multi_latent_attention.py @@ -0,0 +1,309 @@ +# Copyright 2025 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ +"""Test Multi-head Latent Attention (MLA) with various configurations.""" +import os +import time +import socket +import subprocess +from pathlib import Path +import pytest +import numpy as np +from mindformers.tools.logger import logger +from tests.utils.double_benchmark import DoubleBenchmarkStandard, DoubleBenchmarkComparator +from tests.st.test_ut.test_parallel_core.test_training_graph.test_transformer.test_multi_latent_attention.data_gen_utils import GOLDEN_DATA, GPU_DATA + +BATCH_SIZE = 2 +SEQ_LENGTH = 4 +HIDDEN_SIZE = 16 +SINGLE_CARD_TEST_PARAM = 'model_args, golden_data_key' +SINGLE_CARD_TEST_CASES = [ + ( + # case 1 + # Input: + # - model_args: { + # 'q_lora_rank': 8, # Enable q_lora_rank with rank 8 for query projection + # 'use_flash_attn': True, # Use flash attention optimization + # 'q_layernorm': 'RMSNorm', # Use RMSNorm for query layer normalization + # 'k_layernorm': 'RMSNorm' # Use RMSNorm for key layer normalization + # } + # - golden_data_key: 'q8_flash_ql_kl' (reference output key) + # Expected Output: Model output should match reference data stored under 'q8_flash_ql_kl' + { + 'q_lora_rank': 8, + 'use_flash_attn': True, + 'q_layernorm': 'RMSNorm', + 'k_layernorm': 'RMSNorm' + }, + 'q8_flash_ql_kl' + ), + ( + # case 2 + # Input: + # - model_args: { + # 'q_lora_rank': 0, # Disable q_lora_rank for query projection + # 'use_flash_attn': True, # Use flash attention optimization + # 'q_layernorm': 'RMSNorm', # Use RMSNorm for query layer normalization + # 'k_layernorm': 'RMSNorm' # Use RMSNorm for key layer normalization + # } + # - golden_data_key: 'q0_flash_ql_kl' (reference output key) + # Expected Output: Model output should match reference data stored under 'q0_flash_ql_kl' + { + 'q_lora_rank': 0, + 'use_flash_attn': True, + 'q_layernorm': 'RMSNorm', + 'k_layernorm': 'RMSNorm' + }, + 'q0_flash_ql_kl' + ), + ( + # case 3 + # Input: + # - model_args: { + # 'q_lora_rank': 8, # Enable q_lora_rank with rank 8 for query projection + # 'use_flash_attn': True, # Use flash attention optimization + # 'q_layernorm': None, # Disable query layer normalization + # 'k_layernorm': 'RMSNorm' # Use RMSNorm for key layer normalization + # } + # - golden_data_key: 'q8_flash_kl' (reference output key) + # Expected Output: Model output should match reference data stored under 'q8_flash_kl' + { + 'q_lora_rank': 8, + 'use_flash_attn': True, + 'q_layernorm': None, + 'k_layernorm': 'RMSNorm' + }, + 'q8_flash_kl' + ) +] + +MULTI_CARD_TEST_PARAM = 'model_args, golden_data_key, tensor_parallel' +MULTI_CARD_TEST_CASES = [ + ( + # case 1 + # Input: + # - model_args: { + # 'q_lora_rank': 8, # Enable q_lora_rank with rank 8 for query projection + # 'use_flash_attn': True, # Use flash attention optimization + # 'q_layernorm': 'RMSNorm', # Use RMSNorm for query layer normalization + # 'k_layernorm': 'RMSNorm' # Use RMSNorm for key layer normalization + # } + # - golden_data_key: 'multi_q8_flash_ql_kl' (reference output key for multi-card scenario) + # - tensor_parallel: 2 (use 2 GPUs with tensor parallelism) + # Expected Output: Model output from two structures should match with each other + { + 'q_lora_rank': 8, + 'use_flash_attn': True, + 'q_layernorm': 'RMSNorm', + 'k_layernorm': 'RMSNorm' + }, + 'multi_q8_flash_ql_kl', + 2 + ), +] + + +def get_free_port(): + """Getting a random free port.""" + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + s.bind(('0.0.0.0', 0)) + _, port = s.getsockname() + time.sleep(0.05) + return port + + +def build_msrun_command_list( + worker_num, local_worker_num, log_dir, run_script_path, struct, + batch_size, seq_length, hidden_size, # Input shape args + model_args, # Dictionary of model config args + output_path_param, tensor_parallel, + test_name +): + """ Build the msrun command for Multi-head Latent Attention. """ + if worker_num == 1: + cmd_list = [ + "python", + ] + else: + cmd_list = [ + "msrun", + f"--worker_num={worker_num}", + f"--local_worker_num={local_worker_num}", + f"--master_port={get_free_port()}", + f"--log_dir={log_dir}", + "--join=True" + ] + + cmd_list.extend([ + str(run_script_path), + f"--struct={struct}", + f"--batch_size={batch_size}", + f"--seq_length={seq_length}", + f"--hidden_size={hidden_size}", + f"--q_lora_rank={model_args['q_lora_rank']}", + f"--use_flash_attn={model_args['use_flash_attn']}", + f"--q_layernorm={model_args['q_layernorm']}", + f"--k_layernorm={model_args['k_layernorm']}", + f"--output_path={output_path_param}", + f"--tensor_parallel={tensor_parallel}", + f"--test_name={test_name}" + ]) + + logger.info(f"Test case shell command: {' '.join(cmd_list)}") + return cmd_list + + +class TestMultiLatentAttention: + """Test class for Multi-head Latent Attention""" + LOG_DIR_NAME = "msrun_mla_log" + WORKER_LOG_FILENAME = "worker_0.log" + + def setup_method(self): + """Setup method to prepare test environment""" + self.current_dir = Path(__file__).parent.resolve() + self.run_script_path = self.current_dir / "run_multi_latent_attention.py" + + def check_acc(self, output_ms_dict, golden_data_key): + """ + Compare output_ms with GOLDEN_DATA and GPU_DATA. + """ + if not golden_data_key: # Should not happen if expect_error is False + logger.warning("No golden_data_key provided for accuracy check.") + return + + dtype_str = "bfloat16" + if "fp16" in golden_data_key: + dtype_str = "float16" + + standard = DoubleBenchmarkStandard(dtype=dtype_str) + + npu_data = output_ms_dict.get("output") + golden_data = GOLDEN_DATA.get(golden_data_key) + gpu_data = GPU_DATA.get(golden_data_key) + + DoubleBenchmarkComparator.check_pass_or_not( + npu_data=npu_data, + gpu_data=gpu_data, + golden_data=golden_data, + standard=standard + ) + + def run_mla_test( + self, + worker_num, + local_worker_num, + struct, + model_args, + golden_data_key, + tmp_path, + tensor_parallel=1, + ): + """Helper function to run MLA test and check results""" + output_file_path = tmp_path / f"output_mla_ms_{struct}_{'single' if tensor_parallel == 1 else 'multi'}.npz" + log_dir_path = tmp_path / self.LOG_DIR_NAME + log_dir_path.mkdir(parents=True, exist_ok=True) + worker_log_file = log_dir_path / self.WORKER_LOG_FILENAME + + cmd_list = build_msrun_command_list( + worker_num=worker_num, + local_worker_num=local_worker_num, + log_dir=log_dir_path, + run_script_path=self.run_script_path, + struct=struct, + batch_size=BATCH_SIZE, + seq_length=SEQ_LENGTH, + hidden_size=HIDDEN_SIZE, + model_args=model_args, + output_path_param=output_file_path, + tensor_parallel=tensor_parallel, + test_name=golden_data_key + ) + env = os.environ.copy() + env["PYTHONHASHSEED"] = "0" + result = subprocess.run(cmd_list, shell=False, capture_output=True, text=True, check=False, env=env) + + if worker_num != 1: + assert worker_log_file.exists() + assert result.returncode == 0, ( + "Multi-head Latent Attention script failed with non-zero exit code: " + f"{result.returncode}.\nStdout:\n{result.stdout}\nStderr:\n{result.stderr}" + ) + assert output_file_path.exists(), ( + f"Output file {output_file_path} was not created." + ) + + if tensor_parallel == 1: + output_ms_dict = np.load(output_file_path) + self.check_acc(output_ms_dict, golden_data_key) + else: + if struct == 'a2': + prefix, mode_num = tmp_path.name.split('mode') + mega_output_dir = prefix + 'mode' + str(int(mode_num) - 1) + mega_output_file = tmp_path.parent / mega_output_dir / "output_mla_ms_megatron_multi.npz" + output_ms_dict_mega = np.load(mega_output_file, allow_pickle=True)['output'] + output_ms_dict_mind = np.load(output_file_path, allow_pickle=True)['output'] + assert np.allclose(output_ms_dict_mind, output_ms_dict_mega) + + @pytest.mark.level0 + @pytest.mark.platform_arm_ascend910b_training + @pytest.mark.env_onecard + @pytest.mark.parametrize('struct', ['megatron', 'a2']) + @pytest.mark.parametrize( + SINGLE_CARD_TEST_PARAM, + SINGLE_CARD_TEST_CASES + ) + def test_single_card_mla_cases( + self, + struct, + model_args, + golden_data_key, + tmp_path + ): + """Test Multi-head Latent Attention on single card with various configurations.""" + self.run_mla_test( + worker_num=1, + local_worker_num=1, + struct=struct, + model_args=model_args, + golden_data_key=golden_data_key, + tmp_path=tmp_path + ) + + @pytest.mark.level0 + @pytest.mark.platform_arm_ascend910b_training + @pytest.mark.env_single + @pytest.mark.parametrize('struct', ['megatron', 'a2']) + @pytest.mark.parametrize( + MULTI_CARD_TEST_PARAM, + MULTI_CARD_TEST_CASES + ) + def test_multi_card_mla_cases( + self, + struct, + model_args, + golden_data_key, + tensor_parallel, + tmp_path + ): + """Test Multi-head Latent Attention on multiple cards with various configurations.""" + self.run_mla_test( + worker_num=4, + local_worker_num=4, + struct=struct, + model_args=model_args, + golden_data_key=golden_data_key, + tensor_parallel=tensor_parallel, + tmp_path=tmp_path + ) -- Gitee