update

Author: DN6

tests/models/testing_utils/__init__.py

@@ -1,9 +1,10 @@
-from .attention import AttentionTesterMixin, ContextParallelTesterMixin
+from .attention import AttentionTesterMixin
 from .common import BaseModelTesterConfig, ModelTesterMixin
 from .compile import TorchCompileTesterMixin
 from .ip_adapter import IPAdapterTesterMixin
 from .lora import LoraHotSwappingForModelTesterMixin, LoraTesterMixin
 from .memory import CPUOffloadTesterMixin, GroupOffloadTesterMixin, LayerwiseCastingTesterMixin, MemoryTesterMixin
+from .parallelism import ContextParallelTesterMixin
 from .quantization import (
     BitsAndBytesTesterMixin,
     GGUFTesterMixin,

tests/models/testing_utils/attention.py

@@ -13,13 +13,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import os
-
 import pytest
 import torch
-import torch.multiprocessing as mp
 
-from diffusers.models._modeling_parallel import ContextParallelConfig
 from diffusers.models.attention import AttentionModuleMixin
 from diffusers.models.attention_processor import (
     AttnProcessor,
@@ -28,8 +24,6 @@
 from ...testing_utils import (
     assert_tensors_close,
     is_attention,
-    is_context_parallel,
-    require_torch_multi_accelerator,
     torch_device,
 )
 
@@ -71,9 +65,7 @@ def test_fuse_unfuse_qkv_projections(self):
 
         # Get output before fusion
         with torch.no_grad():
-            output_before_fusion = model(**inputs_dict)
-            if isinstance(output_before_fusion, dict):
-                output_before_fusion = output_before_fusion.to_tuple()[0]
+            output_before_fusion = model(**inputs_dict, return_dict=False)[0]
 
         # Fuse projections
         model.fuse_qkv_projections()
@@ -90,9 +82,7 @@ def test_fuse_unfuse_qkv_projections(self):
         if has_fused_projections:
             # Get output after fusion
             with torch.no_grad():
-                output_after_fusion = model(**inputs_dict)
-                if isinstance(output_after_fusion, dict):
-                    output_after_fusion = output_after_fusion.to_tuple()[0]
+                output_after_fusion = model(**inputs_dict, return_dict=False)[0]
 
             # Verify outputs match
             assert_tensors_close(
@@ -115,9 +105,7 @@ def test_fuse_unfuse_qkv_projections(self):
 
             # Get output after unfusion
             with torch.no_grad():
-                output_after_unfusion = model(**inputs_dict)
-                if isinstance(output_after_unfusion, dict):
-                    output_after_unfusion = output_after_unfusion.to_tuple()[0]
+                output_after_unfusion = model(**inputs_dict, return_dict=False)[0]
 
             # Verify outputs still match
             assert_tensors_close(
@@ -195,80 +183,3 @@ def test_attention_processor_count_mismatch_raises_error(self):
             model.set_attn_processor(wrong_processors)
 
         assert "number of processors" in str(exc_info.value).lower(), "Error should mention processor count mismatch"
-
-
-def _context_parallel_worker(rank, world_size, model_class, init_dict, cp_dict, inputs_dict, result_queue):
-    try:
-        # Setup distributed environment
-        os.environ["MASTER_ADDR"] = "localhost"
-        os.environ["MASTER_PORT"] = "12355"
-
-        torch.distributed.init_process_group(
-            backend="nccl",
-            init_method="env://",
-            world_size=world_size,
-            rank=rank,
-        )
-        torch.cuda.set_device(rank)
-        device = torch.device(f"cuda:{rank}")
-
-        model = model_class(**init_dict)
-        model.to(device)
-        model.eval()
-
-        inputs_on_device = {}
-        for key, value in inputs_dict.items():
-            if isinstance(value, torch.Tensor):
-                inputs_on_device[key] = value.to(device)
-            else:
-                inputs_on_device[key] = value
-
-        cp_config = ContextParallelConfig(**cp_dict)
-        model.enable_parallelism(config=cp_config)
-
-        with torch.no_grad():
-            output = model(**inputs_on_device)
-            if isinstance(output, dict):
-                output = output.to_tuple()[0]
-
-        if rank == 0:
-            result_queue.put(("success", output.shape))
-
-    except Exception as e:
-        if rank == 0:
-            result_queue.put(("error", str(e)))
-    finally:
-        if torch.distributed.is_initialized():
-            torch.distributed.destroy_process_group()
-
-
-@is_context_parallel
-@require_torch_multi_accelerator
-class ContextParallelTesterMixin:
-    base_precision = 1e-3
-
-    @pytest.mark.parametrize("cp_type", ["ulysses_degree", "ring_degree"], ids=["ulysses", "ring"])
-    def test_context_parallel_inference(self, cp_type):
-        if not torch.distributed.is_available():
-            pytest.skip("torch.distributed is not available.")
-
-        if not hasattr(self.model_class, "_cp_plan") or self.model_class._cp_plan is None:
-            pytest.skip("Model does not have a _cp_plan defined for context parallel inference.")
-
-        world_size = 2
-        init_dict = self.get_init_dict()
-        inputs_dict = self.get_dummy_inputs()
-        cp_dict = {cp_type: world_size}
-
-        ctx = mp.get_context("spawn")
-        result_queue = ctx.Queue()
-
-        mp.spawn(
-            _context_parallel_worker,
-            args=(world_size, self.model_class, init_dict, cp_dict, inputs_dict, result_queue),
-            nprocs=world_size,
-            join=True,
-        )
-
-        status, result = result_queue.get(timeout=60)
-        assert status == "success", f"Context parallel inference failed: {result}"

tests/models/testing_utils/common.py

@@ -259,7 +259,7 @@ class TestMyModel(MyModelTestConfig, ModelTesterMixin):
             pass
     """
 
-    def test_from_save_pretrained(self, tmp_path, atol=5e-5, rtol=0):
+    def test_from_save_pretrained(self, tmp_path, atol=5e-5, rtol=5e-5):
         torch.manual_seed(0)
         model = self.model_class(**self.get_init_dict())
         model.to(torch_device)
@@ -278,15 +278,8 @@ def test_from_save_pretrained(self, tmp_path, atol=5e-5, rtol=0):
             )
 
         with torch.no_grad():
-            image = model(**self.get_dummy_inputs())
-
-            if isinstance(image, dict):
-                image = image.to_tuple()[0]
-
-            new_image = new_model(**self.get_dummy_inputs())
-
-            if isinstance(new_image, dict):
-                new_image = new_image.to_tuple()[0]
+            image = model(**self.get_dummy_inputs(), return_dict=False)[0]
+            new_image = new_model(**self.get_dummy_inputs(), return_dict=False)[0]
 
         assert_tensors_close(image, new_image, atol=atol, rtol=rtol, msg="Models give different forward passes.")
 
@@ -308,14 +301,8 @@ def test_from_save_pretrained_variant(self, tmp_path, atol=5e-5, rtol=0):
         new_model.to(torch_device)
 
         with torch.no_grad():
-            image = model(**self.get_dummy_inputs())
-            if isinstance(image, dict):
-                image = image.to_tuple()[0]
-
-            new_image = new_model(**self.get_dummy_inputs())
-
-            if isinstance(new_image, dict):
-                new_image = new_image.to_tuple()[0]
+            image = model(**self.get_dummy_inputs(), return_dict=False)[0]
+            new_image = new_model(**self.get_dummy_inputs(), return_dict=False)[0]
 
         assert_tensors_close(image, new_image, atol=atol, rtol=rtol, msg="Models give different forward passes.")
 
@@ -343,13 +330,8 @@ def test_determinism(self, atol=1e-5, rtol=0):
         model.eval()
 
         with torch.no_grad():
-            first = model(**self.get_dummy_inputs())
-            if isinstance(first, dict):
-                first = first.to_tuple()[0]
-
-            second = model(**self.get_dummy_inputs())
-            if isinstance(second, dict):
-                second = second.to_tuple()[0]
+            first = model(**self.get_dummy_inputs(), return_dict=False)[0]
+            second = model(**self.get_dummy_inputs(), return_dict=False)[0]
 
         # Filter out NaN values before comparison
         first_flat = first.flatten()
@@ -369,10 +351,7 @@ def test_output(self, expected_output_shape=None):
 
         inputs_dict = self.get_dummy_inputs()
         with torch.no_grad():
-            output = model(**inputs_dict)
-
-            if isinstance(output, dict):
-                output = output.to_tuple()[0]
+            output = model(**inputs_dict, return_dict=False)[0]
 
         assert output is not None, "Model output is None"
         assert output[0].shape == expected_output_shape or self.output_shape, (
@@ -501,13 +480,8 @@ def test_from_save_pretrained_dtype_inference(self, tmp_path, dtype):
                 assert param.data.dtype == dtype
 
         with torch.no_grad():
-            output = model(**self.get_dummy_inputs())
-            if isinstance(output, dict):
-                output = output.to_tuple()[0]
-
-            output_loaded = model_loaded(**self.get_dummy_inputs())
-            if isinstance(output_loaded, dict):
-                output_loaded = output_loaded.to_tuple()[0]
+            output = model(**self.get_dummy_inputs(), return_dict=False)[0]
+            output_loaded = model_loaded(**self.get_dummy_inputs(), return_dict=False)[0]
 
         assert_tensors_close(output, output_loaded, atol=1e-4, rtol=0, msg=f"Loaded model output differs for {dtype}")
 
@@ -519,7 +493,7 @@ def test_sharded_checkpoints(self, tmp_path):
         model = self.model_class(**config).eval()
         model = model.to(torch_device)
 
-        base_output = model(**inputs_dict)
+        base_output = model(**inputs_dict, return_dict=False)[0]
 
         model_size = compute_module_persistent_sizes(model)[""]
         max_shard_size = int((model_size * 0.75) / (2**10))  # Convert to KB as these test models are small
@@ -539,10 +513,10 @@ def test_sharded_checkpoints(self, tmp_path):
 
         torch.manual_seed(0)
         inputs_dict_new = self.get_dummy_inputs()
-        new_output = new_model(**inputs_dict_new)
+        new_output = new_model(**inputs_dict_new, return_dict=False)[0]
 
         assert_tensors_close(
-            base_output[0], new_output[0], atol=1e-5, rtol=0, msg="Output should match after sharded save/load"
+            base_output, new_output, atol=1e-5, rtol=0, msg="Output should match after sharded save/load"
         )
 
     @require_accelerator
@@ -553,7 +527,7 @@ def test_sharded_checkpoints_with_variant(self, tmp_path):
         model = self.model_class(**config).eval()
         model = model.to(torch_device)
 
-        base_output = model(**inputs_dict)
+        base_output = model(**inputs_dict, return_dict=False)[0]
 
         model_size = compute_module_persistent_sizes(model)[""]
         max_shard_size = int((model_size * 0.75) / (2**10))  # Convert to KB as these test models are small
@@ -578,10 +552,10 @@ def test_sharded_checkpoints_with_variant(self, tmp_path):
 
         torch.manual_seed(0)
         inputs_dict_new = self.get_dummy_inputs()
-        new_output = new_model(**inputs_dict_new)
+        new_output = new_model(**inputs_dict_new, return_dict=False)[0]
 
         assert_tensors_close(
-            base_output[0], new_output[0], atol=1e-5, rtol=0, msg="Output should match after variant sharded save/load"
+            base_output, new_output, atol=1e-5, rtol=0, msg="Output should match after variant sharded save/load"
         )
 
     def test_sharded_checkpoints_with_parallel_loading(self, tmp_path):
@@ -593,7 +567,7 @@ def test_sharded_checkpoints_with_parallel_loading(self, tmp_path):
         model = self.model_class(**config).eval()
         model = model.to(torch_device)
 
-        base_output = model(**inputs_dict)
+        base_output = model(**inputs_dict, return_dict=False)[0]
 
         model_size = compute_module_persistent_sizes(model)[""]
         max_shard_size = int((model_size * 0.75) / (2**10))  # Convert to KB as these test models are small
@@ -628,10 +602,10 @@ def test_sharded_checkpoints_with_parallel_loading(self, tmp_path):
 
             torch.manual_seed(0)
             inputs_dict_parallel = self.get_dummy_inputs()
-            output_parallel = model_parallel(**inputs_dict_parallel)
+            output_parallel = model_parallel(**inputs_dict_parallel, return_dict=False)[0]
 
             assert_tensors_close(
-                base_output[0], output_parallel[0], atol=1e-5, rtol=0, msg="Output should match with parallel loading"
+                base_output, output_parallel, atol=1e-5, rtol=0, msg="Output should match with parallel loading"
             )
 
         finally:
@@ -652,7 +626,7 @@ def test_model_parallelism(self, tmp_path):
         model = model.to(torch_device)
 
         torch.manual_seed(0)
-        base_output = model(**inputs_dict)
+        base_output = model(**inputs_dict, return_dict=False)[0]
 
         model_size = compute_module_sizes(model)[""]
         max_gpu_sizes = [int(p * model_size) for p in self.model_split_percents]
@@ -668,8 +642,8 @@ def test_model_parallelism(self, tmp_path):
             check_device_map_is_respected(new_model, new_model.hf_device_map)
 
             torch.manual_seed(0)
-            new_output = new_model(**inputs_dict)
+            new_output = new_model(**inputs_dict, return_dict=False)[0]
 
             assert_tensors_close(
-                base_output[0], new_output[0], atol=1e-5, rtol=0, msg="Output should match with model parallelism"
+                base_output, new_output, atol=1e-5, rtol=0, msg="Output should match with model parallelism"
             )