fixes

Author: sayakpaul

src/diffusers/quantizers/bitsandbytes/bnb_quantizer.py

@@ -135,6 +135,7 @@ def create_quantized_param(
         target_device: "torch.device",
         state_dict: Dict[str, Any],
         unexpected_keys: Optional[List[str]] = None,
+        **kwargs,
     ):
         import bitsandbytes as bnb
 
@@ -445,6 +446,7 @@ def create_quantized_param(
         target_device: "torch.device",
         state_dict: Dict[str, Any],
         unexpected_keys: Optional[List[str]] = None,
+        **kwargs,
     ):
         import bitsandbytes as bnb
 

src/diffusers/quantizers/torchao/torchao_quantizer.py

@@ -215,6 +215,7 @@ def create_quantized_param(
         target_device: "torch.device",
         state_dict: Dict[str, Any],
         unexpected_keys: List[str],
+        **kwargs,
     ):
         r"""
         Each nn.Linear layer that needs to be quantized is processsed here. First, we set the value the weight tensor,

tests/quantization/bnb/test_4bit.py

@@ -75,6 +75,8 @@ class Base4bitTests(unittest.TestCase):
     # This was obtained on audace so the number might slightly change
     expected_rel_difference = 3.69
 
+    expected_memory_saving_ratio = 0.8
+
     prompt = "a beautiful sunset amidst the mountains."
     num_inference_steps = 10
     seed = 0
@@ -119,8 +121,10 @@ def setUp(self):
         )
 
     def tearDown(self):
-        del self.model_fp16
-        del self.model_4bit
+        if hasattr(self, "model_fp16"):
+            del self.model_fp16
+        if hasattr(self, "model_4bit"):
+            del self.model_4bit
 
         gc.collect()
         torch.cuda.empty_cache()
@@ -159,6 +163,32 @@ def test_memory_footprint(self):
         linear = get_some_linear_layer(self.model_4bit)
         self.assertTrue(linear.weight.__class__ == bnb.nn.Params4bit)
 
+    def test_model_memory_usage(self):
+        # Delete to not let anything interfere.
+        del self.model_4bit, self.model_fp16
+
+        # Re-instantiate.
+        inputs = self.get_dummy_inputs()
+        inputs = {
+            k: v.to(device=torch_device, dtype=torch.float16) for k, v in inputs.items() if not isinstance(v, bool)
+        }
+        model_fp16 = SD3Transformer2DModel.from_pretrained(
+            self.model_name, subfolder="transformer", torch_dtype=torch.float16
+        ).to(torch_device)
+        unquantized_model_memory = get_memory_consumption_stat(model_fp16, inputs)
+        del model_fp16
+
+        nf4_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_compute_dtype=torch.float16,
+        )
+        model_4bit = SD3Transformer2DModel.from_pretrained(
+            self.model_name, subfolder="transformer", quantization_config=nf4_config, torch_dtype=torch.float16
+        )
+        quantized_model_memory = get_memory_consumption_stat(model_4bit, inputs)
+        assert unquantized_model_memory / quantized_model_memory >= self.expected_memory_saving_ratio
+
     def test_original_dtype(self):
         r"""
         A simple test to check if the model succesfully stores the original dtype
@@ -329,29 +359,6 @@ def test_bnb_4bit_errors_loading_incorrect_state_dict(self):
 
             assert key_to_target in str(err_context.exception)
 
-    def test_model_memory_usage(self):
-        # Delete to not let anything interfere.
-        del self.model_4bit, self.model_fp16
-        
-        # Re-instantiate.
-        inputs = self.get_dummy_inputs()
-        model_fp16 = SD3Transformer2DModel.from_pretrained(
-            self.model_name, subfolder="transformer", torch_dtype=torch.float16
-        )
-        unquantized_model_memory = get_memory_consumption_stat(model_fp16, inputs)
-        nf4_config = BitsAndBytesConfig(
-            load_in_4bit=True,
-            bnb_4bit_quant_type="nf4",
-            bnb_4bit_compute_dtype=torch.float16,
-        )
-        model_4bit = SD3Transformer2DModel.from_pretrained(
-            self.model_name, subfolder="transformer", quantization_config=nf4_config, device_map=torch_device
-        )
-        quantized_model_memory = get_memory_consumption_stat(model_4bit, inputs)
-        print(f"{unquantized_model_memory=}, {quantized_model_memory=}")
-        assert (1.0 - (unquantized_model_memory / quantized_model_memory)) >= 100.
-
-
 
 class BnB4BitTrainingTests(Base4bitTests):
     def setUp(self):

