pass config for init later

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

Author: kylesayrs

src/compressed_tensors/modeling/attention.py

@@ -31,7 +31,7 @@
 from torch import Tensor
 from torch.nn import Module
 from torch.utils.hooks import RemovableHandle
-from transformers import AttentionInterface, PreTrainedModel
+from transformers import AttentionInterface, PretrainedConfig, PreTrainedModel
 from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
 
 
@@ -62,8 +62,9 @@ class QuantizedAttentionImpl(InternalModule):
     :param attn_module: parent attention module
     """
 
-    def __init__(self, attn_module: Module):
+    def __init__(self, config: PretrainedConfig, attn_module: Module):
         super().__init__()
+        self.config = config
         self.attn_module = ref(attn_module)  # avoid circular references
         self._qparams_initialized = False
 
@@ -143,7 +144,7 @@ def initialize_hooked_attention(
     :param quantize: initialize attention quantization parameters
     """
     if not hasattr(module, IMPL_ATTR):
-        module.register_module(IMPL_ATTR, QuantizedAttentionImpl(module))
+        module.register_module(IMPL_ATTR, QuantizedAttentionImpl(model.config, module))
         if model.config._attn_implementation != HOOKED_ATTENTION_NAME:
             # assumes only one model at a time
             global _original_impl

src/compressed_tensors/modeling/kvcache.py

@@ -25,7 +25,7 @@
 from torch import Tensor
 from torch.nn import Module
 from torch.utils.hooks import RemovableHandle
-from transformers import Cache, PreTrainedModel
+from transformers import Cache, PretrainedConfig, PreTrainedModel
 
 
 __all__ = [
@@ -53,8 +53,9 @@ class QuantizedKVCache(InternalModule):
     :param attn_module: parent attention module
     """
 
-    def __init__(self, attn_module: Module):
+    def __init__(self, config: PretrainedConfig, attn_module: Module):
         super().__init__()
+        self.config = config
         self.attn_module = ref(attn_module)  # avoid circular reference
         self.past_key_values: Optional[Cache] = None
         self._qparams_initialized = False
@@ -134,7 +135,7 @@ def initialize_hooked_kv_cache(
     :param quantize: initialize kv cache quantization parameters
     """
     if not hasattr(module, KV_CACHE_ATTR):
-        module.register_module(KV_CACHE_ATTR, QuantizedKVCache(module))
+        module.register_module(KV_CACHE_ATTR, QuantizedKVCache(model.config, module))
         module.register_forward_pre_hook(_kv_cache_attention_hook, with_kwargs=True)
 
     kv_cache: QuantizedKVCache = getattr(module, KV_CACHE_ATTR)

src/compressed_tensors/transform/factory/base.py

@@ -108,9 +108,7 @@ def apply_to_model(self, model: Module, use_tqdm=True):
         for module, arg in tqdm.tqdm(modules_args, desc=desc, disable=(not use_tqdm)):
             self._apply_to_module(model, module, arg)
 
-    def _apply_to_module(
-        self, model: Module, module: Module, args: TransformArgs
-    ):
+    def _apply_to_module(self, model: Module, module: Module, args: TransformArgs):
         """
         Create transforms and apply them to the module