add tests

Author: SunMarc

src/diffusers/models/model_loading_utils.py

@@ -239,12 +239,8 @@ def load_model_dict_into_meta(
         # in int/uint/bool and not cast them.
         # TODO: revisit cases when param.dtype == torch.float8_e4m3fn
         if dtype is not None and torch.is_floating_point(param):
-            if (
-                keep_in_fp32_modules is not None
-                and any(
-                    module_to_keep_in_fp32 in param_name.split(".") for module_to_keep_in_fp32 in keep_in_fp32_modules
-                )
-                and dtype == torch.float16
+            if keep_in_fp32_modules is not None and any(
+                module_to_keep_in_fp32 in param_name.split(".") for module_to_keep_in_fp32 in keep_in_fp32_modules
             ):
                 param = param.to(torch.float32)
                 set_module_kwargs["dtype"] = torch.float32

src/diffusers/models/modeling_utils.py

@@ -1011,9 +1011,10 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
                 raise ValueError("`low_cpu_mem_usage` cannot be False or None when using quantization.")
 
         # Check if `_keep_in_fp32_modules` is not None
-        use_keep_in_fp32_modules = (cls._keep_in_fp32_modules is not None) and (
-            (torch_dtype == torch.float16) or hasattr(hf_quantizer, "use_keep_in_fp32_modules")
+        use_keep_in_fp32_modules = (cls._keep_in_fp32_modules is not None) or hasattr(
+            hf_quantizer, "use_keep_in_fp32_modules"
         )
+
         if use_keep_in_fp32_modules:
             keep_in_fp32_modules = cls._keep_in_fp32_modules
             if not isinstance(keep_in_fp32_modules, list):

tests/models/test_modeling_common.py

@@ -37,7 +37,7 @@
 from parameterized import parameterized
 from requests.exceptions import HTTPError
 
-from diffusers.models import UNet2DConditionModel
+from diffusers.models import SD3Transformer2DModel, UNet2DConditionModel
 from diffusers.models.attention_processor import (
     AttnProcessor,
     AttnProcessor2_0,
@@ -334,6 +334,28 @@ def test_weight_overwrite(self):
 
         assert model.config.in_channels == 9
 
+    def test_keep_modules_in_fp32(self):
+        r"""
+        A simple tests to check if the modules under `_keep_in_fp32_modules` are kept in fp32 when we load the model in fp16/bf16
+        Also ensures if inference works.
+        """
+        fp32_modules = SD3Transformer2DModel._keep_in_fp32_modules
+
+        for torch_dtype in [torch.bfloat16, torch.float16]:
+            SD3Transformer2DModel._keep_in_fp32_modules = ["proj_out"]
+
+            model = SD3Transformer2DModel.from_pretrained(
+                "stabilityai/stable-diffusion-3-medium-diffusers", subfolder="transformer", torch_dtype=torch_dtype
+            )
+
+            for name, module in model.named_modules():
+                if isinstance(module, torch.nn.Linear):
+                    if name in model._keep_in_fp32_modules:
+                        self.assertTrue(module.weight.dtype == torch.float32)
+                    else:
+                        self.assertTrue(module.weight.dtype == torch_dtype)
+        SD3Transformer2DModel._keep_in_fp32_modules = fp32_modules
+
 
 class UNetTesterMixin:
     def test_forward_with_norm_groups(self):

tests/quantization/bnb/test_4bit.py

@@ -136,7 +136,7 @@ def setUp(self):
             bnb_4bit_compute_dtype=torch.float16,
         )
         self.model_4bit = SD3Transformer2DModel.from_pretrained(
-            self.model_name, subfolder="transformer", quantization_config=nf4_config
+            self.model_name, subfolder="transformer", quantization_config=nf4_config, device_map=torch_device
         )
 
     def tearDown(self):
@@ -202,7 +202,7 @@ def test_keep_modules_in_fp32(self):
             bnb_4bit_compute_dtype=torch.float16,
         )
         model = SD3Transformer2DModel.from_pretrained(
-            self.model_name, subfolder="transformer", quantization_config=nf4_config
+            self.model_name, subfolder="transformer", quantization_config=nf4_config, device_map=torch_device
         )
 
         for name, module in model.named_modules():
@@ -327,7 +327,7 @@ def test_bnb_4bit_errors_loading_incorrect_state_dict(self):
         with tempfile.TemporaryDirectory() as tmpdirname:
             nf4_config = BitsAndBytesConfig(load_in_4bit=True)
             model_4bit = SD3Transformer2DModel.from_pretrained(
-                self.model_name, subfolder="transformer", quantization_config=nf4_config
+                self.model_name, subfolder="transformer", quantization_config=nf4_config, device_map=torch_device
             )
             model_4bit.save_pretrained(tmpdirname)
             del model_4bit
@@ -362,7 +362,7 @@ def setUp(self):
             bnb_4bit_compute_dtype=torch.float16,
         )
         self.model_4bit = SD3Transformer2DModel.from_pretrained(
-            self.model_name, subfolder="transformer", quantization_config=nf4_config
+            self.model_name, subfolder="transformer", quantization_config=nf4_config, device_map=torch_device
         )
 
     def test_training(self):
@@ -410,7 +410,7 @@ def setUp(self) -> None:
             bnb_4bit_compute_dtype=torch.float16,
         )
         model_4bit = SD3Transformer2DModel.from_pretrained(
-            self.model_name, subfolder="transformer", quantization_config=nf4_config
+            self.model_name, subfolder="transformer", quantization_config=nf4_config, device_map=torch_device
         )
         self.pipeline_4bit = DiffusionPipeline.from_pretrained(
             self.model_name, transformer=model_4bit, torch_dtype=torch.float16
@@ -472,7 +472,7 @@ def test_moving_to_cpu_throws_warning(self):
             bnb_4bit_compute_dtype=torch.float16,
         )
         model_4bit = SD3Transformer2DModel.from_pretrained(
-            self.model_name, subfolder="transformer", quantization_config=nf4_config
+            self.model_name, subfolder="transformer", quantization_config=nf4_config, device_map=torch_device
         )
 
         logger = logging.get_logger("diffusers.pipelines.pipeline_utils")
