Repository: huggingface/transformers
Pull request: another way of fix for #42400

# What does this PR do?

This was originally #40617, which IMO a better approach than #42400 (despite not ideal neither).

With current form, the failure ratio is 0.1% instead of 1%. We can add `@is_flaky(max_attempts=2)` as in #40617 however.

Changed files:
filename: tests/models/ernie4_5_moe/test_modeling_ernie4_5_moe.py
patch: @@ -89,15 +89,14 @@ def test_flash_attn_2_equivalence(self):
                 # higher tolerance, not sure where it stems from
                 assert torch.allclose(logits_fa, logits, atol=1e-2, rtol=1e-2)
 
-    # Ignore copy
-    @unittest.skip("TODO @ArthurZucker investigate later on")
+    @is_flaky(max_attempts=2)
     def test_load_balancing_loss(self):
         r"""
         Let's make sure we can actually compute the loss and do a backward on it.
         """
         config, input_dict = self.model_tester.prepare_config_and_inputs_for_common()
         config.num_labels = 3
-        config.num_experts = 8
+        config.num_experts = 3
         config.expert_interval = 2
         config.output_router_logits = True
         input_ids = input_dict["input_ids"]
@@ -111,7 +110,7 @@ def test_load_balancing_loss(self):
 
         # First, we make sure that adding padding tokens doesn't change the loss
         # loss(input_ids, attention_mask=None) == loss(input_ids + padding, attention_mask=attention_mask_with_padding)
-        pad_length = 1000
+        pad_length = input_ids.shape[1] * 4
         # Add padding tokens (assume that pad_token_id=1) to input_ids
         padding_block = torch.ones(input_ids.shape[0], pad_length, dtype=torch.int32).to(torch_device)
         padded_input_ids = torch.cat((padding_block, input_ids), dim=1)  # this is to simulate padding to the left

