memory usage tests

Author: sayakpaul

tests/quantization/__init__.py

@@ -54,29 +54,8 @@ def get_some_linear_layer(model):
 
 if is_torch_available():
     import torch
-    import torch.nn as nn
 
-    class LoRALayer(nn.Module):
-        """Wraps a linear layer with LoRA-like adapter - Used for testing purposes only
-
-        Taken from
-        https://github.com/huggingface/transformers/blob/566302686a71de14125717dea9a6a45b24d42b37/tests/quantization/bnb/test_4bit.py#L62C5-L78C77
-        """
-
-        def __init__(self, module: nn.Module, rank: int):
-            super().__init__()
-            self.module = module
-            self.adapter = nn.Sequential(
-                nn.Linear(module.in_features, rank, bias=False),
-                nn.Linear(rank, module.out_features, bias=False),
-            )
-            small_std = (2.0 / (5 * min(module.in_features, module.out_features))) ** 0.5
-            nn.init.normal_(self.adapter[0].weight, std=small_std)
-            nn.init.zeros_(self.adapter[1].weight)
-            self.adapter.to(module.weight.device)
-
-        def forward(self, input, *args, **kwargs):
-            return self.module(input, *args, **kwargs) + self.adapter(input)
+    from ..utils import LoRALayer, get_memory_consumption_stat
 
 
 if is_bitsandbytes_available():
@@ -350,6 +329,29 @@ 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_4bit.py

@@ -60,29 +60,8 @@ def get_some_linear_layer(model):
 
 if is_torch_available():
     import torch
-    import torch.nn as nn
 
-    class LoRALayer(nn.Module):
-        """Wraps a linear layer with LoRA-like adapter - Used for testing purposes only
-
-        Taken from
-        https://github.com/huggingface/transformers/blob/566302686a71de14125717dea9a6a45b24d42b37/tests/quantization/bnb/test_8bit.py#L62C5-L78C77
-        """
-
-        def __init__(self, module: nn.Module, rank: int):
-            super().__init__()
-            self.module = module
-            self.adapter = nn.Sequential(
-                nn.Linear(module.in_features, rank, bias=False),
-                nn.Linear(rank, module.out_features, bias=False),
-            )
-            small_std = (2.0 / (5 * min(module.in_features, module.out_features))) ** 0.5
-            nn.init.normal_(self.adapter[0].weight, std=small_std)
-            nn.init.zeros_(self.adapter[1].weight)
-            self.adapter.to(module.weight.device)
-
-        def forward(self, input, *args, **kwargs):
-            return self.module(input, *args, **kwargs) + self.adapter(input)
+    from ..utils import LoRALayer, get_memory_consumption_stat
 
 
 if is_bitsandbytes_available():
@@ -248,7 +227,7 @@ def test_llm_skip(self):
         self.assertTrue(linear.weight.dtype == torch.int8)
         self.assertTrue(isinstance(linear, bnb.nn.Linear8bitLt))
 
-        self.assertTrue(isinstance(model_8bit.proj_out, nn.Linear))
+        self.assertTrue(isinstance(model_8bit.proj_out, torch.nn.Linear))
         self.assertTrue(model_8bit.proj_out.weight.dtype != torch.int8)
 
     def test_config_from_pretrained(self):
@@ -308,6 +287,24 @@ 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/bnb/test_mixed_int8.py

@@ -19,29 +19,8 @@
 
 if is_torch_available():
     import torch
-    import torch.nn as nn
-
-    class LoRALayer(nn.Module):
-        """Wraps a linear layer with LoRA-like adapter - Used for testing purposes only
-
-        Taken from
-        https://github.com/huggingface/transformers/blob/566302686a71de14125717dea9a6a45b24d42b37/tests/quantization/bnb/test_4bit.py#L62C5-L78C77
-        """
-
-        def __init__(self, module: nn.Module, rank: int):
-            super().__init__()
-            self.module = module
-            self.adapter = nn.Sequential(
-                nn.Linear(module.in_features, rank, bias=False),
-                nn.Linear(rank, module.out_features, bias=False),
-            )
-            small_std = (2.0 / (5 * min(module.in_features, module.out_features))) ** 0.5
-            nn.init.normal_(self.adapter[0].weight, std=small_std)
-            nn.init.zeros_(self.adapter[1].weight)
-            self.adapter.to(module.weight.device)
-
-        def forward(self, input, *args, **kwargs):
-            return self.module(input, *args, **kwargs) + self.adapter(input)
+    
+    from ..utils import LoRALayer, get_memory_consumption_stat
 
 
 @nightly
