Repository: huggingface/transformers
Pull request: Fix `GptOssModelTest::test_assisted_decoding_matches_greedy_search_1_same`

# What does this PR do?

This test is flaky:

 - the eps needs to be set to 1.0 to avoid large difference
 - but even with this, we still get something > 5e-5 from time to time on CPU. On GPU, they are all very small , around 1e-9.
   - no idea and just increate the atol/rtol to 1e-4 

Changed files:
filename: tests/generation/test_utils.py
patch: @@ -938,6 +938,10 @@ def test_assisted_decoding_matches_greedy_search(self, assistant_type):
         # - assisted_decoding does not support `use_cache = False`
         # - assisted_decoding does not support `batch_size > 1`
 
+        # No idea why this cause problem!
+        if type(self).__name__ not in ["Gemma3nTextModelTest"]:
+            set_model_tester_for_less_flaky_test(self)
+
         for model_class in self.all_generative_model_classes:
             if model_class._is_stateful:
                 self.skipTest(reason="Stateful models don't support assisted generation")
@@ -959,6 +963,7 @@ def test_assisted_decoding_matches_greedy_search(self, assistant_type):
 
             # enable cache
             config, inputs_dict = self.prepare_config_and_inputs_for_generate(batch_size=1)
+            set_config_for_less_flaky_test(config)
 
             # force eager attention to support output attentions
             if self.has_attentions:
@@ -970,6 +975,7 @@ def test_assisted_decoding_matches_greedy_search(self, assistant_type):
 
             config.is_decoder = True
             model = model_class._from_config(config, attn_implementation="eager").to(torch_device).eval()
+            set_model_for_less_flaky_test(model)
             config = model.config
             # Sets assisted generation arguments such that:
             # a) no EOS is generated, to ensure generation doesn't break early
@@ -1006,8 +1012,15 @@ def test_assisted_decoding_matches_greedy_search(self, assistant_type):
             generation_kwargs.update({"assistant_model": assistant_model})
             output_assisted = model.generate(**generation_kwargs, **inputs_dict, **logits_processor_kwargs)
 
+            # default values of `has_similar_generate_outputs`
+            atol, rtol = 1e-5, 1e-5
+            # `gpt_oss` seems to have larger differences on CPU every other generated tokens, sth. like
+            # 1e-9, 1e-5, 1e-9, 1e-5. While on GPU, they are all very small 1e-9.
+            if model.config.model_type == "gpt_oss" and torch_device == "cpu":
+                atol, rtol = 1e-4, 1e-4
+
             # The two outputs must match and their shape must be as expected
-            self.assertTrue(has_similar_generate_outputs(output_greedy, output_assisted))
+            self.assertTrue(has_similar_generate_outputs(output_greedy, output_assisted, atol=atol, rtol=rtol))
             for output in (output_greedy, output_assisted):
                 self._check_generate_outputs(output, model.config, use_cache=True)
 
@@ -5082,7 +5095,7 @@ def has_similar_generate_outputs(output_1, output_2, atol=1e-5, rtol=1e-5) -> bo
             output_1_first_mismatch_scores = output_1.scores[first_mismatch_idx][batch_idx]
             output_2_first_mismatch_scores = output_2.scores[first_mismatch_idx][batch_idx]
             has_matching_scores = torch.allclose(
-                output_1_first_mismatch_scores, output_2_first_mismatch_scores, rtol=atol, atol=rtol
+                output_1_first_mismatch_scores, output_2_first_mismatch_scores, atol=atol, rtol=rtol
             )
             if not has_matching_scores:
                 break