@@ -502,6 +502,7 @@ def test_pipeline_cuda_placement_works_with_nf4(self):
             subfolder="transformer",
             quantization_config=transformer_nf4_config,
             torch_dtype=torch.float16,
+            device_map=torch_device,
         )
         text_encoder_3_nf4_config = BnbConfig(
             load_in_4bit=True,
@@ -513,6 +514,7 @@ def test_pipeline_cuda_placement_works_with_nf4(self):
             subfolder="text_encoder_3",
             quantization_config=text_encoder_3_nf4_config,
             torch_dtype=torch.float16,
+            device_map=torch_device,
         )
         # CUDA device placement works.
         pipeline_4bit = DiffusionPipeline.from_pretrained(
@@ -527,6 +529,94 @@ def test_pipeline_cuda_placement_works_with_nf4(self):
 
         del pipeline_4bit
 
+    def test_device_map(self):
+        """
+        Test if the quantized model is working properly with "auto".
+        cpu/disk offloading as well doesn't work with bnb.
+        """
+
+        def get_dummy_tensor_inputs(device=None, seed: int = 0):
+            batch_size = 1
+            num_latent_channels = 4
+            num_image_channels = 3
+            height = width = 4
+            sequence_length = 48
+            embedding_dim = 32
+
+            torch.manual_seed(seed)
+            hidden_states = torch.randn((batch_size, height * width, num_latent_channels)).to(
+                device, dtype=torch.bfloat16
+            )
+            torch.manual_seed(seed)
+            encoder_hidden_states = torch.randn((batch_size, sequence_length, embedding_dim)).to(
+                device, dtype=torch.bfloat16
+            )
+
+            torch.manual_seed(seed)
+            pooled_prompt_embeds = torch.randn((batch_size, embedding_dim)).to(device, dtype=torch.bfloat16)
+
+            torch.manual_seed(seed)
+            text_ids = torch.randn((sequence_length, num_image_channels)).to(device, dtype=torch.bfloat16)
+
+            torch.manual_seed(seed)
+            image_ids = torch.randn((height * width, num_image_channels)).to(device, dtype=torch.bfloat16)
+
+            timestep = torch.tensor([1.0]).to(device, dtype=torch.bfloat16).expand(batch_size)
+
+            return {
+                "hidden_states": hidden_states,
+                "encoder_hidden_states": encoder_hidden_states,
+                "pooled_projections": pooled_prompt_embeds,
+                "txt_ids": text_ids,
+                "img_ids": image_ids,
+                "timestep": timestep,
+            }
+
+        inputs = get_dummy_tensor_inputs(torch_device)
+        expected_slice = np.array(
+            [0.47070312, 0.00390625, -0.03662109, -0.19628906, -0.53125, 0.5234375, -0.17089844, -0.59375, 0.578125]
+        )
+
+        # non sharded
+        quantization_config = BitsAndBytesConfig(
+            load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.float16
+        )
+        quantized_model = FluxTransformer2DModel.from_pretrained(
+            "hf-internal-testing/tiny-flux-pipe",
+            subfolder="transformer",
+            quantization_config=quantization_config,
+            device_map="auto",
+            torch_dtype=torch.bfloat16,
+        )
+
+        weight = quantized_model.transformer_blocks[0].ff.net[2].weight
+        self.assertTrue(isinstance(weight, bnb.nn.modules.Params4bit))
+
+        output = quantized_model(**inputs)[0]
+        output_slice = output.flatten()[-9:].detach().float().cpu().numpy()
+        self.assertTrue(np.allclose(output_slice, expected_slice, atol=1e-3, rtol=1e-3))
+
+        # sharded
+
+        quantization_config = BitsAndBytesConfig(
+            load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.float16
+        )
+        quantized_model = FluxTransformer2DModel.from_pretrained(
+            "hf-internal-testing/tiny-flux-sharded",
+            subfolder="transformer",
+            quantization_config=quantization_config,
+            device_map="auto",
+            torch_dtype=torch.bfloat16,
+        )
+
+        weight = quantized_model.transformer_blocks[0].ff.net[2].weight
+        self.assertTrue(isinstance(weight, bnb.nn.modules.Params4bit))
+
+        output = quantized_model(**inputs)[0]
+        output_slice = output.flatten()[-9:].detach().float().cpu().numpy()
+
+        self.assertTrue(np.allclose(output_slice, expected_slice, atol=1e-3, rtol=1e-3))
+
 
 @require_transformers_version_greater("4.44.0")
 class SlowBnb4BitFluxTests(Base4bitTests):
@@ -610,7 +700,10 @@ def test_serialization(self, quant_type="nf4", double_quant=True, safe_serializa
             bnb_4bit_compute_dtype=torch.bfloat16,
         )
         model_0 = SD3Transformer2DModel.from_pretrained(
-            self.model_name, subfolder="transformer", quantization_config=self.quantization_config
+            self.model_name,
+            subfolder="transformer",
+            quantization_config=self.quantization_config,
+            device_map=torch_device,
         )
         self.assertTrue("_pre_quantization_dtype" in model_0.config)
         with tempfile.TemporaryDirectory() as tmpdirname: