Clean up flaky behaviour on Slow CUDA Pytorch Push Tests (#4759)

use max diff to compare model outputs

tests/models/test_modeling_common.py

@@ -195,7 +195,7 @@ class ModelTesterMixin:
     main_input_name = None  # overwrite in model specific tester class
     base_precision = 1e-3
 
-    def test_from_save_pretrained(self):
+    def test_from_save_pretrained(self, expected_max_diff=5e-5):
         init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
 
         model = self.model_class(**init_dict)
@@ -221,8 +221,8 @@ class ModelTesterMixin:
             if isinstance(new_image, dict):
                 new_image = new_image.to_tuple()[0]
 
-        max_diff = (image - new_image).abs().sum().item()
-        self.assertLessEqual(max_diff, 5e-5, "Models give different forward passes")
+        max_diff = (image - new_image).abs().max().item()
+        self.assertLessEqual(max_diff, expected_max_diff, "Models give different forward passes")
 
     def test_getattr_is_correct(self):
         init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
@@ -316,7 +316,7 @@ class ModelTesterMixin:
         assert torch.allclose(output_2, output_5, atol=self.base_precision)
         assert torch.allclose(output_2, output_6, atol=self.base_precision)
 
-    def test_from_save_pretrained_variant(self):
+    def test_from_save_pretrained_variant(self, expected_max_diff=5e-5):
         init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
 
         model = self.model_class(**init_dict)
@@ -351,8 +351,8 @@ class ModelTesterMixin:
             if isinstance(new_image, dict):
                 new_image = new_image.to_tuple()[0]
 
-        max_diff = (image - new_image).abs().sum().item()
-        self.assertLessEqual(max_diff, 5e-5, "Models give different forward passes")
+        max_diff = (image - new_image).abs().max().item()
+        self.assertLessEqual(max_diff, expected_max_diff, "Models give different forward passes")
 
     @require_torch_2
     def test_from_save_pretrained_dynamo(self):

tests/models/test_models_unet_2d_condition.py

@@ -589,10 +589,10 @@ class UNet2DConditionModelTests(ModelTesterMixin, UNetTesterMixin, unittest.Test
         with torch.no_grad():
             new_sample = new_model(**inputs_dict, cross_attention_kwargs={"scale": 0.5}).sample
 
-        assert (sample - new_sample).abs().max() < 1e-4
+        assert (sample - new_sample).abs().max() < 5e-4
 
         # LoRA and no LoRA should NOT be the same
-        assert (sample - old_sample).abs().max() > 1e-4
+        assert (sample - old_sample).abs().max() > 5e-4
 
     def test_lora_save_load_safetensors(self):
         # enable deterministic behavior for gradient checkpointing

tests/pipelines/controlnet/test_controlnet.py

@@ -959,7 +959,7 @@ class ControlNetPipelineSlowTests(unittest.TestCase):
             gc.collect()
             torch.cuda.empty_cache()
 
-        assert np.abs(images[0] - images[1]).sum() < 1e-3
+        assert np.abs(images[0] - images[1]).max() < 1e-3
 
 
 @slow

tests/pipelines/controlnet/test_controlnet_img2img.py

@@ -446,4 +446,4 @@ class ControlNetImg2ImgPipelineSlowTests(unittest.TestCase):
             gc.collect()
             torch.cuda.empty_cache()
 
-        assert np.abs(images[0] - images[1]).sum() < 1e-3
+        assert np.abs(images[0] - images[1]).max() < 1e-3

tests/pipelines/controlnet/test_controlnet_inpaint.py

@@ -593,4 +593,4 @@ class ControlNetInpaintPipelineSlowTests(unittest.TestCase):
             gc.collect()
             torch.cuda.empty_cache()
 
-        assert np.abs(images[0] - images[1]).sum() < 1e-3
+        assert np.abs(images[0] - images[1]).max() < 1e-3

tests/pipelines/test_pipelines.py

@@ -122,7 +122,7 @@ def _test_from_save_pretrained_dynamo(in_queue, out_queue, timeout):
         generator = torch.Generator(device=torch_device).manual_seed(0)
         new_image = new_ddpm(generator=generator, num_inference_steps=5, output_type="numpy").images
 
-        assert np.abs(image - new_image).sum() < 1e-5, "Models don't give the same forward pass"
+        assert np.abs(image - new_image).max() < 1e-5, "Models don't give the same forward pass"
     except Exception:
         error = f"{traceback.format_exc()}"
 
@@ -1543,7 +1543,7 @@ class PipelineSlowTests(unittest.TestCase):
         generator = torch.Generator(device=torch_device).manual_seed(0)
         new_image = new_ddpm(generator=generator, num_inference_steps=5, output_type="numpy").images
 
-        assert np.abs(image - new_image).sum() < 1e-5, "Models don't give the same forward pass"
+        assert np.abs(image - new_image).max() < 1e-5, "Models don't give the same forward pass"
 
     @require_torch_2
     def test_from_save_pretrained_dynamo(self):
@@ -1568,7 +1568,7 @@ class PipelineSlowTests(unittest.TestCase):
         generator = torch.Generator(device=torch_device).manual_seed(0)
         new_image = ddpm_from_hub(generator=generator, num_inference_steps=5, output_type="numpy").images
 
-        assert np.abs(image - new_image).sum() < 1e-5, "Models don't give the same forward pass"
+        assert np.abs(image - new_image).max() < 1e-5, "Models don't give the same forward pass"
 
     def test_from_pretrained_hub_pass_model(self):
         model_path = "google/ddpm-cifar10-32"
@@ -1591,7 +1591,7 @@ class PipelineSlowTests(unittest.TestCase):
         generator = torch.Generator(device=torch_device).manual_seed(0)
         new_image = ddpm_from_hub(generator=generator, num_inference_steps=5, output_type="numpy").images
 
-        assert np.abs(image - new_image).sum() < 1e-5, "Models don't give the same forward pass"
+        assert np.abs(image - new_image).max() < 1e-5, "Models don't give the same forward pass"
 
     def test_output_format(self):
         model_path = "google/ddpm-cifar10-32"

tests/pipelines/test_pipelines_common.py

@@ -296,7 +296,7 @@ class PipelineTesterMixin:
         gc.collect()
         torch.cuda.empty_cache()
 
-    def test_save_load_local(self, expected_max_difference=1e-4):
+    def test_save_load_local(self, expected_max_difference=5e-4):
         components = self.get_dummy_components()
         pipe = self.pipeline_class(**components)
         pipe.to(torch_device)

tests/pipelines/versatile_diffusion/test_versatile_diffusion_dual_guided.py

@@ -76,7 +76,7 @@ class VersatileDiffusionDualGuidedPipelineIntegrationTests(unittest.TestCase):
             output_type="numpy",
         ).images
 
-        assert np.abs(image - new_image).sum() < 1e-5, "Models don't have the same forward pass"
+        assert np.abs(image - new_image).max() < 1e-5, "Models don't have the same forward pass"
 
     def test_inference_dual_guided(self):
         pipe = VersatileDiffusionDualGuidedPipeline.from_pretrained("shi-labs/versatile-diffusion")

tests/pipelines/versatile_diffusion/test_versatile_diffusion_mega.py

@@ -77,7 +77,7 @@ class VersatileDiffusionMegaPipelineIntegrationTests(unittest.TestCase):
             output_type="numpy",
         ).images
 
-        assert np.abs(image - new_image).sum() < 1e-5, "Models don't have the same forward pass"
+        assert np.abs(image - new_image).max() < 1e-5, "Models don't have the same forward pass"
 
     def test_inference_dual_guided_then_text_to_image(self):
         pipe = VersatileDiffusionPipeline.from_pretrained("shi-labs/versatile-diffusion", torch_dtype=torch.float16)

tests/pipelines/versatile_diffusion/test_versatile_diffusion_text_to_image.py

@@ -64,7 +64,7 @@ class VersatileDiffusionTextToImagePipelineIntegrationTests(unittest.TestCase):
             prompt=prompt, generator=generator, guidance_scale=7.5, num_inference_steps=2, output_type="numpy"
         ).images
 
-        assert np.abs(image - new_image).sum() < 1e-5, "Models don't have the same forward pass"
+        assert np.abs(image - new_image).max() < 1e-5, "Models don't have the same forward pass"
 
     def test_inference_text2img(self):
         pipe = VersatileDiffusionTextToImagePipeline.from_pretrained(