diff --git a/ds_pynative.yaml b/ds_pynative.yaml index 5affff04d88d50dfe9d78007e46a2f043c7eaf97..79a0df4419840c79e6f49a067f85e32a5ce26b52 100644 --- a/ds_pynative.yaml +++ b/ds_pynative.yaml @@ -196,6 +196,11 @@ model: mscale: 1 mscale_all_dim: 1 rope_theta: 10000 + # mhc + mhc_residual: False + mhc_expansion_rate: 4 + mhc_sk_iter: 20 + mhc_gating_factor_init: 0.01 # moe config router_dense_type: "float32" gated_linear_unit: True diff --git a/mindformers/parallel_core/transformer_config.py b/mindformers/parallel_core/transformer_config.py index 59e850e281b69f605c89febeab053927874b59c9..bfaad5c33c8aa5b78c2418357e8a1ee72cfc6fe1 100644 --- a/mindformers/parallel_core/transformer_config.py +++ b/mindformers/parallel_core/transformer_config.py @@ -73,6 +73,18 @@ class TransformerConfig(ModelParallelConfig, MFModelConfig): apply_residual_connection_post_layernorm: bool = False """If True, uses the original BERT residule connection ordering.""" + mhc_residual: bool = False + """If True, uses the Manifold-Constrained Hyper-Connections(mHC) for the residual submodule.""" + + mhc_expansion_rate: int = 4 + """Expansion rate of mHC, only work when mhc_residual = True.""" + + mhc_sk_iter: int = 20 + """Sinkhorn-Knopp iteration of mHC, only work when mhc_residual = True.""" + + mhc_gating_factor_init: float = 0.01 + """Gating factor init alpha of mHC, only work when mhc_residual = True.""" + layernorm_epsilon: float = 1e-5 """Epsilon value for any LayerNorm operations.""" diff --git a/mindformers/parallel_core/transformer_config_utils.py b/mindformers/parallel_core/transformer_config_utils.py index 0089ab068e8d3d6b511feb94519ed3f955e0f41c..860cb84a3d217d0610dabc8b890c527f5c081781 100644 --- a/mindformers/parallel_core/transformer_config_utils.py +++ b/mindformers/parallel_core/transformer_config_utils.py @@ -278,6 +278,10 @@ COMMON_CONFIG_MAPPING = { "add_mlp_fc2_bias_linear": "add_mlp_fc2_bias_linear", "rotary_interleaved": "rotary_interleaved", "fp32_residual_connection": "fp32_residual_connection", + "mhc_residual": "mhc_residual", + "mhc_expansion_rate": "mhc_expansion_rate", + "mhc_sk_iter": "mhc_sk_iter", + "mhc_gating_factor_init": "mhc_gating_factor_init", "window_size": "window_size", "window_attn_skip_freq": "window_attn_skip_freq", "model_architecture": "model_architecture", diff --git a/mindformers/pynative/base_models/gpt/gpt_layer_specs.py b/mindformers/pynative/base_models/gpt/gpt_layer_specs.py index 18205cd2829c7577395dfbd3b33bb384770c7aa4..ebaa541f609b0da91c81bc7ea14b50b898a5713d 100644 --- a/mindformers/pynative/base_models/gpt/gpt_layer_specs.py +++ b/mindformers/pynative/base_models/gpt/gpt_layer_specs.py @@ -31,6 +31,7 @@ from mindformers.pynative.transformers.transformer_layer import TransformerLayer from mindformers.parallel_core.transformer_config import TransformerConfig from mindformers.parallel_core.utils.spec_utils import ModuleSpec from mindformers.pynative.base_models.gpt.moe_module_specs import get_moe_module_spec +from mindformers.pynative.layers.mhc_layers import get_mhc_module, MHCMappingSubmodules from mindformers.pynative.transformers.multi_latent_attention import MLASelfAttention, \ MLASelfAttentionSubmodules @@ -101,6 +102,24 @@ def get_gpt_layer_local_spec( self_attention=self_attention, pre_mlp_layernorm=get_norm_cls(normalization, fused_norm), mlp=mlp, + mhc_res_mapping=ModuleSpec( + module=get_mhc_module("res"), + submodules=MHCMappingSubmodules( + mhc_rmsnorm=get_norm_cls(normalization, fused_norm) + ) + ), + mhc_pre_mapping=ModuleSpec( + module=get_mhc_module("pre"), + submodules=MHCMappingSubmodules( + mhc_rmsnorm=get_norm_cls(normalization, fused_norm) + ) + ), + mhc_post_mapping=ModuleSpec( + module=get_mhc_module("post"), + submodules=MHCMappingSubmodules( + mhc_rmsnorm=get_norm_cls(normalization, fused_norm) + ) + ), ), ) @@ -119,8 +138,26 @@ def get_gpt_layer_local_spec( ), ), pre_mlp_layernorm=get_norm_cls(normalization, fused_norm), - mlp=mlp - ) + mlp=mlp, + mhc_res_mapping=ModuleSpec( + module=get_mhc_module("res"), + submodules=MHCMappingSubmodules( + mhc_rmsnorm=get_norm_cls(normalization, fused_norm) + ) + ), + mhc_pre_mapping=ModuleSpec( + module=get_mhc_module("pre"), + submodules=MHCMappingSubmodules( + mhc_rmsnorm=get_norm_cls(normalization, fused_norm) + ) + ), + mhc_post_mapping=ModuleSpec( + module=get_mhc_module("post"), + submodules=MHCMappingSubmodules( + mhc_rmsnorm=get_norm_cls(normalization, fused_norm) + ) + ), + ), ) diff --git a/mindformers/pynative/layers/mhc_layers.py b/mindformers/pynative/layers/mhc_layers.py new file mode 100644 index 0000000000000000000000000000000000000000..3c4f238ff777d56404cc0563504e2876e0ce42bd --- /dev/null +++ b/mindformers/pynative/layers/mhc_layers.py @@ -0,0 +1,149 @@ +# 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. +# ============================================================================ + +__all__ = ["get_mhc_module"] + +import numpy as np +from dataclasses import dataclass +from typing import Union + +import mindspore as ms +from mindspore import nn, Parameter, mint, Tensor, ops +from mindspore.ops import sigmoid +from mindformers.modules.layers import Linear +from mindformers.parallel_core.transformer_config import TransformerConfig +from mindformers.parallel_core.utils.spec_utils import ModuleSpec, build_module +from mindformers.tools.logger import logger + +@dataclass +class MHCMappingSubmodules: + """ + Configuration class for specifying the submodules of a manifold-constrained hyper-connections. + """ + mhc_rmsnorm: Union[ModuleSpec, type] = None + +class MHCMappingBase(nn.Cell): + def __init__(self, config: TransformerConfig, submodules: MHCMappingSubmodules) -> None: + super().__init__() + self.n_stream = config.mhc_expansion_rate + self.sk_iter = config.mhc_sk_iter + self.hidden_size = config.hidden_size + self.seq_length = config.seq_length + self.context_len = self.hidden_size * self.n_stream + self.linear_proj = Linear( + in_channels=self.context_len, + out_channels=self.n_stream, + compute_dtype=config.compute_dtype, + param_init_type=config.params_dtype, + has_bias=False + ) + self.alpha = Parameter(Tensor(config.mhc_gating_factor_init, dtype=ms.float32), name="mhc_alpha") + self.beta = Parameter(mint.zeros((1, 1, 1, self.n_stream))) + self.rms_norm = build_module( + submodules.mhc_rmsnorm, + dim = self.context_len, + eps = config.layernorm_epsilon + ) + self.sigmoid = sigmoid + + def mapping(self, x): + S, B, N, H = x.shape + flatten_x = x.reshape((S, B, 1, N*H)) # (S, B, 1, N*H) + norm_x = self.rms_norm(flatten_x) # (S, B, 1, N*H) + static_map_x = self.linear_proj(norm_x) # (S, B, 1, N) + dynamic_map_x = self.alpha * static_map_x + self.beta # (S, B, 1, N) + return dynamic_map_x + +class MHCResMapping(MHCMappingBase): + def __init__(self, config: TransformerConfig, submodules: MHCMappingSubmodules) -> None: + super().__init__( + config=config, + submodules=submodules + ) + self.res_mapping_size = self.n_stream * self.n_stream + self.linear_proj = Linear( + in_channels=self.context_len, + out_channels=self.res_mapping_size, + compute_dtype=config.compute_dtype, + param_init_type=config.params_dtype, + has_bias=False + ) + self.beta = Parameter(mint.zeros((1, 1, self.n_stream, self.n_stream))) + self.reshape = mint.reshape + + def mapping(self, x): + S, B, N, H = x.shape + flatten_x = x.reshape((S, B, 1, N*H)) # 可以提出去 # (S, B, 1, N*H) + norm_x = self.rms_norm(flatten_x) # (S, B, 1, N*H) + static_map_x = self.linear_proj(norm_x) # (S, B, 1, N*N) + dynamic_map_x = self.alpha * static_map_x.reshape((S, B, N, N)) + self.beta # (S, B, N, N) + return dynamic_map_x + + def sinkhorn_knopp(self, x, iter=20, eps=1e-8): + S, B, N, _ = x.shape # (S, B, N, N) + batch = B * S + x = x.reshape((batch, N, N)) # (S, B, N, N) -> (B*S, N, N) + u = mint.ones((batch, N)) # (B*S, N) + v = mint.ones((batch, N)) # (B*S, N) + X = mint.exp(x) + for it in range(iter): + v_temp = v.unsqueeze(2) + xv = mint.bmm(X, v_temp).squeeze(2) + u = 1.0 / (xv + eps) + + u_temp = u.unsqueeze(2) + xt_u = mint.bmm(X.transpose(1, 2), u_temp).squeeze(2) + v = 1.0 / (xt_u + eps) + + U = ops.diag_embed(u) + V = ops.diag_embed(v) + P = mint.bmm(mint.bmm(U, X), V) + P = P.reshape((S, B, N, N)) # (B*S, N, N) -> (S, B, N, N) + return P, u, v + + def construct(self, x): + map_x = self.mapping(x) + sk_out, _, _ = ops.stop_gradient(self.sinkhorn_knopp(map_x, self.sk_iter)) + out = sk_out.astype(x.dtype) @ x + return out + +class MHCPreMapping(MHCMappingBase): + def __init__(self, config: TransformerConfig, submodules: MHCMappingSubmodules) -> None: + super().__init__(config=config, submodules=submodules) + + def construct(self, x): + map_x = self.mapping(x) + sigmoid_x = self.sigmoid(map_x) + out = sigmoid_x.astype(x.dtype) @ x + return out.squeeze(axis=2) + +class MHCPostMapping(MHCMappingBase): + def __init__(self, config: TransformerConfig, submodules: MHCMappingSubmodules) -> None: + super().__init__(config=config, submodules=submodules) + + def construct(self, x): + unsqeeze_x = x.unsqueeze(2).repeat(1, 1, self.n_stream, 1) + map_x = self.mapping(unsqeeze_x) + sigmoid_x = 2 * self.sigmoid(map_x) + out = sigmoid_x.astype(x.dtype) @ unsqeeze_x + return out.squeeze(axis=2) + +def get_mhc_module(mhc_type): + if mhc_type == "res": + return MHCResMapping + elif mhc_type == "pre": + return MHCPreMapping + else: + return MHCPostMapping \ No newline at end of file diff --git a/mindformers/pynative/transformers/transformer_layer.py b/mindformers/pynative/transformers/transformer_layer.py index bd5a6caa348c27191ebb670fe084396c33053768..a4c9b54df0f444a5a91ea85a0c88fc5e11a472a1 100644 --- a/mindformers/pynative/transformers/transformer_layer.py +++ b/mindformers/pynative/transformers/transformer_layer.py @@ -30,6 +30,7 @@ class TransformerLayerSubmodules: pre_mlp_layernorm (Union[ModuleSpec, type]): Specification for the layer normalization before the MLP. mlp (Union[ModuleSpec, type]): Specification for the MLP in Dense layer. + mhc_mapping (Union[ModuleSpec, type]): Specification for the mHC layers. """ input_layernorm: Union[ModuleSpec, type] = IdentityOp @@ -40,6 +41,10 @@ class TransformerLayerSubmodules: pre_mlp_layernorm: Union[ModuleSpec, type] = IdentityOp mlp: Union[ModuleSpec, type] = IdentityOp + mhc_rmsnorm: Union[ModuleSpec, type] = IdentityOp + mhc_res_mapping: Union[ModuleSpec, type] = IdentityOp + mhc_pre_mapping: Union[ModuleSpec, type] = IdentityOp + mhc_post_mapping: Union[ModuleSpec, type] = IdentityOp class BaseTransformerLayer: @@ -86,7 +91,9 @@ class TransformerLayer(nn.Cell, BaseTransformerLayer): super().__init__() self.config = config self.apply_residual_connection_post_norm = config.apply_residual_connection_post_layernorm + self.apply_manifold_constrained_hyper_connections = config.mhc_residual self.hidden_dropout = config.hidden_dropout if hidden_dropout is None else hidden_dropout + self.batch_size = config.batch_size self.input_layernorm = build_module( submodules.input_layernorm, @@ -131,6 +138,13 @@ class TransformerLayer(nn.Cell, BaseTransformerLayer): self.mlp = build_module(submodules.mlp, config=self.config) + if self.apply_manifold_constrained_hyper_connections: + self.mhc_res_mapping = build_module(submodules.mhc_res_mapping, config=self.config) + self.mhc_pre_mapping = build_module(submodules.mhc_pre_mapping, config=self.config) + self.mhc_post_mapping = build_module(submodules.mhc_post_mapping, config=self.config) + self.mhc_expansion_rate = config.mhc_expansion_rate + self.mhc_sk_iter = config.mhc_sk_iter + # mlp_bda(BiasDropoutFusion) is not supported. self.hidden_states_dropout = Dropout(drop_prob=self.hidden_dropout) @@ -168,6 +182,15 @@ class TransformerLayer(nn.Cell, BaseTransformerLayer): """ # Note: context parameter is currently unused but kept for API compatibility. # It may be used in future cross-attention implementations. + def expand_hidden_states(hidden_states): + hidden_states = hidden_states.unsqueeze(2) + return hidden_states.repeat(1, 1, self.mhc_expansion_rate, 1) + + # Before layernorm + if self.apply_manifold_constrained_hyper_connections: + mhc_hidden_states = expand_hidden_states(hidden_states) + residual = self.mhc_res_mapping(mhc_hidden_states) + hidden_states = self.mhc_pre_mapping(mhc_hidden_states) # Layer norm at the beginning input_layernorm_output = self.input_layernorm(hidden_states) @@ -199,7 +222,19 @@ class TransformerLayer(nn.Cell, BaseTransformerLayer): # Dropout dropout_output = self.hidden_states_dropout(attention_output) - norm_input = self.add(residual, dropout_output) + # Add residual + if self.apply_manifold_constrained_hyper_connections: + mhc_out = self.mhc_post_mapping(dropout_output) + norm_input = self.add(residual, mhc_out) + else: + norm_input = self.add(residual, dropout_output) + + + # Before layernorm + if self.apply_manifold_constrained_hyper_connections: + mhc_hidden_states = expand_hidden_states(norm_input) + residual = self.mhc_res_mapping(mhc_hidden_states) + norm_input = self.mhc_pre_mapping(mhc_hidden_states) # Layer norm post the self attention pre_mlp_layernorm_output = self.pre_mlp_layernorm(norm_input) @@ -217,8 +252,13 @@ class TransformerLayer(nn.Cell, BaseTransformerLayer): # Dropout dropout_output = self.hidden_states_dropout(mlp_output) - - output = self.add(residual, dropout_output) + + # Add residual + if self.apply_manifold_constrained_hyper_connections: + mhc_out = self.mhc_post_mapping(dropout_output) + output = self.add(residual, mhc_out) + else: + output = self.add(residual, dropout_output) # Note: context parameter is returned for API compatibility but currently unused. # It may be deprecated in future versions. return output, context