tests/quantization/bnb/test_mixed_int8.py

@@ -81,6 +81,8 @@ class Base8bitTests(unittest.TestCase):
     # This was obtained on audace so the number might slightly change
     expected_rel_difference = 1.94
 
+    expected_memory_saving_ratio = 0.7
+
     prompt = "a beautiful sunset amidst the mountains."
     num_inference_steps = 10
     seed = 0
@@ -121,8 +123,10 @@ def setUp(self):
         )
 
     def tearDown(self):
-        del self.model_fp16
-        del self.model_8bit
+        if hasattr(self, "model_fp16"):
+            del self.model_fp16
+        if hasattr(self, "model_8bit"):
+            del self.model_8bit
 
         gc.collect()
         torch.cuda.empty_cache()
@@ -161,6 +165,28 @@ def test_memory_footprint(self):
         linear = get_some_linear_layer(self.model_8bit)
         self.assertTrue(linear.weight.__class__ == bnb.nn.Int8Params)
 
+    def test_model_memory_usage(self):
+        # Delete to not let anything interfere.
+        del self.model_8bit, self.model_fp16
+
+        # Re-instantiate.
+        inputs = self.get_dummy_inputs()
+        inputs = {
+            k: v.to(device=torch_device, dtype=torch.float16) for k, v in inputs.items() if not isinstance(v, bool)
+        }
+        model_fp16 = SD3Transformer2DModel.from_pretrained(
+            self.model_name, subfolder="transformer", torch_dtype=torch.float16
+        ).to(torch_device)
+        unquantized_model_memory = get_memory_consumption_stat(model_fp16, inputs)
+        del model_fp16
+
+        config = BitsAndBytesConfig(load_in_8bit=True)
+        model_8bit = SD3Transformer2DModel.from_pretrained(
+            self.model_name, subfolder="transformer", quantization_config=config, torch_dtype=torch.float16
+        )
+        quantized_model_memory = get_memory_consumption_stat(model_8bit, inputs)
+        assert unquantized_model_memory / quantized_model_memory >= self.expected_memory_saving_ratio
+
     def test_original_dtype(self):
         r"""
         A simple test to check if the model succesfully stores the original dtype
@@ -287,24 +313,6 @@ def test_device_and_dtype_assignment(self):
         # Check that this does not throw an error
         _ = self.model_fp16.cuda()
 
-    def test_model_memory_usage(self):
-        # Delete to not let anything interfere.
-        del self.model_4bit, self.model_fp16
-        
-        # Re-instantiate.
-        inputs = self.get_dummy_inputs()
-        model_fp16 = SD3Transformer2DModel.from_pretrained(
-            self.model_name, subfolder="transformer", torch_dtype=torch.float16
-        )
-        unquantized_model_memory = get_memory_consumption_stat(model_fp16, inputs)
-        config = BitsAndBytesConfig(load_in_8bit=True)
-        model_8bit = SD3Transformer2DModel.from_pretrained(
-            self.model_name, subfolder="transformer", quantization_config=config, device_map=torch_device
-        )
-        quantized_model_memory = get_memory_consumption_stat(model_8bit, inputs)
-        print(f"{unquantized_model_memory=}, {quantized_model_memory=}")
-        assert (1.0 - (unquantized_model_memory / quantized_model_memory)) >= 100.
-
 
 class Bnb8bitDeviceTests(Base8bitTests):
     def setUp(self) -> None:

tests/quantization/quanto/test_quanto.py

@@ -19,7 +19,7 @@
 
 if is_torch_available():
     import torch
-    
+
     from ..utils import LoRALayer, get_memory_consumption_stat
 
 
@@ -64,15 +64,20 @@ def test_quanto_layers(self):
                 assert isinstance(module, QLinear)
 
     def test_quanto_memory_usage(self):
-        unquantized_model = self.model_cls.from_pretrained(self.model_id, torch_dtype=self.torch_dtype)
         inputs = self.get_dummy_inputs()
+        inputs = {
+            k: v.to(device=torch_device, dtype=torch.bfloat16) for k, v in inputs.items() if not isinstance(v, bool)
+        }
+
+        unquantized_model = self.model_cls.from_pretrained(self.model_id, torch_dtype=self.torch_dtype)
+        unquantized_model.to(torch_device)
         unquantized_model_memory = get_memory_consumption_stat(unquantized_model, inputs)
 
         quantized_model = self.model_cls.from_pretrained(**self.get_dummy_model_init_kwargs())
+        quantized_model.to(torch_device)
         quantized_model_memory = get_memory_consumption_stat(quantized_model, inputs)
 
-        print(f"{unquantized_model_memory=}, {quantized_model_memory=}")
-        assert (1.0 - (unquantized_model_memory / quantized_model_memory)) >= self.expected_memory_reduction
+        assert unquantized_model_memory / quantized_model_memory >= self.expected_memory_reduction
 
     def test_keep_modules_in_fp32(self):
         r"""
