remove apply_quantization_status

Signed-off-by: Brian Dellabetta <[email protected]>

Author: brian-dellabetta

src/compressed_tensors/quantization/lifecycle/apply.py

@@ -45,7 +45,6 @@
 __all__ = [
     "load_pretrained_quantization_parameters",
     "apply_quantization_config",
-    "apply_quantization_status",
     "find_name_or_class_matches",
 ]
 
@@ -163,8 +162,14 @@ def apply_quantization_config(
             )
             replace_module(model, name, compressed_linear)
 
-        # apply current quantization status to each targeted submodule
-        apply_quantization_status(submodule, config.quantization_status)
+        else:
+            initialize_module_for_quantization(
+                submodule,
+                force_zero_point=config.quantization_status
+                != QuantizationStatus.COMPRESSED,
+            )
+
+        submodule.quantization_status = config.quantization_status
 
 
 def process_quantization_config(config: QuantizationConfig) -> QuantizationConfig:
@@ -203,21 +208,6 @@ def process_kv_cache_config(
     return config
 
 
-def apply_quantization_status(module: Module, status: QuantizationStatus):
-    """
-    Applies in place the quantization lifecycle up to the given status
-
-    :param module: module to apply quantization to
-    :param status: status to update the module to
-    """
-
-    force_zero_point_init = status != QuantizationStatus.COMPRESSED
-
-    initialize_module_for_quantization(module, force_zero_point=force_zero_point_init)
-
-    module.quantization_status = status
-
-
 @deprecated(
     message="This function is deprecated and will be removed in a future release."
     "Please use `match_targets` from `compressed_tensors.utils.match` instead."

src/compressed_tensors/quantization/lifecycle/initialize.py

@@ -75,14 +75,14 @@ def initialize_module_for_quantization(
     :param force_zero_point: whether to force initialization of a zero point for
         symmetric quantization
     """
-    _clear_all_qparams(module)
-
     # TODO: don't initialize parameters when running decompression
     scheme = scheme or getattr(module, "quantization_scheme", None)
     if scheme is None:
         # no scheme passed and layer not targeted for quantization - skip
         return
 
+    _clear_all_qparams(module)
+
     if is_attention_module(module):
         # quantized actions based on calltime status
         _initialize_attn_scales(module)

tests/test_quantization/lifecycle/test_apply.py

@@ -26,6 +26,7 @@
     QuantizationStatus,
 )
 from compressed_tensors.quantization.lifecycle import apply_quantization_config
+from compressed_tensors.utils import match_named_modules
 from tests.testing_utils import requires_accelerate
 from transformers import AutoModelForCausalLM
 
@@ -103,60 +104,27 @@ def test_target_prioritization(mock_frozen):
 
 def test_apply_quantization_config_tinyllama():
     quant_config = get_sample_tinyllama_quant_config(
-        status=QuantizationStatus.CALIBRATION
+        status=QuantizationStatus.INITIALIZED
     )
     model = get_tinyllama_model()
 
     # check that model is not already quantized
     for module in model.modules():
         _test_layer_quantization_status(module, inputs=False, weights=False)
 
-    count_layer_names = ("Linear", "Embeddidng", "LlamaRotaryEmbedding")
-    count_layer_num = defaultdict(int)
-
-    for name, module in model.named_modules():
-        if name in quant_config.ignore:
-            continue
-        module_type = module.__class__.__name__
-        if module_type in count_layer_names:
-            count_layer_num[module_type] += 1
-
-    assert len(count_layer_num) > 0, f"None of {count_layer_names} found in model"
-    assert all(value > 0 for value in count_layer_num.values())
-
     # apply quant config to model
     apply_quantization_config(model, quant_config)
 
     # check for correct application of quant config
-    for name, module in model.named_modules():
-        if name in quant_config.ignore:
-            continue
-        module_type = module.__class__.__name__
-        if module_type in count_layer_names:
-            count_layer_num[module_type] -= 1
-            _inputs = module_type == "Linear"
-            _weights = not module_type == "LlamaRotaryEmbedding"
-            _test_layer_quantization_status(module, inputs=_inputs, weights=_weights)
-
-    assert all(
-        value == 0 for value in count_layer_num.values()
-    ), "Not all values are zero"
-
-    # test quantization compression
-    # sample forward pass to fill scales, zps
-    model(torch.zeros((1, 1), dtype=int), torch.zeros((1, 1), dtype=int))
-    quant_config.quantization_status = QuantizationStatus.COMPRESSED
-    apply_quantization_config(model, quant_config)
-    for name, module in model.named_modules():
-        if name in quant_config.ignore:
-            continue
-        module_type = module.__class__.__name__
-        if module_type == "Linear":
+    for quant_scheme in quant_config.config_groups.values():
+        for name, module in match_named_modules(
+            model, quant_scheme.targets, quant_config.ignore
+        ):
             _test_layer_quantization_status(
                 module,
-                inputs=True,
-                weights=True,
-                expected_status=QuantizationStatus.COMPRESSED,
+                inputs=quant_scheme.input_activations is not None,
+                weights=quant_scheme.weights is not None,
+                expected_status=QuantizationStatus.INITIALIZED,
             )