From 965ce5788d4a3cb431508b01a43bab931fd5a10e Mon Sep 17 00:00:00 2001 From: Lysandre Date: Fri, 29 Aug 2025 16:27:08 +0200 Subject: [PATCH 1/2] Isolate Ernie fixes --- .../ernie4_5_moe/modeling_ernie4_5_moe.py | 109 ++++++++---------- .../ernie4_5_moe/modular_ernie4_5_moe.py | 77 +++++++------ 2 files changed, 88 insertions(+), 98 deletions(-) diff --git a/src/transformers/models/ernie4_5_moe/modeling_ernie4_5_moe.py b/src/transformers/models/ernie4_5_moe/modeling_ernie4_5_moe.py index 36f1f25c32be..d4d45875b78b 100644 --- a/src/transformers/models/ernie4_5_moe/modeling_ernie4_5_moe.py +++ b/src/transformers/models/ernie4_5_moe/modeling_ernie4_5_moe.py @@ -29,7 +29,6 @@ from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub from ...masking_utils import create_causal_mask -from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update @@ -67,7 +66,7 @@ def __init__(self, config, intermediate_size=None): super().__init__() self.config = config self.hidden_size = config.hidden_size - self.intermediate_size = intermediate_size if intermediate_size is not None else config.intermediate_size + self.intermediate_size = config.intermediate_size if intermediate_size is None else intermediate_size self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.use_bias) self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.use_bias) @@ -288,6 +287,46 @@ def forward(self, hidden_states): return hidden_states + self.e_score_correction_bias.squeeze() +class Ernie4_5_MoeNaiveMoe(nn.ModuleList): + def __init__(self, config): + super().__init__() + self.top_k = config.num_experts_per_tok + self.num_experts = config.num_local_experts + for _ in range(self.num_experts): + self += [Ernie4_5_MoeMLP(config, config.moe_intermediate_size)] + + def forward(self, hidden_states, e_score_correction_bias, device_type): + # router_logits: (batch * sequence_length, n_experts) + with torch.autocast(device_type=device_type, enabled=False): # Force float32 + routing_weights = F.softmax(hidden_states, dim=1, dtype=torch.float) + _, selected_experts = torch.topk(routing_weights + e_score_correction_bias, self.top_k, dim=-1) + routing_weights = torch.gather(routing_weights, dim=-1, index=selected_experts) + routing_weights = routing_weights / torch.clamp( + routing_weights.sum(dim=-1, keepdim=True), min=self.norm_min + ) + routing_weights = routing_weights.to(hidden_states.dtype) + + final_hidden_states = torch.zeros_like(hidden_states, device=hidden_states.device) + # One hot encode the selected experts to create an expert mask + # this will be used to easily index which expert is going to be sollicitated + expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=self.num_experts).permute(2, 1, 0) + + expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero() + for expert_idx in expert_hit: + expert_layer = self[expert_idx] + idx, top_x = torch.where(expert_mask[expert_idx].squeeze(0)) + # Index the correct hidden states and compute the expert hidden state for + # the current expert. We need to make sure to multiply the output hidden + # states by `routing_weights` on the corresponding tokens (top-1 and top-2) + current_state = hidden_states[None, top_x].reshape(-1, hidden_states.shape[-1]) + current_hidden_states = expert_layer(current_state) * routing_weights[top_x, idx, None] + + # However `index_add_` only support torch tensors for indexing so we'll use + # the `top_x` tensor here. + final_hidden_states.index_add_(0, top_x, current_hidden_states.to(hidden_states.dtype)) + return final_hidden_states + + class Ernie4_5_MoeSparseMoeBlock(nn.Module): """ This implementation is @@ -313,9 +352,7 @@ def __init__(self, config): # gating self.gate = nn.Linear(config.hidden_size, config.moe_num_experts, bias=False, dtype=torch.float32) - self.experts = nn.ModuleList( - [Ernie4_5_MoeMLP(config, config.moe_intermediate_size) for _ in range(config.moe_num_experts)] - ) + self.experts = Ernie4_5_MoeNaiveMoe(config) self.norm_min = config.moe_norm_min # (optional) shared experts for all forwards @@ -343,38 +380,10 @@ def forward( # router_logits: (batch * sequence_length, n_experts) router_logits = self.gate(hidden_states.float()) - routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float) - _, selected_experts = torch.topk(self.moe_statics(routing_weights), self.top_k, dim=-1) - routing_weights = torch.gather(routing_weights, dim=-1, index=selected_experts) - routing_weights = routing_weights / torch.clamp( - routing_weights.sum(dim=-1, keepdim=True), min=self.norm_min - ) - routing_weights = routing_weights.to(hidden_states.dtype) - - final_hidden_states = torch.zeros( - (batch_size * sequence_length, hidden_dim), dtype=hidden_states.dtype, device=hidden_states.device + final_hidden_states = self.experts( + hidden_states, self.moe_statics.e_score_correction_bias.squeeze(), device_type ) - # One hot encode the selected experts to create an expert mask - # this will be used to easily index which expert is going to be sollicitated - expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=self.num_experts).permute(2, 1, 0) - - # Loop over all available experts in the model and perform the computation on each expert - expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero() - for expert_idx in expert_hit: - expert_layer = self.experts[expert_idx] - idx, top_x = torch.where(expert_mask[expert_idx].squeeze(0)) - - # Index the correct hidden states and compute the expert hidden state for - # the current expert. We need to make sure to multiply the output hidden - # states by `routing_weights` on the corresponding tokens (top-1 and top-2) - current_state = hidden_states[None, top_x].reshape(-1, hidden_dim) - current_hidden_states = expert_layer(current_state) * routing_weights[top_x, idx, None] - - # However `index_add_` only support torch tensors for indexing so we'll use - # the `top_x` tensor here. - final_hidden_states.index_add_(0, top_x, current_hidden_states.to(hidden_states.dtype)) - # Add (optional) shared experts to the result if self.shared_experts is not None: final_hidden_states = final_hidden_states + shared_output @@ -411,32 +420,8 @@ def forward( position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[tuple[torch.Tensor]] = None, cache_position: Optional[torch.LongTensor] = None, - **kwargs: Unpack[FlashAttentionKwargs], + **kwargs: Unpack[TransformersKwargs], ) -> torch.FloatTensor: - """ - Args: - hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` - attention_mask (`torch.FloatTensor`, *optional*): attention mask of size - `(batch, sequence_length)` where padding elements are indicated by 0. - output_attentions (`bool`, *optional*): - Whether or not to return the attentions tensors of all attention layers. See `attentions` under - returned tensors for more detail. - output_router_logits (`bool`, *optional*): - Whether or not to return the logits of all the routers. They are useful for computing the router loss, - and should not be returned during inference. - use_cache (`bool`, *optional*): - If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding - (see `past_key_values`). - past_key_values (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states - cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): - Indices depicting the position of the input sequence tokens in the sequence. - position_embeddings (`tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*): - Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`, - with `head_dim` being the embedding dimension of each attention head. - kwargs (`dict`, *optional*): - Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code - into the model - """ residual = hidden_states hidden_states = self.input_layernorm(hidden_states) @@ -457,11 +442,9 @@ def forward( residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = self.mlp(hidden_states) - # For the MoE layers, we need to unpack if isinstance(hidden_states, tuple): - hidden_states, _ = hidden_states + hidden_states = hidden_states[0] hidden_states = residual + hidden_states - return hidden_states diff --git a/src/transformers/models/ernie4_5_moe/modular_ernie4_5_moe.py b/src/transformers/models/ernie4_5_moe/modular_ernie4_5_moe.py index 754193330e6f..815887a0592e 100644 --- a/src/transformers/models/ernie4_5_moe/modular_ernie4_5_moe.py +++ b/src/transformers/models/ernie4_5_moe/modular_ernie4_5_moe.py @@ -96,7 +96,46 @@ def forward(self, hidden_states): return hidden_states + self.e_score_correction_bias.squeeze() -# FIXME: update the expert +class Ernie4_5_MoeNaiveMoe(nn.ModuleList): + def __init__(self, config): + super().__init__() + self.top_k = config.num_experts_per_tok + self.num_experts = config.num_local_experts + for _ in range(self.num_experts): + self += [Ernie4_5_MoeMLP(config, config.moe_intermediate_size)] + + def forward(self, hidden_states, e_score_correction_bias, device_type): + # router_logits: (batch * sequence_length, n_experts) + with torch.autocast(device_type=device_type, enabled=False): # Force float32 + routing_weights = F.softmax(hidden_states, dim=1, dtype=torch.float) + _, selected_experts = torch.topk(routing_weights + e_score_correction_bias, self.top_k, dim=-1) + routing_weights = torch.gather(routing_weights, dim=-1, index=selected_experts) + routing_weights = routing_weights / torch.clamp( + routing_weights.sum(dim=-1, keepdim=True), min=self.norm_min + ) + routing_weights = routing_weights.to(hidden_states.dtype) + + final_hidden_states = torch.zeros_like(hidden_states, device=hidden_states.device) + # One hot encode the selected experts to create an expert mask + # this will be used to easily index which expert is going to be sollicitated + expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=self.num_experts).permute(2, 1, 0) + + expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero() + for expert_idx in expert_hit: + expert_layer = self[expert_idx] + idx, top_x = torch.where(expert_mask[expert_idx].squeeze(0)) + # Index the correct hidden states and compute the expert hidden state for + # the current expert. We need to make sure to multiply the output hidden + # states by `routing_weights` on the corresponding tokens (top-1 and top-2) + current_state = hidden_states[None, top_x].reshape(-1, hidden_states.shape[-1]) + current_hidden_states = expert_layer(current_state) * routing_weights[top_x, idx, None] + + # However `index_add_` only support torch tensors for indexing so we'll use + # the `top_x` tensor here. + final_hidden_states.index_add_(0, top_x, current_hidden_states.to(hidden_states.dtype)) + return final_hidden_states + + class Ernie4_5_MoeSparseMoeBlock(nn.Module): """ This implementation is @@ -122,9 +161,7 @@ def __init__(self, config): # gating self.gate = nn.Linear(config.hidden_size, config.moe_num_experts, bias=False, dtype=torch.float32) - self.experts = nn.ModuleList( - [Ernie4_5_MoeMLP(config, config.moe_intermediate_size) for _ in range(config.moe_num_experts)] - ) + self.experts = Ernie4_5_MoeNaiveMoe(config) self.norm_min = config.moe_norm_min # (optional) shared experts for all forwards @@ -152,37 +189,7 @@ def forward( # router_logits: (batch * sequence_length, n_experts) router_logits = self.gate(hidden_states.float()) - routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float) - _, selected_experts = torch.topk(self.moe_statics(routing_weights), self.top_k, dim=-1) - routing_weights = torch.gather(routing_weights, dim=-1, index=selected_experts) - routing_weights = routing_weights / torch.clamp( - routing_weights.sum(dim=-1, keepdim=True), min=self.norm_min - ) - routing_weights = routing_weights.to(hidden_states.dtype) - - final_hidden_states = torch.zeros( - (batch_size * sequence_length, hidden_dim), dtype=hidden_states.dtype, device=hidden_states.device - ) - - # One hot encode the selected experts to create an expert mask - # this will be used to easily index which expert is going to be sollicitated - expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=self.num_experts).permute(2, 1, 0) - - # Loop over all available experts in the model and perform the computation on each expert - expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero() - for expert_idx in expert_hit: - expert_layer = self.experts[expert_idx] - idx, top_x = torch.where(expert_mask[expert_idx].squeeze(0)) - - # Index the correct hidden states and compute the expert hidden state for - # the current expert. We need to make sure to multiply the output hidden - # states by `routing_weights` on the corresponding tokens (top-1 and top-2) - current_state = hidden_states[None, top_x].reshape(-1, hidden_dim) - current_hidden_states = expert_layer(current_state) * routing_weights[top_x, idx, None] - - # However `index_add_` only support torch tensors for indexing so we'll use - # the `top_x` tensor here. - final_hidden_states.index_add_(0, top_x, current_hidden_states.to(hidden_states.dtype)) + final_hidden_states = self.experts(hidden_states, self.moe_statics.e_score_correction_bias.squeeze(), device_type) # Add (optional) shared experts to the result if self.shared_experts is not None: From d50af34dc36bdac313b8b65c557c71e6c7441586 Mon Sep 17 00:00:00 2001 From: Vasqu Date: Fri, 29 Aug 2025 16:39:05 +0200 Subject: [PATCH 2/2] fix moe --- .../ernie4_5_moe/modeling_ernie4_5_moe.py | 29 ++++++++++------- .../ernie4_5_moe/modular_ernie4_5_moe.py | 31 ++++++++++++------- 2 files changed, 36 insertions(+), 24 deletions(-) diff --git a/src/transformers/models/ernie4_5_moe/modeling_ernie4_5_moe.py b/src/transformers/models/ernie4_5_moe/modeling_ernie4_5_moe.py index d4d45875b78b..e770e8559d0a 100644 --- a/src/transformers/models/ernie4_5_moe/modeling_ernie4_5_moe.py +++ b/src/transformers/models/ernie4_5_moe/modeling_ernie4_5_moe.py @@ -290,16 +290,16 @@ def forward(self, hidden_states): class Ernie4_5_MoeNaiveMoe(nn.ModuleList): def __init__(self, config): super().__init__() - self.top_k = config.num_experts_per_tok - self.num_experts = config.num_local_experts + self.top_k = config.moe_k + self.num_experts = config.moe_num_experts + self.norm_min = config.moe_norm_min for _ in range(self.num_experts): - self += [Ernie4_5_MoeMLP(config, config.moe_intermediate_size)] + self += [Ernie4_5_MoeMLP(config)] - def forward(self, hidden_states, e_score_correction_bias, device_type): - # router_logits: (batch * sequence_length, n_experts) + def forward(self, hidden_states, routing_weights, routing_bias, device_type): with torch.autocast(device_type=device_type, enabled=False): # Force float32 - routing_weights = F.softmax(hidden_states, dim=1, dtype=torch.float) - _, selected_experts = torch.topk(routing_weights + e_score_correction_bias, self.top_k, dim=-1) + routing_weights = F.softmax(routing_weights, dim=1, dtype=torch.float) + _, selected_experts = torch.topk(routing_weights + routing_bias, self.top_k, dim=-1) routing_weights = torch.gather(routing_weights, dim=-1, index=selected_experts) routing_weights = routing_weights / torch.clamp( routing_weights.sum(dim=-1, keepdim=True), min=self.norm_min @@ -307,10 +307,12 @@ def forward(self, hidden_states, e_score_correction_bias, device_type): routing_weights = routing_weights.to(hidden_states.dtype) final_hidden_states = torch.zeros_like(hidden_states, device=hidden_states.device) + # One hot encode the selected experts to create an expert mask # this will be used to easily index which expert is going to be sollicitated expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=self.num_experts).permute(2, 1, 0) + # Loop over all available experts in the model and perform the computation on each expert expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero() for expert_idx in expert_hit: expert_layer = self[expert_idx] @@ -344,6 +346,7 @@ class Ernie4_5_MoeSparseMoeBlock(nn.Module): def __init__(self, config): super().__init__() + self.hidden_dim = config.hidden_size self.num_experts = config.moe_num_experts self.top_k = config.moe_k @@ -353,7 +356,6 @@ def __init__(self, config): # gating self.gate = nn.Linear(config.hidden_size, config.moe_num_experts, bias=False, dtype=torch.float32) self.experts = Ernie4_5_MoeNaiveMoe(config) - self.norm_min = config.moe_norm_min # (optional) shared experts for all forwards self.shared_experts = None @@ -364,8 +366,8 @@ def forward( self, hidden_states: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - batch_size, sequence_length, hidden_dim = hidden_states.shape - hidden_states = hidden_states.view(-1, hidden_dim) + batch_size, sequence_length, _ = hidden_states.shape + hidden_states = hidden_states.view(-1, self.hidden_dim) # (Optional) shared experts if self.shared_experts is not None: @@ -381,14 +383,17 @@ def forward( router_logits = self.gate(hidden_states.float()) final_hidden_states = self.experts( - hidden_states, self.moe_statics.e_score_correction_bias.squeeze(), device_type + hidden_states, + routing_weights=router_logits, + routing_bias=self.moe_statics.e_score_correction_bias.squeeze(), + device_type=device_type, ) # Add (optional) shared experts to the result if self.shared_experts is not None: final_hidden_states = final_hidden_states + shared_output - final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim) + final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, self.hidden_dim) return final_hidden_states, router_logits diff --git a/src/transformers/models/ernie4_5_moe/modular_ernie4_5_moe.py b/src/transformers/models/ernie4_5_moe/modular_ernie4_5_moe.py index 815887a0592e..697f710417a0 100644 --- a/src/transformers/models/ernie4_5_moe/modular_ernie4_5_moe.py +++ b/src/transformers/models/ernie4_5_moe/modular_ernie4_5_moe.py @@ -99,16 +99,16 @@ def forward(self, hidden_states): class Ernie4_5_MoeNaiveMoe(nn.ModuleList): def __init__(self, config): super().__init__() - self.top_k = config.num_experts_per_tok - self.num_experts = config.num_local_experts + self.top_k = config.moe_k + self.num_experts = config.moe_num_experts + self.norm_min = config.moe_norm_min for _ in range(self.num_experts): - self += [Ernie4_5_MoeMLP(config, config.moe_intermediate_size)] + self += [Ernie4_5_MoeMLP(config)] - def forward(self, hidden_states, e_score_correction_bias, device_type): - # router_logits: (batch * sequence_length, n_experts) + def forward(self, hidden_states, routing_weights, routing_bias, device_type): with torch.autocast(device_type=device_type, enabled=False): # Force float32 - routing_weights = F.softmax(hidden_states, dim=1, dtype=torch.float) - _, selected_experts = torch.topk(routing_weights + e_score_correction_bias, self.top_k, dim=-1) + routing_weights = F.softmax(routing_weights, dim=1, dtype=torch.float) + _, selected_experts = torch.topk(routing_weights + routing_bias, self.top_k, dim=-1) routing_weights = torch.gather(routing_weights, dim=-1, index=selected_experts) routing_weights = routing_weights / torch.clamp( routing_weights.sum(dim=-1, keepdim=True), min=self.norm_min @@ -116,10 +116,12 @@ def forward(self, hidden_states, e_score_correction_bias, device_type): routing_weights = routing_weights.to(hidden_states.dtype) final_hidden_states = torch.zeros_like(hidden_states, device=hidden_states.device) + # One hot encode the selected experts to create an expert mask # this will be used to easily index which expert is going to be sollicitated expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=self.num_experts).permute(2, 1, 0) + # Loop over all available experts in the model and perform the computation on each expert expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero() for expert_idx in expert_hit: expert_layer = self[expert_idx] @@ -153,6 +155,7 @@ class Ernie4_5_MoeSparseMoeBlock(nn.Module): def __init__(self, config): super().__init__() + self.hidden_dim = config.hidden_size self.num_experts = config.moe_num_experts self.top_k = config.moe_k @@ -162,7 +165,6 @@ def __init__(self, config): # gating self.gate = nn.Linear(config.hidden_size, config.moe_num_experts, bias=False, dtype=torch.float32) self.experts = Ernie4_5_MoeNaiveMoe(config) - self.norm_min = config.moe_norm_min # (optional) shared experts for all forwards self.shared_experts = None @@ -173,8 +175,8 @@ def forward( self, hidden_states: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - batch_size, sequence_length, hidden_dim = hidden_states.shape - hidden_states = hidden_states.view(-1, hidden_dim) + batch_size, sequence_length, _ = hidden_states.shape + hidden_states = hidden_states.view(-1, self.hidden_dim) # (Optional) shared experts if self.shared_experts is not None: @@ -189,13 +191,18 @@ def forward( # router_logits: (batch * sequence_length, n_experts) router_logits = self.gate(hidden_states.float()) - final_hidden_states = self.experts(hidden_states, self.moe_statics.e_score_correction_bias.squeeze(), device_type) + final_hidden_states = self.experts( + hidden_states, + routing_weights=router_logits, + routing_bias=self.moe_statics.e_score_correction_bias.squeeze(), + device_type=device_type, + ) # Add (optional) shared experts to the result if self.shared_experts is not None: final_hidden_states = final_hidden_states + shared_output - final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim) + final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, self.hidden_dim) return final_hidden_states, router_logits