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 2976beba1033..7211e5a35310 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 @@ -623,8 +623,10 @@ def load_balancing_loss_func( # Compute the average probability of routing to these experts router_prob_per_expert = torch.mean(routing_weights, dim=0) else: - batch_size, sequence_length = attention_mask.shape - num_hidden_layers = concatenated_gate_logits.shape[0] // (batch_size * sequence_length) + num_hidden_layers = len(gate_logits) + batch_size = attention_mask.shape[0] + sequence_length = gate_logits[0].shape[0] // batch_size + attention_mask = attention_mask[:, -sequence_length:] # Compute the mask that masks all padding tokens as 0 with the same shape of expert_mask expert_attention_mask = ( diff --git a/src/transformers/models/gpt_oss/modeling_gpt_oss.py b/src/transformers/models/gpt_oss/modeling_gpt_oss.py index 0d5c936e8adc..f132693fb0b7 100644 --- a/src/transformers/models/gpt_oss/modeling_gpt_oss.py +++ b/src/transformers/models/gpt_oss/modeling_gpt_oss.py @@ -570,8 +570,10 @@ def load_balancing_loss_func( # Compute the average probability of routing to these experts router_prob_per_expert = torch.mean(routing_weights, dim=0) else: - batch_size, sequence_length = attention_mask.shape - num_hidden_layers = concatenated_gate_logits.shape[0] // (batch_size * sequence_length) + num_hidden_layers = len(gate_logits) + batch_size = attention_mask.shape[0] + sequence_length = gate_logits[0].shape[0] // batch_size + attention_mask = attention_mask[:, -sequence_length:] # Compute the mask that masks all padding tokens as 0 with the same shape of expert_mask expert_attention_mask = ( diff --git a/src/transformers/models/minimax/modeling_minimax.py b/src/transformers/models/minimax/modeling_minimax.py index 633e053e2d54..c37625a0a66c 100644 --- a/src/transformers/models/minimax/modeling_minimax.py +++ b/src/transformers/models/minimax/modeling_minimax.py @@ -773,8 +773,10 @@ def load_balancing_loss_func( # Compute the average probability of routing to these experts router_prob_per_expert = torch.mean(routing_weights, dim=0) else: - batch_size, sequence_length = attention_mask.shape - num_hidden_layers = concatenated_gate_logits.shape[0] // (batch_size * sequence_length) + num_hidden_layers = len(gate_logits) + batch_size = attention_mask.shape[0] + sequence_length = gate_logits[0].shape[0] // batch_size + attention_mask = attention_mask[:, -sequence_length:] # Compute the mask that masks all padding tokens as 0 with the same shape of expert_mask expert_attention_mask = ( diff --git a/src/transformers/models/mixtral/modeling_mixtral.py b/src/transformers/models/mixtral/modeling_mixtral.py index 2412092aeb86..5a54605bdc8f 100644 --- a/src/transformers/models/mixtral/modeling_mixtral.py +++ b/src/transformers/models/mixtral/modeling_mixtral.py @@ -529,8 +529,10 @@ def load_balancing_loss_func( # Compute the average probability of routing to these experts router_prob_per_expert = torch.mean(routing_weights, dim=0) else: - batch_size, sequence_length = attention_mask.shape - num_hidden_layers = concatenated_gate_logits.shape[0] // (batch_size * sequence_length) + num_hidden_layers = len(gate_logits) + batch_size = attention_mask.shape[0] + sequence_length = gate_logits[0].shape[0] // batch_size + attention_mask = attention_mask[:, -sequence_length:] # Compute the mask that masks all padding tokens as 0 with the same shape of expert_mask expert_attention_mask = ( diff --git a/src/transformers/models/mixtral/modular_mixtral.py b/src/transformers/models/mixtral/modular_mixtral.py index d897824c4cff..93771ce612be 100644 --- a/src/transformers/models/mixtral/modular_mixtral.py +++ b/src/transformers/models/mixtral/modular_mixtral.py @@ -101,8 +101,10 @@ def load_balancing_loss_func( # Compute the average probability of routing to these experts router_prob_per_expert = torch.mean(routing_weights, dim=0) else: - batch_size, sequence_length = attention_mask.shape - num_hidden_layers = concatenated_gate_logits.shape[0] // (batch_size * sequence_length) + num_hidden_layers = len(gate_logits) + batch_size = attention_mask.shape[0] + sequence_length = gate_logits[0].shape[0] // batch_size + attention_mask = attention_mask[:, -sequence_length:] # Compute the mask that masks all padding tokens as 0 with the same shape of expert_mask expert_attention_mask = ( diff --git a/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py b/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py index 2056e7c76a3a..f2e9da654de8 100644 --- a/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py +++ b/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py @@ -552,8 +552,10 @@ def load_balancing_loss_func( # Compute the average probability of routing to these experts router_prob_per_expert = torch.mean(routing_weights, dim=0) else: - batch_size, sequence_length = attention_mask.shape - num_hidden_layers = concatenated_gate_logits.shape[0] // (batch_size * sequence_length) + num_hidden_layers = len(gate_logits) + batch_size = attention_mask.shape[0] + sequence_length = gate_logits[0].shape[0] // batch_size + attention_mask = attention_mask[:, -sequence_length:] # Compute the mask that masks all padding tokens as 0 with the same shape of expert_mask expert_attention_mask = ( diff --git a/src/transformers/models/qwen3_next/modeling_qwen3_next.py b/src/transformers/models/qwen3_next/modeling_qwen3_next.py index 7d2b60d943e2..9db3fa0a17ac 100644 --- a/src/transformers/models/qwen3_next/modeling_qwen3_next.py +++ b/src/transformers/models/qwen3_next/modeling_qwen3_next.py @@ -1113,8 +1113,10 @@ def load_balancing_loss_func( # Compute the average probability of routing to these experts router_prob_per_expert = torch.mean(routing_weights, dim=0) else: - batch_size, sequence_length = attention_mask.shape - num_hidden_layers = concatenated_gate_logits.shape[0] // (batch_size * sequence_length) + num_hidden_layers = len(gate_logits) + batch_size = attention_mask.shape[0] + sequence_length = gate_logits[0].shape[0] // batch_size + attention_mask = attention_mask[:, -sequence_length:] # Compute the mask that masks all padding tokens as 0 with the same shape of expert_mask expert_attention_mask = (