[compressed-tensors] keep weights compressed unless dequantize=True

Author: ArthurZucker

src/transformers/quantizers/quantizer_compressed_tensors.py

@@ -111,26 +111,12 @@ def _process_model_before_weight_loading(self, model, **kwargs):
                 ignore=ct_config.ignore,
                 modules_to_not_convert=self.modules_to_not_convert,
             )
-            apply_quantization_config(model, ct_config, run_compressed=False)
-            # Packed formats (e.g. int4 `weight_packed`) need the compressed module layout to
-            # receive the checkpoint tensors.
-            if self.quantization_config.is_quantization_compressed:
-                self.compressor.compress_model(model=model)
-            return
 
-        # No kernel runs these weights compressed: the `deserialize` load operation decompresses them, and
-        # compressed-tensors only sets up what quantizes at inference (activations, kv cache)
-        from compressed_tensors.quantization import QuantizationStatus
-
-        ct_config.quantization_status = QuantizationStatus.FROZEN
-        for group in ct_config.config_groups.values():
-            group.weights = None
-        ct_config.config_groups = {
-            name: group
-            for name, group in ct_config.config_groups.items()
-            if group.input_activations is not None or group.output_activations is not None
-        }
         apply_quantization_config(model, ct_config, run_compressed=False)
+        # Packed formats (e.g. int4 `weight_packed`) need the compressed module layout to
+        # receive the checkpoint tensors.
+        if self.quantization_config.is_quantization_compressed:
+            self.compressor.compress_model(model=model)
 
     @property
     def is_trainable(self):
@@ -182,5 +168,5 @@ def get_weight_conversions(self):
     }
 
     def get_load_operations(self):
-        # Without kernels nothing runs compressed, so deserializing decompresses as well
-        return {"deserialize": CompressedTensorsDequantize, "dequantize": CompressedTensorsDequantize}
+        # No `deserialize`: the serialized tensors load as is into the params compressed-tensors creates
+        return {"dequantize": CompressedTensorsDequantize}

tests/quantization/compressed_tensors_integration/test_compressed_models.py

@@ -105,8 +105,9 @@ def test_no_warnings_for_all_models(self):
 @require_compressed_tensors
 @require_torch
 class DequantizeTest(unittest.TestCase):
-    """Non-FP8 schemes have no kernels, so they run dequantized: the weights are decompressed while loading.
-    `dequantize=True` (or the deprecated `run_compressed=False`) also drops the activation quantization."""
+    """Non-FP8 schemes have no kernels, so they run dequantized. `dequantize=True` (or the
+    deprecated `run_compressed=False`) does it at load time; by default the weights are left
+    compressed and compressed-tensors decompresses them on the first forward pass."""
 
     tinyllama_w4a16 = "nm-testing/tinyllama-w4a16-compressed"
     tinyllama_w8a8 = "nm-testing/tinyllama-w8a8-compressed"
@@ -120,19 +121,26 @@ def tearDown(self):
         backend_empty_cache(torch_device)
         gc.collect()
 
-    def test_default_load_decompresses_while_loading(self):
+    def test_default_load_stays_compressed_until_first_forward(self):
         from compressed_tensors import QuantizationStatus
 
         for stub in self.stubs:
             model = AutoModelForCausalLM.from_pretrained(stub)
 
+            compressed_count = sum(
+                1 for m in model.modules() if getattr(m, "quantization_status", None) == QuantizationStatus.COMPRESSED
+            )
+            self.assertGreater(compressed_count, 0, "weights must not be decompressed at load time")
+
+            model(input_ids=torch.tensor([[0, 1, 2]]))
+
+            # The hook compressed-tensors registered decompressed the model on that forward pass.
             compressed_count = sum(
                 1 for m in model.modules() if getattr(m, "quantization_status", None) == QuantizationStatus.COMPRESSED
             )
             self.assertEqual(compressed_count, 0, "no modules should be left in COMPRESSED state")
             self.assertFalse(any("weight_packed" in name for name, _ in model.named_parameters()))
             self.assertTrue(model.model.layers[0].self_attn.q_proj.weight.is_floating_point())
-            model(input_ids=torch.tensor([[0, 1, 2]]))
 
     def test_deprecated_run_compressed_maps_to_dequantize(self):
         # `run_compressed=False` is deprecated and must behave as `dequantize=True`.
@@ -166,17 +174,10 @@ def test_dequantize_outputs_match_default(self):
         self.assertEqual(tokenizer.decode(output_default[0]), tokenizer.decode(output_dequantized[0]))
 
     def test_dequantize_drops_activation_quantization(self):
-        """`dequantize=True` loads a dense model: w8a8 also stops quantizing its activations, so its outputs differ
-        from the default (quantized activations) load."""
-        from transformers import AutoTokenizer
-
-        tokenizer = AutoTokenizer.from_pretrained(self.tinyllama_w8a8)
-        input_ids = tokenizer(self.prompt, return_tensors="pt").input_ids
+        """`dequantize=True` loads a plain dense model: no compressed-tensors scheme is left, so w8a8 also stops
+        quantizing its activations."""
         model = AutoModelForCausalLM.from_pretrained(
             self.tinyllama_w8a8, quantization_config=CompressedTensorsConfig(dequantize=True)
         )
-        output = model.generate(input_ids, max_new_tokens=20, do_sample=False)
-        self.assertEqual(
-            tokenizer.decode(output[0]),
-            "<s> Paris is the capital of which country?\n\n  * a) France\n  * b) Germany\n  * c) Italy\n",
-        )
+        self.assertFalse(any(hasattr(module, "quantization_scheme") for module in model.modules()))
+        self.assertTrue(model.model.layers[0].self_attn.q_proj.weight.is_floating_point())

tests/quantization/compressed_tensors_integration/test_compressed_tensors.py

@@ -244,20 +244,23 @@ def test_tinyllama_fp8_per_tensor_save_reload(self):
             self.assertLessEqual(torch.exp(outputs.loss), 20.0)
 
     def test_use_optimized_inference_is_opt_in(self):
-        """Without `use_optimized_inference=True`, an FP8 checkpoint runs no kernels: its weights are
-        decompressed while loading."""
+        """Without `use_optimized_inference=True`, an FP8 checkpoint takes the regular compressed-tensors
+        route: no kernels, and with `dequantize=False` the weights are left compressed until the
+        first forward pass decompresses them."""
         from transformers.integrations.compressed_tensors import CompressedTensorsFP8Linear
 
         model = AutoModelForCausalLM.from_pretrained(self.tinyllama_fp8, device_map=torch_device)
         module = model.model.layers[0].self_attn.q_proj
 
         self.assertNotIsInstance(module, CompressedTensorsFP8Linear)
-        self.assertEqual(module.weight.dtype, model.dtype, "weights are decompressed while loading")
+        self.assertEqual(module.weight.dtype, torch.float8_e4m3fn, "weights must not be decompressed at load time")
 
         tokenizer = AutoTokenizer.from_pretrained(self.tinyllama_fp8)
         inputs = tokenizer(self.prompt, return_tensors="pt").to(model.device)
         with torch.no_grad():
             outputs = model(**inputs, labels=inputs["input_ids"])
+        # compressed-tensors decompressed the model on that first forward pass
+        self.assertEqual(module.weight.dtype, model.dtype)
         self.assertLessEqual(torch.exp(outputs.loss), 20.0)
 
     def test_non_fp8_model_unaffected(self):