Approach changed

Signed-off-by: amitraj <quic_amitraj@quicinc.com>

Author: quic-amitraj

QEfficient/transformers/custom_attention.py

@@ -12,8 +12,6 @@
 from torch import nn
 from transformers.models.bert.modeling_bert import BertSelfAttention
 
-from QEfficient.utils.constants import BLOCK_SIZE
-
 
 class QEffBertSelfAttention(BertSelfAttention):
     def forward(
@@ -25,7 +23,7 @@ def forward(
         encoder_attention_mask: Optional[torch.FloatTensor] = None,
         past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
         output_attentions: Optional[bool] = False,
-        block_size: int = BLOCK_SIZE,
+        block_size: int = None,
     ) -> Tuple[torch.Tensor]:
         mixed_query_layer = self.query(hidden_states)
 

QEfficient/transformers/models/modeling_auto.py

@@ -57,10 +57,11 @@ def __repr__(self) -> str:
     def from_pretrained(cls, pretrained_model_name_or_path: str, is_tlm: bool = False, *args, **kwargs):
         if kwargs.get("attn_implementation", None) not in {None, "eager"}:
             logger.warning('Updating attn_implementation="eager"')
-            kwargs.update({"attn_implementation": "eager"})
+
         if kwargs.get("low_cpu_mem_usage", None):
             logger.warning("Updating low_cpu_mem_usage=False")
-            kwargs.update({"low_cpu_mem_usage": False})
+
+        kwargs.update({"attn_implementation": "eager", "low_cpu_mem_usage": False})
 
         model = cls._hf_auto_class.from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
         return cls(model, is_tlm=is_tlm)
@@ -430,20 +431,16 @@ class QEFFAutoModel(QEFFTransformersBase):
     _pytorch_transforms = [CustomOpsTransform, AwqToMatmulNbitsTransform, GPTQToMatmulNbitsTransform]
     _onnx_transforms = [FP16ClipTransform, SplitTensorsTransform]
 
-    def __init__(self, model: nn.Module, **kwargs):
-        if kwargs.get("block_size", None):
-            constants.BLOCK_SIZE = kwargs.get("block_size")
-            self._pytorch_transforms.append(BlockAttentionTransorm)
-            kwargs.update({"attn_implementation": "custom"})
-            kwargs.pop("block_size")
-
+    def __init__(self, model: nn.Module, block_size: Optional[int] = None, **kwargs):
+        if block_size:
+            BlockAttentionTransorm.apply(model, block_size=block_size)
         super().__init__(model)
         self.model.config.use_cache = True
         self.num_layers = model.config.num_hidden_layers
 
     @classmethod
     @with_replaced_quantizers
-    def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
+    def from_pretrained(cls, pretrained_model_name_or_path, block_size: Optional[int] = None, *args, **kwargs):
         """
         This method serves as the easiest entry point into using QEfficient. The interface is designed to be similar to transformers.AutoModel.
         Once the model is initialized, you can use other methods such as export, compile, and generate on the same object.
@@ -470,28 +467,20 @@ def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
             # You can now execute the model
             model.generate(inputs)
         """
-        if kwargs.get("block_size", None):
-            constants.BLOCK_SIZE = kwargs.get("block_size")
-            cls._pytorch_transforms.append(BlockAttentionTransorm)
-            kwargs.update({"attn_implementation": "custom"})
-            kwargs.pop("block_size")
-
-        if kwargs.get("attn_implementation", None) not in {None, "eager", "custom"}:
+        if kwargs.get("attn_implementation", None) not in {None, "eager"}:
             logger.warning('Updating attn_implementation="eager"')
-            kwargs.update({"attn_implementation": "eager"})
 
         if kwargs.get("low_cpu_mem_usage", None):
             logger.warning("Updating low_cpu_mem_usage=False")
-            kwargs.update({"low_cpu_mem_usage": False})
 
+        kwargs.update({"attn_implementation": "eager", "low_cpu_mem_usage": False, "add_pooling_layer": False})
         try:
-            kwargs.update({"add_pooling_layer": False})
             model = cls._hf_auto_class.from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
             warnings.warn("Removing pooling layer from the model if exist")
         except TypeError:
             kwargs.pop("add_pooling_layer", None)
             model = cls._hf_auto_class.from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
-        return cls(model)
+        return cls(model, block_size)
 
     @property
     def model_hash(self) -> str:

QEfficient/transformers/models/pytorch_transforms.py

@@ -5,6 +5,7 @@
 #
 # -----------------------------------------------------------------------------
 
+from functools import partial
 from types import MethodType
 from typing import Tuple
 
@@ -355,6 +356,13 @@ class BlockAttentionTransorm(ModuleMappingTransform):
     }
 
     @classmethod
-    def apply(cls, model: nn.Module) -> Tuple[nn.Module, bool]:
-        model, transformed = super().apply(model)
+    def apply(cls, model: nn.Module, block_size) -> Tuple[nn.Module, bool]:
+        transformed = False
+        for module in model.modules():
+            if repl_module := cls._module_mapping.get(type(module)):
+                module.__class__ = repl_module
+                # Bind the partial function to the instance
+                module.forward = MethodType(partial(repl_module.forward, block_size=block_size), module)
+                transformed = True
+                break
         return model, transformed

QEfficient/utils/_utils.py

@@ -8,6 +8,7 @@
 import json
 import os
 import subprocess
+from contextlib import contextmanager
 from typing import Any, Dict, List, Optional, Tuple, Union
 
 import requests
@@ -394,3 +395,13 @@ def create_json(file_path: str, json_data: object):
             json.dump(json_data, file, indent=4)
     except Exception as e:
         print(f"Failed to create JSON File {file_path}: {e}")
+
+
+@contextmanager
+def temporarily_remove_key(d, key):
+    value = d.pop(key, None)
+    try:
+        yield
+    finally:
+        if value is not None:
+            d[key] = value

QEfficient/utils/constants.py

@@ -49,7 +49,6 @@ def get_models_dir():
 ONNX_EXPORT_EXAMPLE_FBS = 4
 ONNX_EXPORT_EXAMPLE_NLK = 2  # Number of Logits to Keep
 ONNX_EXPORT_OPSET = 13
-BLOCK_SIZE = 32
 
 COMPILER = ["/opt/qti-aic/exec/qaic-exec", "-aic-hw", "-aic-hw-version=2.0"]
 

tests/peft/test_peft_onnx_transforms.py

@@ -24,7 +24,7 @@ def test_adapter_weights_to_inputs_transform():
     <
         float[32, 32] layer1_{adapter_name}_weight = [ "location" : "{external_tensors_file}" ],
         float[32, 32] layer2_{adapter_name}_weight = [ "location" : "{external_tensors_file}" ]
-    >
+    >f
     {{
         layer1output = MatMul (input, layer1_{adapter_name}_weight)
         output = MatMul (layer1output, layer2_{adapter_name}_weight)