@@ -292,14 +297,14 @@ def test_training(self):
 
 
 class FluxTransformerFloat8WeightsTest(FluxTransformerQuantoMixin, unittest.TestCase):
-    expected_memory_reduction = 0.3
+    expected_memory_reduction = 0.6
 
     def get_dummy_init_kwargs(self):
         return {"weights_dtype": "float8"}
 
 
 class FluxTransformerInt8WeightsTest(FluxTransformerQuantoMixin, unittest.TestCase):
-    expected_memory_reduction = 0.3
+    expected_memory_reduction = 0.6
     _test_torch_compile = True
 
     def get_dummy_init_kwargs(self):

tests/quantization/torchao/test_torchao.py

@@ -483,16 +483,21 @@ def test_memory_footprint(self):
             # there is additional overhead of scales and zero points
             self.assertTrue(total_bf16 < total_int4wo)
 
-    def test_memory_usage(self):
+    def test_model_memory_usage(self):
         model_id = "hf-internal-testing/tiny-flux-pipe"
-        inputs = self.get_dummy_inputs()
+        expected_memory_saving_ratio = 2.0
+
+        inputs = self.get_dummy_tensor_inputs(device=torch_device)
+
         transformer_bf16 = self.get_dummy_components(None, model_id=model_id)["transformer"]
+        transformer_bf16.to(torch_device)
         unquantized_model_memory = get_memory_consumption_stat(transformer_bf16, inputs)
+        del transformer_bf16
 
         transformer_int8wo = self.get_dummy_components(TorchAoConfig("int8wo"), model_id=model_id)["transformer"]
+        transformer_int8wo.to(torch_device)
         quantized_model_memory = get_memory_consumption_stat(transformer_int8wo, inputs)
-        print(f"{unquantized_model_memory=}, {quantized_model_memory=}")
-        assert (1.0 - (unquantized_model_memory / quantized_model_memory)) >= 100.
+        assert unquantized_model_memory / quantized_model_memory >= expected_memory_saving_ratio
 
     def test_wrong_config(self):
         with self.assertRaises(ValueError):

tests/quantization/utils.py

@@ -26,8 +26,7 @@ def __init__(self, module: nn.Module, rank: int):
 
         def forward(self, input, *args, **kwargs):
             return self.module(input, *args, **kwargs) + self.adapter(input)
-        
-    
+
     @torch.no_grad()
     @torch.inference_mode()
     def get_memory_consumption_stat(model, inputs):