@@ -86,19 +65,14 @@ def test_quanto_layers(self):
 
     def test_quanto_memory_usage(self):
         unquantized_model = self.model_cls.from_pretrained(self.model_id, torch_dtype=self.torch_dtype)
-        unquantized_model_memory = unquantized_model.get_memory_footprint() / 1024**3
-
-        model = self.model_cls.from_pretrained(**self.get_dummy_model_init_kwargs())
         inputs = self.get_dummy_inputs()
+        unquantized_model_memory = get_memory_consumption_stat(unquantized_model, inputs)
 
-        torch.cuda.reset_peak_memory_stats()
-        torch.cuda.empty_cache()
+        quantized_model = self.model_cls.from_pretrained(**self.get_dummy_model_init_kwargs())
+        quantized_model_memory = get_memory_consumption_stat(quantized_model, inputs)
 
-        model.to(torch_device)
-        with torch.no_grad():
-            model(**inputs)
-        max_memory = torch.cuda.max_memory_allocated() / 1024**3
-        assert (1.0 - (max_memory / unquantized_model_memory)) >= self.expected_memory_reduction
+        print(f"{unquantized_model_memory=}, {quantized_model_memory=}")
+        assert (1.0 - (unquantized_model_memory / quantized_model_memory)) >= self.expected_memory_reduction
 
     def test_keep_modules_in_fp32(self):
         r"""

tests/quantization/quanto/__init__.py

@@ -50,27 +50,7 @@
     import torch
     import torch.nn as nn
 
-    class LoRALayer(nn.Module):
-        """Wraps a linear layer with LoRA-like adapter - Used for testing purposes only
-
-        Taken from
-        https://github.com/huggingface/transformers/blob/566302686a71de14125717dea9a6a45b24d42b37/tests/quantization/bnb/test_4bit.py#L62C5-L78C77
-        """
-
-        def __init__(self, module: nn.Module, rank: int):
-            super().__init__()
-            self.module = module
-            self.adapter = nn.Sequential(
-                nn.Linear(module.in_features, rank, bias=False),
-                nn.Linear(rank, module.out_features, bias=False),
-            )
-            small_std = (2.0 / (5 * min(module.in_features, module.out_features))) ** 0.5
-            nn.init.normal_(self.adapter[0].weight, std=small_std)
-            nn.init.zeros_(self.adapter[1].weight)
-            self.adapter.to(module.weight.device)
-
-        def forward(self, input, *args, **kwargs):
-            return self.module(input, *args, **kwargs) + self.adapter(input)
+    from ..utils import LoRALayer, get_memory_consumption_stat
 
 
 if is_torchao_available():
@@ -503,6 +483,17 @@ 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):
+        model_id = "hf-internal-testing/tiny-flux-pipe"
+        inputs = self.get_dummy_inputs()
+        transformer_bf16 = self.get_dummy_components(None, model_id=model_id)["transformer"]
+        unquantized_model_memory = get_memory_consumption_stat(transformer_bf16, inputs)
+
+        transformer_int8wo = self.get_dummy_components(TorchAoConfig("int8wo"), model_id=model_id)["transformer"]
+        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.
+
     def test_wrong_config(self):
         with self.assertRaises(ValueError):
             self.get_dummy_components(TorchAoConfig("int42"))

tests/quantization/quanto/test_quanto.py

@@ -0,0 +1,39 @@
+from diffusers.utils import is_torch_available
+
+
+if is_torch_available():
+    import torch
+    import torch.nn as nn
+
+    class LoRALayer(nn.Module):
+        """Wraps a linear layer with LoRA-like adapter - Used for testing purposes only
+
+        Taken from
+        https://github.com/huggingface/transformers/blob/566302686a71de14125717dea9a6a45b24d42b37/tests/quantization/bnb/test_4bit.py#L62C5-L78C77
+        """
+
+        def __init__(self, module: nn.Module, rank: int):
+            super().__init__()
+            self.module = module
+            self.adapter = nn.Sequential(
+                nn.Linear(module.in_features, rank, bias=False),
+                nn.Linear(rank, module.out_features, bias=False),
+            )
+            small_std = (2.0 / (5 * min(module.in_features, module.out_features))) ** 0.5
+            nn.init.normal_(self.adapter[0].weight, std=small_std)
+            nn.init.zeros_(self.adapter[1].weight)
+            self.adapter.to(module.weight.device)
+
+        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):
+        torch.cuda.reset_peak_memory_stats()
+        torch.cuda.empty_cache()
+
+        model(**inputs)
+        max_memory_mem_allocated = torch.cuda.max_memory_allocated()
+        return max_memory_mem_allocated