diff --git a/ds_pynative.yaml b/ds_pynative.yaml index 79a0df4419840c79e6f49a067f85e32a5ce26b52..5da787bb7c4636c5ae7dc44f3454fa6ae3b76d41 100644 --- a/ds_pynative.yaml +++ b/ds_pynative.yaml @@ -95,7 +95,7 @@ train_dataset: &train_dataset pad: -1 # The token id of `pad` in the dataset. data_path: # Megatron dataset sampling ratio and path. - '1' - - "/home/l00913161/data/deepseek-datasets/mmap_deepseekv3_datasets_text_document" + - "/home/w00932055/dsv4/deepseek-datasets/mmap_deepseekv3_datasets_text_document" input_columns: ["input_ids", "labels", "loss_mask", "position_ids"] construct_args_key: ["input_ids", "labels", "loss_mask", "position_ids"] num_parallel_workers: 8 @@ -202,6 +202,8 @@ model: mhc_sk_iter: 20 mhc_gating_factor_init: 0.01 # moe config + hash_routed_layer: 2 + hash_size: 3 # feasible when hash_router_layer > 0 router_dense_type: "float32" gated_linear_unit: True moe_intermediate_size: 2048 @@ -244,4 +246,4 @@ profile_start_step: 5 profile_stop_step: 7 init_start_profile: False profile_communication: False -profile_memory: True \ No newline at end of file +profile_memory: True diff --git a/mindformers/parallel_core/transformer_config.py b/mindformers/parallel_core/transformer_config.py index bfaad5c33c8aa5b78c2418357e8a1ee72cfc6fe1..b8126c9729f79b3c122539ca6490c70cdcca35a4 100644 --- a/mindformers/parallel_core/transformer_config.py +++ b/mindformers/parallel_core/transformer_config.py @@ -219,6 +219,12 @@ class TransformerConfig(ModelParallelConfig, MFModelConfig): # MoE #################### + hash_routed_layer: Optional[int] = None + """Number of hash-routed MoE layers""" + + hash_size: Optional[int] = None + """Number of experts to be routed by hash functions""" + moe_shared_expert_intermediate_size: Optional[int] = None """ Shared expert total ffn hidden size. diff --git a/mindformers/parallel_core/transformer_config_utils.py b/mindformers/parallel_core/transformer_config_utils.py index 860cb84a3d217d0610dabc8b890c527f5c081781..77ea9c942edaa395b8a57b700026ae57e72ae6ce 100644 --- a/mindformers/parallel_core/transformer_config_utils.py +++ b/mindformers/parallel_core/transformer_config_utils.py @@ -287,6 +287,8 @@ COMMON_CONFIG_MAPPING = { "model_architecture": "model_architecture", "num_encoder_layers": "num_encoder_layers", "num_decoder_layers": "num_decoder_layers", + "hash_routed_layer": "hash_routed_layer", + "hash_size": "hash_size", # Flash Attention # not changes diff --git a/mindformers/pynative/base_models/gpt/gpt_layer_specs.py b/mindformers/pynative/base_models/gpt/gpt_layer_specs.py index ebaa541f609b0da91c81bc7ea14b50b898a5713d..bea96b0222b1156baed56b47fbca4b438fd43a71 100644 --- a/mindformers/pynative/base_models/gpt/gpt_layer_specs.py +++ b/mindformers/pynative/base_models/gpt/gpt_layer_specs.py @@ -38,6 +38,7 @@ from mindformers.pynative.transformers.multi_latent_attention import MLASelfAtte def get_mlp_module_spec( num_experts: Optional[int] = None, moe_grouped_gemm: Optional[bool] = True, + hash_routed_experts: Optional[bool] = False, ) -> ModuleSpec: """Helper function to get module spec for MLP/MoE""" mlp = MLP @@ -53,6 +54,7 @@ def get_mlp_module_spec( return get_moe_module_spec( num_experts=num_experts, moe_grouped_gemm=moe_grouped_gemm, + hash_routed_experts=hash_routed_experts, ) def get_gpt_layer_local_spec( @@ -62,6 +64,7 @@ def get_gpt_layer_local_spec( multi_latent_attention: Optional[bool] = False, fused_norm: Optional[bool] = True, normalization: Optional[str] = "RMSNorm", + hash_routed_experts: Optional[bool] = False, ) -> ModuleSpec: """Use this spec for an implementation using only modules in Megatron-Core. @@ -80,6 +83,7 @@ def get_gpt_layer_local_spec( mlp = get_mlp_module_spec( num_experts=num_experts, moe_grouped_gemm=moe_grouped_gemm, + hash_routed_experts=hash_routed_experts, ) if multi_latent_attention: @@ -182,13 +186,27 @@ def get_gpt_decoder_block_spec( multi_latent_attention=config.multi_latent_attention, fused_norm=config.fused_norm, ) + + hash_moe_layer_spec = get_gpt_layer_local_spec( + num_experts=config.num_moe_experts, + moe_grouped_gemm=config.moe_grouped_gemm, + qk_layernorm=config.qk_layernorm, + multi_latent_attention=config.multi_latent_attention, + fused_norm=config.fused_norm, + hash_routed_experts=True, + ) + # Parse config.moe_layer_freq to determine the pattern of expert/dense layers. - # 0 stands for dense layers, 1 stands for expert layers. + # 0 stands for dense layers, 1 stands for hash routed expert layers, 2 stands for expert layers. # For integer N: Creates a pattern with one expert layer every N layers. # For string pattern: Evaluates the str directly (e.g. "[1,0,1]" for alternating expert/dense). if config.first_k_dense_replace: - moe_layer_pattern = [0] * config.first_k_dense_replace + \ - [1] * (config.num_layers - config.first_k_dense_replace) + dense_layer_num = config.first_k_dense_replace + hash_routed_moe_layer_num = config.hash_routed_layer + moe_layer_num = config.num_layers - dense_layer_num - hash_routed_moe_layer_num + moe_layer_pattern = [0] * dense_layer_num + \ + [1] * hash_routed_moe_layer_num + \ + [2] * moe_layer_num elif isinstance(config.moe_layer_freq, int): moe_layer_pattern = [1 if (i % config.moe_layer_freq == 0) else 0 for i in range(config.num_layers)] elif isinstance(config.moe_layer_freq, list): @@ -205,8 +223,10 @@ def get_gpt_decoder_block_spec( # Create the layer specs for the model. layer_specs = [] for layer_number in range(config.num_layers): - if moe_layer_pattern[layer_number] == 1: + if moe_layer_pattern[layer_number] == 2: layer_specs.append(moe_layer_spec) + elif moe_layer_pattern[layer_number] == 1: + layer_specs.append(hash_moe_layer_spec) elif moe_layer_pattern[layer_number] == 0: layer_specs.append(dense_layer_spec) else: diff --git a/mindformers/pynative/base_models/gpt/gpt_model.py b/mindformers/pynative/base_models/gpt/gpt_model.py index 95cefef4f502213737ed5d8483a6501145e528e6..89727bb2c85b6011955f2fafe0b60994960bc80f 100644 --- a/mindformers/pynative/base_models/gpt/gpt_model.py +++ b/mindformers/pynative/base_models/gpt/gpt_model.py @@ -353,7 +353,8 @@ class GPTModel(nn.Cell): attn_mask, rotary_pos_emb, prefix_keys_values, - actual_seq_len + actual_seq_len, + input_ids, ) return hidden_states, rotary_pos_emb diff --git a/mindformers/pynative/base_models/gpt/moe_module_specs.py b/mindformers/pynative/base_models/gpt/moe_module_specs.py index 9d540addb1105615113aeb276e7493b2eb851231..c10742e82cd1775e1dceb449d1eadbd89b8017e8 100644 --- a/mindformers/pynative/base_models/gpt/moe_module_specs.py +++ b/mindformers/pynative/base_models/gpt/moe_module_specs.py @@ -19,13 +19,14 @@ from mindformers.pynative.layers.linear import Linear from mindformers.pynative.transformers.mlp import MLPSubmodules from mindformers.pynative.transformers.moe.shared_experts import SharedExpertMLP from mindformers.parallel_core.utils.spec_utils import ModuleSpec -from mindformers.pynative.transformers.moe.moe_layer import MoELayer +from mindformers.pynative.transformers.moe.moe_layer import MoELayer, HashRoutedMoELayer from mindformers.pynative.transformers.moe.experts import GroupedMLP def get_moe_module_spec( num_experts: Optional[int] = None, moe_grouped_gemm: Optional[bool] = False, + hash_routed_experts: Optional[bool] = False, ) -> ModuleSpec: """Helper function to get module spec for MoE""" if num_experts is None: @@ -37,17 +38,22 @@ def get_moe_module_spec( shared_experts = SharedExpertMLP - moe_module_spec = ModuleSpec( - module=MoELayer, - # submodules=MoESubmodules( - # experts=GroupedMLP, - # shared_experts=ModuleSpec( - # module=shared_experts, - # submodules=MLPSubmodules( - # linear_fc1=Linear, - # linear_fc2=Linear - # ), - # ) - # ) - ) + if hash_routed_experts: + moe_module_spec = ModuleSpec( + module=HashRoutedMoELayer, + ) + else: + moe_module_spec = ModuleSpec( + module=MoELayer, + # submodules=MoESubmodules( + # experts=GroupedMLP, + # shared_experts=ModuleSpec( + # module=shared_experts, + # submodules=MLPSubmodules( + # linear_fc1=Linear, + # linear_fc2=Linear + # ), + # ) + # ) + ) return moe_module_spec diff --git a/mindformers/pynative/transformers/mlp.py b/mindformers/pynative/transformers/mlp.py index 20956f7241544a3446bfee3f8771e577dd36a6de..478778d4289876b0c5042779bb6a83ad4e371eda 100644 --- a/mindformers/pynative/transformers/mlp.py +++ b/mindformers/pynative/transformers/mlp.py @@ -104,7 +104,7 @@ class MLP(nn.Cell): self.add = mint.add self.transpose = mint.transpose - def construct(self, hidden_states: Tensor) -> tuple[Tensor, Tensor, float]: + def construct(self, hidden_states: Tensor, input_ids: Tensor = None) -> tuple[Tensor, Tensor, float]: """ Construct function of mlp block. """ # [seq_len, bs, hidden_size] -> [seq_len, bs, ffn_hidden_size] intermediate_parallel, bias_parallel = self.linear_fc1(hidden_states) diff --git a/mindformers/pynative/transformers/moe/experts.py b/mindformers/pynative/transformers/moe/experts.py index 54c452272500720fead31ce4b2531af81f1887c5..142da70905c4c0a295be6250f82cc25234e68f09 100644 --- a/mindformers/pynative/transformers/moe/experts.py +++ b/mindformers/pynative/transformers/moe/experts.py @@ -226,3 +226,11 @@ class GroupedMLP(nn.Cell): fc2_output = GroupedMatmul(split_item=3, group_type=0)( [intermediate_parallel], [w2], None, None, None, None, None, tokens_per_expert)[0] return fc2_output + +class HashGroupedMLP(GroupedMLP): + """ + Extended GroupedMLP to support Multi-Hash (Concatenation) logic. + """ + def __init__(self, config: TransformerConfig): + super().__init__(config) + self.top_k = config.hash_size \ No newline at end of file diff --git a/mindformers/pynative/transformers/moe/moe_layer.py b/mindformers/pynative/transformers/moe/moe_layer.py index 2290d92f5c4fe36d0ed3acf98fba2e9c3a8cdaab..7197e98efef0c8db72d3eb0c43b58c38fd625436 100644 --- a/mindformers/pynative/transformers/moe/moe_layer.py +++ b/mindformers/pynative/transformers/moe/moe_layer.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================ -"""MoE Layer implementation.""" +"""MoE Layer implementation with Multi-Hash Support.""" from mindspore import nn, Tensor, mint, ops import mindspore as ms from mindspore.common.parameter import Parameter @@ -20,8 +20,8 @@ from mindspore.common.parameter import Parameter from mindformers.parallel_core.transformer_config import TransformerConfig from mindformers.pynative.layers.linear import Linear from mindformers.pynative.transformers.mlp import MLPSubmodules -from .router import TopKRouter -from .experts import GroupedMLP +from .router import TopKRouter, HashRouter +from .experts import GroupedMLP, HashGroupedMLP from .shared_experts import SharedExpertMLP @@ -85,7 +85,7 @@ class MoELayer(nn.Cell): self.cast = ops.cast self.ones_like = mint.ones_like - def construct(self, hidden_states: Tensor): + def construct(self, hidden_states: Tensor, input_ids: Tensor = None): """ Forward pass for MoELayer. Args: @@ -112,3 +112,48 @@ class MoELayer(nn.Cell): final_out = out_experts return final_out, None + +class HashRoutedMoELayer(MoELayer): + """ + MoE Layer that combines Hash Router, Grouped Experts, and Shared Experts. + """ + def __init__(self, config: TransformerConfig): + super().__init__(config) + + # Router + self.router = HashRouter(config) + + # Experts + self.experts = HashGroupedMLP(config) + + self.expert_bias = None + + def construct(self, hidden_states: Tensor, input_ids: Tensor = None): + """ + Forward pass for MoELayer. + Args: + hidden_states (Tensor): Input tensor of shape (bs, slen, dim) + input_ids (Tensor, optional): Input token IDs (bs, slen), required if using Hash Router. + """ + bs, slen, dim = hidden_states.shape + + if input_ids is None: + raise ValueError("`input_ids` cannot be None when using HashRouter.") + + input_ids_flat = self.reshape(input_ids, (-1,)) + top_scores, selected_experts_indices, _ = self.router(input_ids_flat) # (bs*seqlen, 1) + + routed_output = self.experts(hidden_states, top_scores, selected_experts_indices) + + shared_output = None + if self.shared_experts is not None: + shared_output, _ = self.shared_experts(hidden_states) + + out_experts = self.reshape(routed_output, (bs, slen, dim)) + + if shared_output is not None: + final_out = self.add(shared_output, out_experts) + else: + final_out = out_experts + + return final_out, None \ No newline at end of file diff --git a/mindformers/pynative/transformers/moe/router.py b/mindformers/pynative/transformers/moe/router.py index 8c8d469c38490e9a0b33d0e13961b5310cf7026d..b534693542b3e4ddbe7b667d685a6448f5c08bf5 100644 --- a/mindformers/pynative/transformers/moe/router.py +++ b/mindformers/pynative/transformers/moe/router.py @@ -7,6 +7,9 @@ """Mixture of Experts (MoE) modules for pynative mode.""" from typing import Tuple, Optional +import numpy as np + +import mindspore as ms from mindspore import nn, Tensor, mint, ops from mindspore.common.parameter import Parameter @@ -220,3 +223,27 @@ class TopKRouter(nn.Cell): ) return top_scores, selected_experts_indices, num_tokens_per_expert + +class HashRouter(nn.Cell): + def __init__(self, config) -> None: + super().__init__() + self.vocab_size = config.vocab_size + self.num_moe_experts = config.num_moe_experts + self.hash_size = config.hash_size + self.hash_map = Tensor(np.random.randint( + 0, self.num_moe_experts, size=(self.vocab_size, self.hash_size) + ), dtype=ms.int) + self.histc = mint.histc + + def construct(self, input_ids): + expert_indices = self.hash_map[input_ids] + + num_tokens_per_expert = self.histc( + expert_indices, + bins=self.num_moe_experts, + min=0, + max=self.num_moe_experts, + ) + router_probs = mint.ones_like(expert_indices, dtype=ms.float32) + + return router_probs, expert_indices, num_tokens_per_expert diff --git a/mindformers/pynative/transformers/transformer_block.py b/mindformers/pynative/transformers/transformer_block.py index 866ff0018c9b4a09af8fa7d9960663207ae9c3f8..e333374af29d4c105cb0852020de745c204a2a1a 100644 --- a/mindformers/pynative/transformers/transformer_block.py +++ b/mindformers/pynative/transformers/transformer_block.py @@ -136,7 +136,8 @@ class TransformerBlock(nn.Cell): attention_mask: Tensor, rotary_pos_emb: Tensor = None, prefix_keys_values=None, - actual_seq_len=None): + actual_seq_len=None, + input_ids: Tensor = None): """ Construct function of transformer block. @@ -148,6 +149,7 @@ class TransformerBlock(nn.Cell): prefix_keys_values (optional): List of prefix key-value tensors for each layer. Each element should be a tuple or list of (key, value) tensors. Default: None. actual_seq_len (optional): Actual sequence length for variable-length sequences. Default: None. + input_ids (optional): Input index, only required when using hash router. Default: None. Returns: Tuple[Tensor, Tensor]: A tuple containing: @@ -161,7 +163,8 @@ class TransformerBlock(nn.Cell): attention_mask, rotary_pos_emb=rotary_pos_emb, prefix_keys_values=prefix_kv, - actual_seq_len=actual_seq_len + actual_seq_len=actual_seq_len, + input_ids=input_ids, ) # final layernorm. diff --git a/mindformers/pynative/transformers/transformer_layer.py b/mindformers/pynative/transformers/transformer_layer.py index a4c9b54df0f444a5a91ea85a0c88fc5e11a472a1..769b3ad4c3f07f28d4616ef821e88004ee24c7df 100644 --- a/mindformers/pynative/transformers/transformer_layer.py +++ b/mindformers/pynative/transformers/transformer_layer.py @@ -157,7 +157,8 @@ class TransformerLayer(nn.Cell, BaseTransformerLayer): context=None, rotary_pos_emb=None, prefix_keys_values=None, - actual_seq_len=None + actual_seq_len=None, + input_ids=None, ): """ Perform a forward pass through the transformer layer. @@ -173,6 +174,7 @@ class TransformerLayer(nn.Cell, BaseTransformerLayer): rotary_pos_emb (Tensor, optional): Rotary positional embeddings. Default: None. prefix_keys_values (Tensor, optional): Prefix key-value cache for attention. Default: None. actual_seq_len (int, optional): Actual sequence length for variable-length sequences. Default: None. + input_ids (optional): Input index, only required when using hash router. Default: None. Returns: Tuple[Tensor, Tensor, float]: A tuple containing: @@ -245,7 +247,7 @@ class TransformerLayer(nn.Cell, BaseTransformerLayer): else: residual = norm_input - mlp_output, mlp_output_bias = self.mlp(pre_mlp_layernorm_output) + mlp_output, mlp_output_bias = self.mlp(pre_mlp_layernorm_output, input_ids) if mlp_output_bias is not None: mlp_output = self.add(mlp_output, mlp_output_bias)