fix kv cache serialization, add tests

Signed-off-by: Kyle Sayers <[email protected]>

Author: kylesayrs

src/compressed_tensors/quantization/quant_config.py

@@ -171,6 +171,7 @@ def from_pretrained(
         :param model: model to calculate quantization scheme of
         :return: filled out QuantizationScheme for the input model
         """
+        from compressed_tensors.modeling import IMPL_ATTR, KV_CACHE_ATTR
         from compressed_tensors.quantization.lifecycle.initialize import (
             is_attention_module,
         )
@@ -196,24 +197,35 @@ def from_pretrained(
         for name, submodule in model.named_modules():
             layer_type: str = module_type(submodule)
 
-            if is_module_quantized(submodule):
+            # add config group if quantized non-attention or attention quant
+            has_config_group = is_module_quantized(submodule) and (
+                not is_attention_module(submodule) or hasattr(submodule, IMPL_ATTR)
+            )
+            # only add kvcache if quant attention (which always implies kvcache)
+            has_kv_cache = is_module_quantized(submodule) and is_attention_module(
+                submodule
+            )
+
+            if has_config_group:
                 # add to running set of schemes/layer_type_names
                 model_status = getattr(submodule, "quantization_status", model_status)
                 quantization_type_names.add(layer_type)
                 if submodule.quantization_scheme not in quantization_schemes:
                     quantization_schemes.append(submodule.quantization_scheme)
 
-                # attention quantization implies kv cache quantization
-                if is_attention_module(submodule):
-                    kv_cache_scheme = submodule.quantization_scheme.input_activations
+            if has_kv_cache:
+                model_status = getattr(submodule, "quantization_status", model_status)
+                kv_cache_scheme = submodule.quantization_scheme.input_activations
 
-            else:
+            if not has_config_group:
                 # add non-quantized layers to the ignore list
                 if layer_type not in ignore:
                     ignore[layer_type] = []
                 ignore[layer_type].append(name)
 
-        if len(quantization_schemes) == 0:  # No quantized layers
+        if (
+            len(quantization_schemes) == 0 and kv_cache_scheme is None
+        ):  # No quantized layers
             return None
 
         # create ignore list, only include layers whose class has ever been targeted

tests/test_modeling/test_attention_and_cache.py

@@ -56,7 +56,10 @@ def test_attention_cache():
     k_called = [False for _ in range(len(layers))]
     v_called = [False for _ in range(len(layers))]
 
+    # apply attention quantization
     _apply_attention(model, layers, q_called, k_called, v_called)
+
+    # check attention quantization
     outputs = model(**inputs)
     assert torch.equal(outputs.logits, true_outputs.logits)
     assert all(q_called) and all(k_called) and all(v_called)

tests/test_quantization/lifecycle/test_apply.py

@@ -131,6 +131,64 @@ def test_apply_quantization_config_tinyllama():
             )
 
 
+@pytest.mark.parametrize(
+    "config",
+    [
+        QuantizationConfig(
+            config_groups={
+                "linear": QuantizationScheme(
+                    targets=["Linear"],
+                    input_activations=QuantizationArgs(
+                        num_bits=8, type="float", strategy="tensor"
+                    ),
+                )
+            }
+        ),
+        QuantizationConfig(
+            config_groups={
+                "linear": QuantizationScheme(
+                    targets=["Linear"],
+                    input_activations=QuantizationArgs(
+                        num_bits=8, type="float", strategy="tensor"
+                    ),
+                )
+            },
+            ignore=[
+                "model.layers.0.self_attn.q_proj",
+                "model.layers.1.self_attn.k_proj",
+                "model.layers.2.self_attn.v_proj",
+            ],
+        ),
+        QuantizationConfig(
+            config_groups={},
+            kv_cache_scheme=QuantizationArgs(
+                num_bits=8, type="float", strategy="tensor"
+            ),
+        ),
+        QuantizationConfig(
+            config_groups={
+                "attention": QuantizationScheme(
+                    targets=["LlamaAttention"],
+                    input_activations=QuantizationArgs(
+                        num_bits=8, type="float", strategy="tensor"
+                    ),
+                )
+            },
+            kv_cache_scheme=QuantizationArgs(
+                num_bits=8, type="float", strategy="tensor"
+            ),
+        ),
+    ],
+)
+def test_from_pretrained(config: QuantizationConfig):
+    model = AutoModelForCausalLM.from_pretrained("nm-testing/llama2.c-stories15M")
+    apply_quantization_config(model, config)
+    _config = QuantizationConfig.from_pretrained(model)
+    assert list(_config.config_groups.values()) == list(config.config_groups.values())
+    assert _config.kv_cache_scheme == config.kv_cache_scheme
+    assert _config.ignore == config.ignore
+
+
 def test_serialize_config_tinyllama():
     quant_config = get_sample_tinyllama_quant_config()
     model = get_tinyllama_model()

tests/test_quantization/test_quant_config.py

@@ -16,11 +16,14 @@
 from compressed_tensors.quantization import (
     DEFAULT_QUANTIZATION_FORMAT,
     DEFAULT_QUANTIZATION_METHOD,
+    QuantizationArgs,
     QuantizationConfig,
     QuantizationScheme,
     QuantizationStatus,
+    apply_quantization_config,
 )
 from pydantic import ValidationError
+from transformers import AutoModelForCausalLM
 
 
 def test_basic_config():