update

Author: DN6

docs/source/en/quantization/quanto.md

@@ -47,6 +47,19 @@ image = pipe(
 image.save("output.png")
 ```
 
+## Skipping Quantization on specific modules
+
+It is possible to skip applying quantization on certain modules using the `modules_to_not_convert` argument in the `QuantoConfig`. Please ensure that the modules passed in to this argument match the keys of the modules in the `state_dict`  
+
+```python
+import torch
+from diffusers import FluxTransformer2DModel, QuantoConfig
+
+model_id = "black-forest-labs/FLUX.1-dev"
+quantization_config = QuantoConfig(weights="float8", modules_to_not_convert=["proj_out"])
+transformer = FluxTransformer2DModel.from_pretrained(model_id, quantization_config=quantization_config, torch_dtype=torch.bfloat16)
+```
+
 ## Using `from_single_file` with the Quanto Backend
 
 ```python

src/diffusers/models/modeling_utils.py

@@ -1036,11 +1036,11 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
                         )
 
                     named_buffers = model.named_buffers()
-
                     unexpected_keys = load_model_dict_into_meta(
                         model,
                         state_dict,
                         device=param_device,
+                        dtype=torch_dtype,
                         model_name_or_path=pretrained_model_name_or_path,
                         hf_quantizer=hf_quantizer,
                         keep_in_fp32_modules=keep_in_fp32_modules,

src/diffusers/quantizers/quanto/utils.py

@@ -2,9 +2,11 @@
 
 import torch.nn as nn
 
-from ...utils import is_accelerate_available
+from ...utils import is_accelerate_available, logging
 
 
+logger = logging.get_logger(__name__)
+
 if is_accelerate_available():
     from accelerate import init_empty_weights
 
@@ -47,5 +49,13 @@ def _replace_layers(model, quantization_config, modules_to_not_convert):
         return model
 
     model = _replace_layers(model, quantization_config, modules_to_not_convert)
+    has_been_replaced = any(isinstance(replaced_module, QLinear) for _, replaced_module in model.named_modules())
+
+    if not has_been_replaced:
+        logger.warning(
+            f"{model.__class__.__name__} does not appear to have any `nn.Linear` modules. Quantization will not be applied."
+            " Please check your model architecture, or submit an issue on Github if you think this is a bug."
+            " https://github.com/huggingface/diffusers"
+        )
 
     return model

tests/quantization/quanto/test_quanto.py

@@ -1,3 +1,4 @@
+import tempfile
 import unittest
 
 import torch
@@ -9,7 +10,6 @@
 from diffusers.utils import is_optimum_quanto_available
 from diffusers.utils.testing_utils import (
     nightly,
-    numpy_cosine_similarity_distance,
     require_accelerate,
     require_big_gpu_with_torch_cuda,
     torch_device,
@@ -29,6 +29,7 @@ class QuantoBaseTesterMixin:
     torch_dtype = torch.bfloat16
     expected_memory_use_in_gb = 5
     keep_in_fp32_module = ""
+    modules_to_not_convert = ""
 
     def get_dummy_init_kwargs(self):
         return {"weights": "float8"}
@@ -76,6 +77,22 @@ def test_keep_modules_in_fp32(self):
                     assert module.weight.dtype == torch.float32
         self.model_cls._keep_in_fp32_modules = _keep_in_fp32_modules
 
+    def test_modules_to_not_convert(self):
+        init_kwargs = self.get_dummy_model_init_kwargs()
+
+        quantization_config_kwargs = self.get_dummy_init_kwargs()
+        quantization_config_kwargs.update({"modules_to_not_convert": self.modules_to_not_convert})
+        quantization_config = QuantoConfig(**quantization_config_kwargs)
+
+        init_kwargs.update({"quantization_config": quantization_config})
+
+        model = self.model_cls.from_pretrained(**init_kwargs)
+        model.to("cuda")
+
+        for name, module in model.named_modules():
+            if name in self.modules_to_not_convert:
+                assert not isinstance(module, QLinear)
+
     def test_dtype_assignment(self):
         model = self.model_cls.from_pretrained(**self.get_dummy_model_init_kwargs())
         assert (model.get_memory_footprint() / 1024**3) < self.expected_memory_use_in_gb
@@ -99,12 +116,35 @@ def test_dtype_assignment(self):
         # This should work
         model.to("cuda")
 
+    def test_serialization(self):
+        model = self.model_cls.from_pretrained(**self.get_dummy_model_init_kwargs())
+        inputs = self.get_dummy_inputs()
+
+        model.to(torch_device)
+        with torch.no_grad():
+            model_output = model(**inputs)
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            model.save_pretrained(tmp_dir)
+            saved_model = self.model_cls.from_pretrained(
+                tmp_dir,
+                torch_dtype=torch.bfloat16,
+            )
+
+        saved_model.to(torch_device)
+        with torch.no_grad():
+            saved_model_output = saved_model(**inputs)
+
+        max_diff = torch.abs(model_output - saved_model_output).max()
+        assert max_diff < 1e-5
+
 
 class FluxTransformerQuantoMixin(QuantoBaseTesterMixin):
     model_id = "hf-internal-testing/tiny-flux-transformer"
     model_cls = FluxTransformer2DModel
     torch_dtype = torch.bfloat16
     keep_in_fp32_module = "proj_out"
+    modules_to_not_convert = ["proj_out"]
 
     def get_dummy_inputs(self):
         return {
@@ -130,14 +170,21 @@ def get_dummy_inputs(self):
         }
 
 
-class FluxTransformerFloat8(FluxTransformerQuantoMixin, unittest.TestCase):
+class FluxTransformerFloat8WeightsTest(FluxTransformerQuantoMixin, unittest.TestCase):
     expected_memory_use_in_gb = 10
 
     def get_dummy_init_kwargs(self):
         return {"weights": "float8"}
 
 
-class FluxTransformerInt8(FluxTransformerQuantoMixin, unittest.TestCase):
+class FluxTransformerFloat8WeightsAndActivationTest(FluxTransformerQuantoMixin, unittest.TestCase):
+    expected_memory_use_in_gb = 10
+
+    def get_dummy_init_kwargs(self):
+        return {"weights": "float8", "activations": "float8"}
+
+
+class FluxTransformerInt8WeightsTest(FluxTransformerQuantoMixin, unittest.TestCase):
     expected_memory_use_in_gb = 10
 
     def get_dummy_init_kwargs(self):
@@ -157,20 +204,42 @@ def test_torch_compile(self):
         with torch.no_grad():
             compiled_model_output = compiled_model(**inputs).sample
 
-        max_diff = numpy_cosine_similarity_distance(
-            model_output.cpu().flatten(), compiled_model_output.cpu().flatten()
-        )
+        max_diff = torch.abs(model_output - compiled_model_output).max()
+        assert max_diff < 1e-4
+
+
+class FluxTransformerInt8WeightsAndActivationTest(FluxTransformerQuantoMixin, unittest.TestCase):
+    expected_memory_use_in_gb = 10
+
+    def get_dummy_init_kwargs(self):
+        return {"weights": "int8", "activations": "int8"}
+
+    def test_torch_compile(self):
+        model = self.model_cls.from_pretrained(**self.get_dummy_model_init_kwargs())
+        compiled_model = torch.compile(model, mode="max-autotune", fullgraph=True)
+        inputs = self.get_dummy_inputs()
+
+        model.to(torch_device)
+        with torch.no_grad():
+            model_output = model(**inputs).sample
+        model.to("cpu")
+
+        compiled_model.to(torch_device)
+        with torch.no_grad():
+            compiled_model_output = compiled_model(**inputs).sample
+
+        max_diff = torch.abs(model_output - compiled_model_output).max()
         assert max_diff < 1e-4
 
 
-class FluxTransformerInt4(FluxTransformerQuantoMixin, unittest.TestCase):
+class FluxTransformerInt4WeightsTest(FluxTransformerQuantoMixin, unittest.TestCase):
     expected_memory_use_in_gb = 6
 
     def get_dummy_init_kwargs(self):
         return {"weights": "int4"}
 
 
-class FluxTransformerInt2(FluxTransformerQuantoMixin, unittest.TestCase):
+class FluxTransformerInt2WeightsTest(FluxTransformerQuantoMixin, unittest.TestCase):
     expected_memory_use_in_gb = 6
 
     def get_dummy_init_kwargs(self):