use qkv hooks

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

Author: kylesayrs

src/llmcompressor/modifiers/quantization/calibration.py

@@ -1,4 +1,5 @@
-from typing import Any, Dict, Optional, Tuple
+from functools import partial
+from typing import TYPE_CHECKING, Any, Dict, Optional, Set, Tuple
 
 import torch
 from compressed_tensors.quantization import (
@@ -13,18 +14,24 @@
 from compressed_tensors.utils import align_module_device, update_parameter_data
 from loguru import logger
 from torch.nn import Module
+from torch.utils.hooks import RemovableHandle
 
 from llmcompressor.modifiers.quantization.cache import QuantizedKVParameterCache
 from llmcompressor.observers import Observer
 from llmcompressor.utils.helpers import getattr_chain
 
+if TYPE_CHECKING:
+    from llmcompressor.modifiers.utils.hooks import HooksMixin
+
+
 DEFAULT_MAXSHRINK = 0.20
 DEFAULT_PATIENCE = 5
 DEFAULT_AVERAGING_CONSTANT = 0.01
 DEFAULT_GRID = 100.0
 DEFAULT_NORM = 2.4
 
 __all__ = [
+    "register_calibrate_attn_hooks",
     "initialize_observer",
     "update_weight_zp_scale",
     "calibrate_input_hook",
@@ -205,14 +212,30 @@ def calibrate_activations(module: Module, value: torch.Tensor, base_name: str):
     )
 
 
-def calibrate_input_hook(module: Module, args: Any):
+def register_calibrate_attn_hooks(
+    modifier: HooksMixin, attention_impl
+) -> Set[RemovableHandle]:
+    return {
+        modifier.register_hook(
+            attention_impl, partial(calibrate_input_hook, basename="q"), "query"
+        ),
+        modifier.register_hook(
+            attention_impl, partial(calibrate_input_hook, basename="k"), "key"
+        ),
+        modifier.register_hook(
+            attention_impl, partial(calibrate_input_hook, basename="v"), "value"
+        ),
+    }
+
+
+def calibrate_input_hook(module: Module, args: Any, base_name: str = "input"):
     """
     Hook to calibrate input activations.
     Will call the observers to update the scales/zp before applying
     input QDQ in the module's forward pass.
     """
     args = args[0] if isinstance(args, tuple) else args
-    calibrate_activations(module, value=args, base_name="input")
+    calibrate_activations(module, value=args, base_name=base_name)
 
 
 def calibrate_output_hook(module: Module, _args: Any, output: torch.Tensor):
@@ -290,14 +313,6 @@ def initialize_attention_observers(module: Module):
         initialize_observer(module, "v", input_args)
 
 
-def calibrate_attention(
-    module: Module, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor
-):
-    calibrate_activations(module, value=query, base_name="q")
-    calibrate_activations(module, value=key, base_name="k")
-    calibrate_activations(module, value=value, base_name="v")
-
-
 def apply_calibration_status(module: Module):
     scheme = getattr(module, "quantization_scheme", None)
     if not scheme:

src/llmcompressor/modifiers/quantization/quantization/mixin.py

@@ -244,7 +244,7 @@ def _initialize_hooks(self, model: torch.nn.Module) -> Set[RemovableHandle]:
 
         # TODO: attnq
         # attention_impl = enable_compressed_attention(model)
-        # hooks.add(self.register_hook(attention_impl, calibrate_attention, "calib"))
+        # hooks |= register_calibrate_attn_hooks(self, attention_impl)
 
         for module in model.modules():
             if not hasattr(module, "quantization_scheme"):

src/llmcompressor/modifiers/transform/spinquant/base.py

@@ -215,7 +215,40 @@ def _create_r2_scheme(self, model: PreTrainedModel) -> TransformScheme:
         )
 
     def _create_r3_scheme(self) -> TransformScheme:
-        raise NotImplementedError()
+        return (
+            TransformScheme(
+                type=self.transform_type,
+                randomize=self.randomize,
+                requires_grad=self.learnable,
+                apply=[
+                    TransformArgs(
+                        targets=[self.mappings.attn],
+                        location="attn_q",
+                    ),
+                    TransformArgs(
+                        targets=[self.mappings.attn],
+                        location="attn_k",
+                    ),
+                ],
+            ),
+        )
 
     def _create_r4_scheme(self) -> TransformScheme:
-        raise NotImplementedError()
+        return (
+            TransformScheme(
+                type=self.transform_type,
+                randomize=self.randomize,
+                requires_grad=self.learnable,
+                apply=[
+                    TransformArgs(
+                        targets=[*self.mappings.mlp_out],
+                        location="input",
+                    ),
+                    TransformArgs(
+                        targets=[*self.mappings.mlp_out],
+                        location="weight_input",
+                        inverse=True,
+                    ),
+                ],
+            ),
+        )

src/llmcompressor/modifiers/transform/spinquant/mappings.py

@@ -10,6 +10,7 @@
 class SpinQuantMapping(BaseModel):
     embedding: str
 
+    attn: str
     attn_q: str
     attn_k: str
     attn_v: str
@@ -31,6 +32,7 @@ def cast_to_list(cls, value):
 
 _default_mappings = SpinQuantMapping(
     embedding="re:.*embed_tokens$",
+    attn="re:.*self_attn$",
     attn_q="re:.*q_proj$",
     attn_k="re:.*k_proj$",
     attn_v="re:.*v_proj$",