Fix load_balancing_loss_func incompatible with past_key_values

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 = (

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 = (

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 = (

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 = (

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 = (

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 = (

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 = (