Use device_map="auto" in single file tests to support large models on limited GPU memory (#13816)

* fix flux tests OOM on 24G GPU

Signed-off-by: jiqing-feng <jiqing.feng@intel.com>

* revert wrong change

Signed-off-by: jiqing-feng <jiqing.feng@intel.com>

---------

Signed-off-by: jiqing-feng <jiqing.feng@intel.com>
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>

Author: jiqing-feng

tests/models/testing_utils/single_file.py

@@ -107,8 +107,8 @@ def teardown_method(self):
         backend_empty_cache(torch_device)
 
     def test_single_file_model_config(self):
-        pretrained_kwargs = {"device": torch_device, **self.pretrained_model_kwargs}
-        single_file_kwargs = {"device": torch_device}
+        pretrained_kwargs = {"device_map": "auto", **self.pretrained_model_kwargs}
+        single_file_kwargs = {"device_map": "auto"}
 
         if self.torch_dtype:
             pretrained_kwargs["torch_dtype"] = self.torch_dtype
@@ -127,8 +127,8 @@ def test_single_file_model_config(self):
             )
 
     def test_single_file_model_parameters(self):
-        pretrained_kwargs = {"device_map": str(torch_device), **self.pretrained_model_kwargs}
-        single_file_kwargs = {"device": torch_device}
+        pretrained_kwargs = {"device_map": "auto", **self.pretrained_model_kwargs}
+        single_file_kwargs = {"device_map": "auto"}
 
         if self.torch_dtype:
             pretrained_kwargs["torch_dtype"] = self.torch_dtype
@@ -259,7 +259,7 @@ def test_checkpoint_variant_loading(self):
             backend_empty_cache(torch_device)
 
     def test_single_file_loading_with_device_map(self):
-        single_file_kwargs = {"device_map": torch_device}
+        single_file_kwargs = {"device_map": "auto"}
 
         if self.torch_dtype:
             single_file_kwargs["torch_dtype"] = self.torch_dtype

tests/models/transformers/test_models_transformer_flux.py

@@ -346,6 +346,10 @@ def alternate_ckpt_paths(self):
     def pretrained_model_name_or_path(self):
         return "black-forest-labs/FLUX.1-dev"
 
+    @property
+    def torch_dtype(self):
+        return torch.bfloat16
+
 
 class TestFluxTransformerBitsAndBytes(FluxTransformerTesterConfig, BitsAndBytesTesterMixin):
     """BitsAndBytes quantization tests for Flux Transformer."""

tests/single_file/test_model_flux_transformer_single_file.py

@@ -15,6 +15,8 @@
 
 import gc
 
+import torch
+
 from diffusers import (
     FluxTransformer2DModel,
 )
@@ -38,9 +40,9 @@ class TestFluxTransformer2DModelSingleFile(SingleFileModelTesterMixin):
     repo_id = "black-forest-labs/FLUX.1-dev"
     subfolder = "transformer"
 
-    def test_device_map_cuda(self):
+    def test_device_map_auto(self):
         backend_empty_cache(torch_device)
-        model = self.model_class.from_single_file(self.ckpt_path, device_map="cuda")
+        model = self.model_class.from_single_file(self.ckpt_path, device_map="auto", torch_dtype=torch.bfloat16)
 
         del model
         gc.collect()