From e6d153a4954f26780e38c771f6c5f6b190f41b04 Mon Sep 17 00:00:00 2001 From: yanglong_unimelb Date: Wed, 25 Feb 2026 17:12:32 +0800 Subject: [PATCH] [bugfix] mhc only apply on moe --- .../base_models/gpt/gpt_layer_specs.py | 56 ++--- mindformers/pynative/layers/mhc_layers.py | 195 +++++++++--------- .../transformers/transformer_layer.py | 105 +++++++--- 3 files changed, 201 insertions(+), 155 deletions(-) diff --git a/mindformers/pynative/base_models/gpt/gpt_layer_specs.py b/mindformers/pynative/base_models/gpt/gpt_layer_specs.py index b39502ff6..4f83f85f0 100644 --- a/mindformers/pynative/base_models/gpt/gpt_layer_specs.py +++ b/mindformers/pynative/base_models/gpt/gpt_layer_specs.py @@ -31,7 +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.layers.mhc_layers import get_mhc_module from mindformers.pynative.transformers.multi_latent_attention import MLASelfAttention, \ MLASelfAttentionSubmodules from mindformers.pynative.transformers.dsa import DSAttention, DSAttentionSubmodules, \ @@ -153,24 +153,15 @@ 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) - ) - ), + # mhc_rmsnorm_mla=get_norm_cls("RMSNorm", fused_norm), + mhc_rmsnorm_moe=get_norm_cls("RMSNorm", fused_norm), + mhc_reduce_rmsnorm=get_norm_cls("RMSNorm", fused_norm), + # mhc_res_mapping_mla=get_mhc_module("res"), + # mhc_pre_mapping_mla=get_mhc_module("pre"), + # mhc_post_mapping_mla=get_mhc_module("post"), + mhc_res_mapping_moe=get_mhc_module("res"), + mhc_pre_mapping_moe=get_mhc_module("pre"), + mhc_post_mapping_moe=get_mhc_module("post"), ), ) @@ -190,24 +181,15 @@ def get_gpt_layer_local_spec( ), 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) - ) - ), + # mhc_rmsnorm_mla=get_norm_cls("RMSNorm", fused_norm), + mhc_rmsnorm_moe=get_norm_cls("RMSNorm", fused_norm), + mhc_reduce_rmsnorm=get_norm_cls("RMSNorm", fused_norm), + # mhc_res_mapping_mla=get_mhc_module("res"), + # mhc_pre_mapping_mla=get_mhc_module("pre"), + # mhc_post_mapping_mla=get_mhc_module("post"), + mhc_res_mapping_moe=get_mhc_module("res"), + mhc_pre_mapping_moe=get_mhc_module("pre"), + mhc_post_mapping_moe=get_mhc_module("post"), ), ) diff --git a/mindformers/pynative/layers/mhc_layers.py b/mindformers/pynative/layers/mhc_layers.py index 3c4f238ff..325539d42 100644 --- a/mindformers/pynative/layers/mhc_layers.py +++ b/mindformers/pynative/layers/mhc_layers.py @@ -16,26 +16,13 @@ __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 mindspore import nn, Parameter, mint, Tensor, ops, float32 +# from mindformers.modules.layers import Linear +from mindspore.mint.nn 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: +class MHCPrePostMapping(nn.Cell): + def __init__(self, config: TransformerConfig, sigmoid_ratio: int = 1) -> None: super().__init__() self.n_stream = config.mhc_expansion_rate self.sk_iter = config.mhc_sk_iter @@ -43,104 +30,124 @@ class MHCMappingBase(nn.Cell): 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 + in_features=self.context_len, + out_features=self.n_stream, + bias=False, + dtype=float32 ) - self.alpha = Parameter(Tensor(config.mhc_gating_factor_init, dtype=ms.float32), name="mhc_alpha") + self.alpha = Parameter(Tensor(config.mhc_gating_factor_init)) 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 - + self.sigmoid_ratio = sigmoid_ratio + self.sigmoid = ops.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) + flatten_x = x.reshape((S, B, 1, N*H)).astype(float32) # (S, B, 1, N*H) + static_map_x = self.linear_proj(flatten_x).astype(x.dtype) # (S, B, 1, N) dynamic_map_x = self.alpha * static_map_x + self.beta # (S, B, 1, N) return dynamic_map_x + + def construct(self, x): + map_x = self.mapping(x) # (S, B, 1, N) + sigmoid_x = self.sigmoid_ratio * self.sigmoid(map_x) # (S, B, 1, N) + return sigmoid_x.astype(x.dtype) -class MHCResMapping(MHCMappingBase): - def __init__(self, config: TransformerConfig, submodules: MHCMappingSubmodules) -> None: - super().__init__( - config=config, - submodules=submodules - ) +class MHCPreMapping(MHCPrePostMapping): + def __init__(self, config: TransformerConfig) -> None: + super().__init__(config=config, sigmoid_ratio=1) + +class MHCPostMapping(MHCPrePostMapping): + def __init__(self, config: TransformerConfig) -> None: + super().__init__(config=config, sigmoid_ratio=2) + +class MHCResMapping(MHCPrePostMapping): + def __init__(self, config: TransformerConfig) -> None: + super().__init__(config=config) 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 + in_features=self.context_len, + out_features=self.res_mapping_size, + bias=False, + dtype=float32 ) - self.beta = Parameter(mint.zeros((1, 1, self.n_stream, self.n_stream))) + beta = np.eye(self.n_stream) + self.beta = Parameter(Tensor(beta[None, None, ...], dtype=float32)) self.reshape = mint.reshape def mapping(self, x): + """ + Compute mapping parameters for residual MHC. + Note: Input x should already be RMSNorm-normalized according to paper formula (7). + + Args: + x: Input tensor with shape (S, B, N, H), already normalized by RMSNorm + + Returns: + Mapping matrix with shape (S, B, N, N) + """ 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) + flatten_x = x.reshape((S, B, 1, N*H)) # (S, B, 1, N*H) + static_map_x = self.linear_proj(flatten_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 logsumexp(self, x, axis, keepdims=False): + """ + Numerically stable log-sum-exp computation. + + Args: + x: Input tensor + axis: Axis along which to compute log-sum-exp + keepdims: Whether to keep dimensions + + Returns: + log(sum(exp(x))) computed in a numerically stable way + """ + x_max, _ = ops.max(x, axis=axis, keepdims=True) + x_exp = mint.exp(x - x_max) + x_sumexp = ops.sum(x_exp, dim=axis, keepdim=True) + x_logsumexp = ops.log(x_sumexp + 1e-8) + result = x_logsumexp + x_max + if not keepdims: + result = ops.squeeze(result, axis=axis) + return result - 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) + def sinkhorn_knopp(self, h, iter=20): + """ + Log-domain Sinkhorn-Knopp algorithm for doubly-stochastic matrix projection. + This implementation operates entirely in log-space for better numerical stability. + + Args: + h: Input matrix with shape (S, B, N, N) - treated as log-domain values + iter: Number of Sinkhorn iterations + + Returns: + P: Doubly-stochastic matrix with shape (S, B, N, N) + """ + for _ in range(iter): + h = h - self.logsumexp(h, axis=3, keepdims=True) + h = h - self.logsumexp(h, axis=2, keepdims=True) + P = mint.exp(h) # (S, B, N, N) + return P -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) + map_x = self.mapping(x) # (S, B, N, N) + sk_out = ops.stop_gradient(self.sinkhorn_knopp(map_x, self.sk_iter)) + return sk_out def get_mhc_module(mhc_type): + """ + Get MHC module class based on type. + + Args: + mhc_type: "res", "pre", or "post" + + Returns: + Module class with correct sigmoid_ratio: + - res: MHCResMapping (uses Sinkhorn, no sigmoid) + - pre: MHCPreMapping (sigmoid_ratio=1, H^pre = σ(H̃^pre)) + - post: MHCPostMapping (sigmoid_ratio=2, H^post = 2σ(H̃^post)) + """ if mhc_type == "res": return MHCResMapping elif mhc_type == "pre": diff --git a/mindformers/pynative/transformers/transformer_layer.py b/mindformers/pynative/transformers/transformer_layer.py index bfb3d709d..40bb092ab 100644 --- a/mindformers/pynative/transformers/transformer_layer.py +++ b/mindformers/pynative/transformers/transformer_layer.py @@ -5,13 +5,15 @@ """Transformer Layer""" from dataclasses import dataclass from typing import Union -from mindspore import nn, mint +from mindspore import nn, mint, ops, float32, Parameter, Tensor from mindformers.parallel_core.utils.spec_utils import ModuleSpec, build_module from mindformers.parallel_core.transformer_config import TransformerConfig from mindformers.pynative.layers.dropout import Dropout from mindformers.pynative.layers.identity_op import IdentityOp from mindformers.pynative.transformers.moe.moe_layer import HashRoutedMoELayer +import numpy as np + @dataclass class TransformerLayerSubmodules: @@ -42,10 +44,15 @@ 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 + # mhc_rmsnorm_mla: Union[ModuleSpec, type] = IdentityOp + mhc_rmsnorm_moe: Union[ModuleSpec, type] = IdentityOp + mhc_reduce_rmsnorm: Union[ModuleSpec, type] = IdentityOp + # mhc_res_mapping_mla: Union[ModuleSpec, type] = IdentityOp + # mhc_pre_mapping_mla: Union[ModuleSpec, type] = IdentityOp + # mhc_post_mapping_mla: Union[ModuleSpec, type] = IdentityOp + mhc_res_mapping_moe: Union[ModuleSpec, type] = IdentityOp + mhc_pre_mapping_moe: Union[ModuleSpec, type] = IdentityOp + mhc_post_mapping_moe: Union[ModuleSpec, type] = IdentityOp class BaseTransformerLayer: @@ -95,6 +102,8 @@ class TransformerLayer(nn.Cell, BaseTransformerLayer): 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.layer_number = layer_number + self.bmm = ops.BatchMatMul() self.input_layernorm = build_module( submodules.input_layernorm, @@ -140,9 +149,28 @@ 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) + # RMSNorm dimension should be hidden_size * expansion_rate because input is expanded + # self.mhc_rmsnorm_mla = build_module( + # submodules.mhc_rmsnorm_mla, + # dim=config.hidden_size, + # eps=config.layernorm_epsilon + # ) + self.mhc_rmsnorm_moe = build_module( + submodules.mhc_rmsnorm_moe, + dim=config.hidden_size, + eps=config.layernorm_epsilon + ) + self.mhc_reduce_rmsnorm = build_module( + submodules.mhc_reduce_rmsnorm, + dim=config.hidden_size, + eps=config.layernorm_epsilon + ) + # self.mhc_res_mapping_mla = build_module(submodules.mhc_res_mapping_mla, config=self.config) + # self.mhc_pre_mapping_mla = build_module(submodules.mhc_pre_mapping_mla, config=self.config) + # self.mhc_post_mapping_mla = build_module(submodules.mhc_post_mapping_mla, config=self.config) + self.mhc_res_mapping_moe = build_module(submodules.mhc_res_mapping_moe, config=self.config) + self.mhc_pre_mapping_moe = build_module(submodules.mhc_pre_mapping_moe, config=self.config) + self.mhc_post_mapping_moe = build_module(submodules.mhc_post_mapping_moe, config=self.config) self.mhc_expansion_rate = config.mhc_expansion_rate self.mhc_sk_iter = config.mhc_sk_iter @@ -188,14 +216,24 @@ 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) + hidden_states = hidden_states.unsqueeze(axis=2) + return hidden_states.repeat(1, 1, self.mhc_expansion_rate, 1).astype(float32) + print(f"-------------- Transformer Layer {self.layer_number} ----------------") + print(f"@ Original hidden_states AMax: {ops.abs(hidden_states).max()} {ops.abs(hidden_states).mean()}") # 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) + # if self.apply_manifold_constrained_hyper_connections: + # mhc_hidden_states = expand_hidden_states(hidden_states).astype(float32) # (S, B, H) -> (S, B, N, H) + # mhc_hidden_states = self.mhc_rmsnorm_mla(mhc_hidden_states) + + # mhc_res = self.mhc_res_mapping_mla(mhc_hidden_states) # (S, B, N, N) + # mhc_pre = self.mhc_pre_mapping_mla(mhc_hidden_states) # (S, B, 1, N) + # mhc_post = self.mhc_post_mapping_mla(mhc_hidden_states) # (S, B, 1, N) + + # mhc_res_hidden = self.bmm(mhc_res, mhc_hidden_states) # (S, B, N, N) @ (S, B, N, H) -> (S, B, N, H) + # mhc_pre_hidden = self.bmm(mhc_pre, mhc_hidden_states) # (S, B, 1, N) @ (S, B, N, H) -> (S, B, 1, H) + # mhc_post = mhc_post.squeeze(axis=2).unsqueeze(-1) # (S, B, 1, N) -> (S, B, N) -> (S, B, N, 1) + # hidden_states = mhc_pre_hidden.squeeze(axis=2) # (S, B, 1, H) -> (S, B, H) # Layer norm at the beginning input_layernorm_output = self.input_layernorm(hidden_states) @@ -203,6 +241,8 @@ class TransformerLayer(nn.Cell, BaseTransformerLayer): # Residual connection if self.apply_residual_connection_post_norm: residual = input_layernorm_output + # elif self.apply_manifold_constrained_hyper_connections: + # residual = mhc_res_hidden else: residual = hidden_states @@ -227,19 +267,31 @@ class TransformerLayer(nn.Cell, BaseTransformerLayer): # Dropout dropout_output = self.hidden_states_dropout(attention_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) + # # Add residual + # if self.apply_manifold_constrained_hyper_connections: + # dropout_expand = dropout_output.unsqueeze(axis=2) # (S, B, H) -> (S, B, 1, H) + # mhc_post_hidden = self.bmm(mhc_post, dropout_expand) # (S, B, N, 1) @ (S, B, 1, H) -> (S, B, N, H) + # norm_input = self.add(residual, mhc_post_hidden) + # norm_input = norm_input.sum(dim=2) # (S, B, N, H) -> (S, B, H) + # else: + # norm_input = self.add(residual, dropout_output) + + 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) + mhc_hidden_states = self.mhc_rmsnorm_moe(mhc_hidden_states) + + mhc_res = self.mhc_res_mapping_moe(mhc_hidden_states) # (S, B, N, N) + mhc_pre = self.mhc_pre_mapping_moe(mhc_hidden_states) # (S, B, 1, N) + mhc_post = self.mhc_post_mapping_moe(mhc_hidden_states) # (S, B, 1, N) + + mhc_res_hidden = self.bmm(mhc_res, mhc_hidden_states) # (S, B, N, N) @ (S, B, N, H) -> (S, B, N, H) + mhc_pre_hidden = self.bmm(mhc_pre, mhc_hidden_states) # (S, B, 1, N) @ (S, B, N, H) -> (S, B, 1, H) + mhc_post = mhc_post.squeeze(axis=2).unsqueeze(-1) # (S, B, 1, N) -> (S, B, N) -> (S, B, N, 1) + norm_input = mhc_pre_hidden.squeeze(axis=2) # (S, B, 1, H) -> (S, B, H) # Layer norm post the self attention pre_mlp_layernorm_output = self.pre_mlp_layernorm(norm_input) @@ -247,6 +299,8 @@ class TransformerLayer(nn.Cell, BaseTransformerLayer): # Residual connection if self.apply_residual_connection_post_norm: residual = pre_mlp_layernorm_output + elif self.apply_manifold_constrained_hyper_connections: + residual = mhc_res_hidden else: residual = norm_input @@ -269,10 +323,13 @@ class TransformerLayer(nn.Cell, BaseTransformerLayer): # Add residual if self.apply_manifold_constrained_hyper_connections: - mhc_out = self.mhc_post_mapping(dropout_output) - output = self.add(residual, mhc_out) + dropout_expand = dropout_output.unsqueeze(axis=2) # (S, B, H) -> (S, B, 1, H) + mhc_post_hidden = mhc_post @ dropout_expand # (S, B, N, 1) @ (S, B, 1, H) -> (S, B, N, H) + output = self.add(residual, mhc_post_hidden) + output = self.mhc_reduce_rmsnorm(output.sum(dim=2)) # (S, B, N, H) -> (S, B, H) 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. + print(f"@ transformer layer output AMax: {ops.abs(output).max()} {ops.abs(output).mean()}") return output, context, aux_loss -- Gitee