[BugFix] fix RMSNorm rms_norm_esp (#2797)

* fix rms

* add vl

* fix

* add vl

* fix

* fix
This commit is contained in:
lizexu123
2025-07-10 20:02:24 +08:00
committed by GitHub
parent ce5adec877
commit 8c660a0dfb
8 changed files with 18 additions and 17 deletions

View File

@@ -84,6 +84,7 @@ class ModelConfig(PretrainedConfig):
head_dim: Optional[int] = None,
tie_word_embeddings: bool = False,
is_quantized: bool = False,
rms_norm_eps: float = 1e-5,
**kwargs,
):
super().__init__(**kwargs)
@@ -123,6 +124,7 @@ class ModelConfig(PretrainedConfig):
self.dtype = dtype
self.tie_word_embeddings = tie_word_embeddings
self.is_quantized = is_quantized
self.rms_norm_eps = rms_norm_eps
@dataclass

View File

@@ -288,14 +288,14 @@ class Ernie4_5_DecoderLayer(nn.Layer):
self.input_layernorm = RMSNorm(
fd_config,
hidden_size=fd_config.model_config.hidden_size,
eps=1e-5,
eps=fd_config.model_config.rms_norm_eps,
prefix=f"{prefix}.input_layernorm",
)
self.post_attention_layernorm = RMSNorm(
fd_config,
hidden_size=fd_config.model_config.hidden_size,
eps=1e-5,
eps=fd_config.model_config.rms_norm_eps,
prefix=f"{prefix}.post_attention_layernorm",
)
@@ -366,7 +366,7 @@ class Ernie4_5_Model(nn.Layer):
self.norm = RMSNorm(
fd_config,
hidden_size=fd_config.model_config.hidden_size,
eps=1e-5,
eps=fd_config.model_config.rms_norm_eps,
prefix=f"{fd_config.model_config.prefix_name}.norm",
)

View File

@@ -275,14 +275,14 @@ class Ernie4_5_MTPModel(nn.Layer):
self.enorm = RMSNorm(
fd_config,
hidden_size=fd_config.model_config.hidden_size,
eps=1e-5,
eps=fd_config.model_config.rms_norm_eps,
prefix="ernie.mtp_emb_norm.0",
)
self.hnorm = RMSNorm(
fd_config,
hidden_size=fd_config.model_config.hidden_size,
eps=1e-5,
eps=fd_config.model_config.rms_norm_eps,
prefix="ernie.mtp_hidden_norm.0",
)

View File

@@ -273,14 +273,14 @@ class Ernie4_5_VLDecoderLayer(nn.Layer):
self.input_layernorm = RMSNorm(
fd_config,
hidden_size=fd_config.model_config.hidden_size,
eps=1e-5,
eps=fd_config.model_config.rms_norm_eps,
prefix=f"{prefix}.input_layernorm",
)
self.post_attention_layernorm = RMSNorm(
fd_config,
hidden_size=fd_config.model_config.hidden_size,
eps=1e-5,
eps=fd_config.model_config.rms_norm_eps,
prefix=f"{prefix}.post_attention_layernorm",
)
@@ -358,7 +358,7 @@ class Ernie4_5_VLModel(nn.Layer):
self.norm = RMSNorm(
fd_config,
hidden_size=fd_config.model_config.hidden_size,
eps=1e-5,
eps=fd_config.model_config.rms_norm_eps,
prefix=f"{fd_config.model_config.prefix_name}.norm",
)

View File

@@ -161,14 +161,14 @@ class Qwen2DecoderLayer(nn.Layer):
self.input_layernorm = RMSNorm(
fd_config,
hidden_size=fd_config.model_config.hidden_size,
eps=1e-6,
eps=fd_config.model_config.rms_norm_eps,
prefix=f"{prefix}.input_layernorm",
)
self.post_attention_layernorm = RMSNorm(
fd_config,
hidden_size=fd_config.model_config.hidden_size,
eps=1e-6,
eps=fd_config.model_config.rms_norm_eps,
prefix=f"{prefix}.post_attention_layernorm",
)
@@ -248,7 +248,7 @@ class Qwen2Model(nn.Layer):
self.norm = RMSNorm(
fd_config,
hidden_size=fd_config.model_config.hidden_size,
eps=1e-5,
eps=fd_config.model_config.rms_norm_eps,
prefix=f"{fd_config.model_config.prefix_name}.norm",
)

View File

@@ -79,12 +79,12 @@ class Qwen3Attention(nn.Layer):
self.q_norm = RMSNorm(fd_config=fd_config,
hidden_size=fd_config.model_config.head_dim,
eps=1e-6,
eps=fd_config.model_config.rms_norm_eps,
prefix=f"{prefix}.q_norm",
begin_norm_axis=2)
self.k_norm = RMSNorm(fd_config=fd_config,
hidden_size=fd_config.model_config.head_dim,
eps=1e-6,
eps=fd_config.model_config.rms_norm_eps,
prefix=f"{prefix}.k_norm",
begin_norm_axis=2)
@@ -183,7 +183,7 @@ class Qwen3Model(nn.Layer):
self.norm = RMSNorm(
fd_config,
hidden_size=fd_config.model_config.hidden_size,
eps=1e-6,
eps=fd_config.model_config.rms_norm_eps,
prefix=f"{fd_config.model_config.prefix_name}.norm",
)

View File

@@ -121,12 +121,12 @@ class Qwen3Attention(nn.Layer):
self.q_norm = RMSNorm(fd_config,
hidden_size=self.head_dim,
eps=1e-6,
eps=fd_config.model_config.rms_norm_eps,
prefix=f"{prefix}.q_norm",
begin_norm_axis=2)
self.k_norm = RMSNorm(fd_config,
hidden_size=self.head_dim,
eps=1e-6,
eps=fd_config.model_config.rms_norm_eps,
prefix=f"{prefix}.k_norm",
begin_norm_axis=2)

View File

@@ -594,7 +594,6 @@ def initialize_fd_config(config_or_args) -> FDConfig:
model_config_dict, _ = ModelConfig.get_config_dict(config_or_args.model_name_or_path)
# Handle MoE related configs
if 'num_experts' in model_config_dict:
model_config_dict['moe_num_experts'] = model_config_dict.pop('num_experts')