Repository: huggingface/transformers
Pull request: Fix flaky `JambaModelTest.test_load_balancing_loss`

# What does this PR do?

Fix this flaky test which has

```bash
>       torch.testing.assert_close(result.aux_loss.cpu(), torch.tensor(2, dtype=torch.float32), rtol=1e-2, atol=1e-2)
E       AssertionError: Scalars are not close!
```

Running 10000 times and all pass.

Also makes this much much faster by reducing model size and padded sequence length: running the same tests 20 times

    20 passed, 2 warnings in 11.35s v.s. 20 passed, 2 warnings in 75.72s (0:01:15) 



### Full log

```
___________________ JambaModelTest.test_load_balancing_loss ____________________
[gw3] linux -- Python 3.9.23 /usr/local/bin/python3

self = <tests.models.jamba.test_modeling_jamba.JambaModelTest testMethod=test_load_balancing_loss>

    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 = 16
        config.output_router_logits = True
        input_ids = input_dict["input_ids"]
        attention_mask = input_ids.ne(config.pad_token_id).to(torch_device)
        model = JambaForCausalLM(config)
        model.to(torch_device)
        model.eval()
        result = model(input_ids, attention_mask=attention_mask)
        bs, seqlen = input_ids.shape
        self.assertEqual(result.router_logits[0].shape, (bs * seqlen, config.num_experts))
>       torch.testing.assert_close(result.aux_loss.cpu(), torch.tensor(2, dtype=torch.float32), rtol=1e-2, atol=1e-2)
E       AssertionError: Scalars are not close!
E       
E       Expected 2.0 but got 2.0415196418762207.
E       Absolute difference: 0.0415196418762207 (up to 0.01 allowed)
E       Relative difference: 0.02075982093811035 (up to 0.01 allowed)
```


https://app.circleci.com/pipelines/github/huggingface/transformers/144153/workflows/ce26cb2f-f3b1-45c5-bca9-60868311cd7f/jobs/1905517

Changed files:
filename: tests/models/jamba/test_modeling_jamba.py
patch: @@ -24,6 +24,7 @@
     DeviceProperties,
     Expectations,
     get_device_properties,
+    is_flaky,
     require_bitsandbytes,
     require_flash_attn,
     require_torch,
@@ -104,10 +105,10 @@ def __init__(
         use_labels=True,
         vocab_size=99,
         hidden_size=32,
-        num_hidden_layers=5,
+        num_hidden_layers=2,
         attn_layer_offset=1,
         attn_layer_period=8,
-        num_attention_heads=4,
+        num_attention_heads=2,
         num_key_value_heads=2,
         intermediate_size=37,
         hidden_act="gelu",
@@ -385,13 +386,15 @@ def test_decoder_model_past_with_large_inputs(self):
         config_and_inputs = self.model_tester.prepare_config_and_inputs_for_decoder()
         self.model_tester.create_and_check_decoder_model_past_large_inputs(*config_and_inputs)
 
+    # After #40617, we still have 0.01 % of failure rate here.
+    @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 = 16
+        config.num_experts = 3
         config.output_router_logits = True
         input_ids = input_dict["input_ids"]
         attention_mask = input_ids.ne(config.pad_token_id).to(torch_device)
@@ -401,11 +404,13 @@ def test_load_balancing_loss(self):
         result = model(input_ids, attention_mask=attention_mask)
         bs, seqlen = input_ids.shape
         self.assertEqual(result.router_logits[0].shape, (bs * seqlen, config.num_experts))
+        # After #40617, we still have 0.01 % of failure rate here.
         torch.testing.assert_close(result.aux_loss.cpu(), torch.tensor(2, dtype=torch.float32), rtol=1e-2, atol=1e-2)
 
         # 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
+        # (This length is selected from experiments)
+        pad_length = input_ids.shape[1] * 4
         # Add padding tokens to input_ids
         padding_block = config.pad_token_id * torch.ones(input_ids.shape[0], pad_length, dtype=torch.int32).to(
             torch_device
@@ -421,6 +426,7 @@ def test_load_balancing_loss(self):
         include_padding_result = model(padded_input_ids, attention_mask=None)
 
         # This is to mimic torch.testing.assert_not_close
+        # After #40617, we still have 0.003 % of failure rate here.
         self.assertNotAlmostEqual(include_padding_result.aux_loss.item(), result.aux_loss.item())
 
     def test_initialization(self):