Fix lint

- Fix transformers version
- Refine pass quantization tagging function
- Rebase

Author: Joey Tsai

backends/qualcomm/quantizer/custom_annotation.py

@@ -22,7 +22,9 @@
 from torch.fx import Node
 
 
-def annotate_matmul_16a8w(gm: torch.fx.GraphModule) -> None:  # noqa: C901
+def annotate_matmul_16a8w(
+    gm: torch.fx.GraphModule, traverse_input1=True
+) -> None:  # noqa: C901
     """
     This function is specific for matmul op 16a8w.
     """
@@ -99,7 +101,8 @@ def annotate_matmul_input1(node: Node):
     for node in gm.graph.nodes:
         if node.op == "call_function" and node.target == torch.ops.aten.matmul.default:
             annotate_matmul(node, quantization_config_16a8w)
-            annotate_matmul_input1(node.args[1])
+            if traverse_input1:
+                annotate_matmul_input1(node.args[1])
 
 
 def custom_annotate_llama_matmul_16a8w(gm: torch.fx.GraphModule) -> None:  # noqa: C901

examples/qualcomm/oss_scripts/llama3_2/llama.py

@@ -8,9 +8,9 @@
 import json
 import logging
 import os
-
 import sys
 import time
+from functools import partial
 from multiprocessing.connection import Client
 
 import torch
@@ -319,8 +319,10 @@ def compile(args):
 
     if args.model_mode == "kv":
         use_kv_cache = output_new_cache_only = True
+        matmul_annotate_func = partial(annotate_matmul_16a8w, traverse_input1=True)
     elif args.model_mode == "batch_prefill":
         use_kv_cache = output_new_cache_only = False
+        matmul_annotate_func = partial(annotate_matmul_16a8w, traverse_input1=False)
     elif args.model_mode == "hybrid":
         raise NotImplementedError(
             f"model_mode {args.model_mode} is not implemented yet."
@@ -385,7 +387,10 @@ def compile(args):
         start_quantize_ts = time.time()
         single_llama.quantize(
             quant_dtype,
-            custom_annotations=(annotate_matmul_16a8w,),
+            custom_annotations=(
+                custom_annotate_llama_last_conv_16a8w,
+                matmul_annotate_func,
+            ),
         )
         end_quantize_ts = time.time()
         logging.info(f"Time for quantizing: {end_quantize_ts - start_quantize_ts}")

install_requirements.py

@@ -137,7 +137,7 @@ def python_is_compatible():
     "timm==1.0.7",
     f"torchaudio==2.5.0.{NIGHTLY_VERSION}" if USE_PYTORCH_NIGHTLY else "torchaudio",
     "torchsr==1.0.4",
-    "transformers==4.42.4", # TODO update back to 4.46.1 once the error is fixed
+    "transformers==4.46.1",
 ]
 
 # pip packages needed for development.