filename: tests/models/minimax/test_modeling_minimax.py
patch: @@ -18,6 +18,7 @@
 from transformers import is_torch_available
 from transformers.testing_utils import (
     Expectations,
+    is_flaky,
     require_torch,
     require_torch_accelerator,
     slow,
@@ -63,13 +64,14 @@ def is_pipeline_test_to_skip(
     ):
         return True
 
+    @is_flaky(max_attempts=2)
     def test_load_balancing_loss(self):
         r"""
         Let's make sure we can actually compute the loss and do a backward on it.
         """
         config, input_dict = self.model_tester.prepare_config_and_inputs_for_common()
         config.num_labels = 3
-        config.num_local_experts = 8
+        config.num_local_experts = 3
         config.output_router_logits = True
         input_ids = input_dict["input_ids"]
         attention_mask = input_ids.ne(1).to(torch_device)
@@ -82,7 +84,7 @@ def test_load_balancing_loss(self):
 
         # First, we make sure that adding padding tokens doesn't change the loss
         # loss(input_ids, attention_mask=None) == loss(input_ids + padding, attention_mask=attention_mask_with_padding)
-        pad_length = 1000
+        pad_length = input_ids.shape[1] * 4
         # Add padding tokens (assume that pad_token_id=1) to input_ids
         padding_block = torch.ones(input_ids.shape[0], pad_length, dtype=torch.int32).to(torch_device)
         padded_input_ids = torch.cat((padding_block, input_ids), dim=1)  # this is to simulate padding to the left

filename: tests/models/mixtral/test_modeling_mixtral.py
patch: @@ -20,6 +20,7 @@
 from transformers import is_torch_available
 from transformers.testing_utils import (
     Expectations,
+    is_flaky,
     require_flash_attn,
     require_torch,
     require_torch_accelerator,
@@ -68,14 +69,14 @@ def is_pipeline_test_to_skip(
     def test_flash_attn_2_inference_equivalence_right_padding(self):
         self.skipTest(reason="Mixtral flash attention does not support right padding")
 
-    # Ignore copy
+    @is_flaky(max_attempts=2)
     def test_load_balancing_loss(self):
         r"""
         Let's make sure we can actually compute the loss and do a backward on it.
         """
         config, input_dict = self.model_tester.prepare_config_and_inputs_for_common()
         config.num_labels = 3
-        config.num_local_experts = 8
+        config.num_local_experts = 3
         config.output_router_logits = True
         input_ids = input_dict["input_ids"]
         attention_mask = input_ids.ne(1).to(torch_device)
@@ -88,7 +89,7 @@ def test_load_balancing_loss(self):
 
         # First, we make sure that adding padding tokens doesn't change the loss
         # loss(input_ids, attention_mask=None) == loss(input_ids + padding, attention_mask=attention_mask_with_padding)
-        pad_length = 1000
+        pad_length = input_ids.shape[1] * 4
         # Add padding tokens (assume that pad_token_id=1) to input_ids
         padding_block = torch.ones(input_ids.shape[0], pad_length, dtype=torch.int32).to(torch_device)
         padded_input_ids = torch.cat((padding_block, input_ids), dim=1)  # this is to simulate padding to the left

filename: tests/models/qwen2_moe/test_modeling_qwen2_moe.py
patch: @@ -20,6 +20,7 @@
 from transformers import AutoTokenizer, is_torch_available, set_seed
 from transformers.testing_utils import (
     cleanup,
+    is_flaky,
     require_flash_attn,
     require_torch,
     require_torch_gpu,
@@ -72,14 +73,14 @@ def is_pipeline_test_to_skip(
     def test_flash_attn_2_inference_equivalence_right_padding(self):
         self.skipTest(reason="Qwen2Moe flash attention does not support right padding")
 
-    # Ignore copy
+    @is_flaky(max_attempts=2)
     def test_load_balancing_loss(self):
         r"""
         Let's make sure we can actually compute the loss and do a backward on it.
         """
         config, input_dict = self.model_tester.prepare_config_and_inputs_for_common()
         config.num_labels = 3
-        config.num_experts = 8
+        config.num_experts = 3
         config.expert_interval = 2
         config.output_router_logits = True
         input_ids = input_dict["input_ids"]
@@ -93,7 +94,7 @@ def test_load_balancing_loss(self):
 
         # First, we make sure that adding padding tokens doesn't change the loss
         # loss(input_ids, attention_mask=None) == loss(input_ids + padding, attention_mask=attention_mask_with_padding)
-        pad_length = 1000
+        pad_length = input_ids.shape[1] * 4
         # Add padding tokens (assume that pad_token_id=1) to input_ids
         padding_block = torch.ones(input_ids.shape[0], pad_length, dtype=torch.int32).to(torch_device)
         padded_input_ids = torch.cat((padding_block, input_ids), dim=1)  # this is to simulate padding to the left

filename: tests/models/qwen3_moe/test_modeling_qwen3_moe.py
patch: @@ -20,6 +20,7 @@
 from transformers import AutoTokenizer, BitsAndBytesConfig, is_torch_available, set_seed
 from transformers.testing_utils import (
     cleanup,
+    is_flaky,
     require_bitsandbytes,
     require_flash_attn,
     require_torch,
@@ -63,14 +64,14 @@ def is_pipeline_test_to_skip(
     ):
         return True
 
-    # Ignore copy
+    @is_flaky(max_attempts=2)
     def test_load_balancing_loss(self):
         r"""
         Let's make sure we can actually compute the loss and do a backward on it.
         """
         config, input_dict = self.model_tester.prepare_config_and_inputs_for_common()
         config.num_labels = 3
-        config.num_experts = 8
+        config.num_experts = 3
         config.expert_interval = 2
         config.output_router_logits = True
         input_ids = input_dict["input_ids"]
@@ -84,7 +85,7 @@ def test_load_balancing_loss(self):
 
         # First, we make sure that adding padding tokens doesn't change the loss
         # loss(input_ids, attention_mask=None) == loss(input_ids + padding, attention_mask=attention_mask_with_padding)
-        pad_length = 1000
+        pad_length = input_ids.shape[1] * 4
         # Add padding tokens (assume that pad_token_id=1) to input_ids
         padding_block = torch.ones(input_ids.shape[0], pad_length, dtype=torch.int32).to(torch_device)
         padded_input_ids = torch.cat((padding_block, input_ids), dim=1)  # this is to simulate padding to the left