Add torchao int4 weight only quantization as an option (#34)

- Add torchao int4 weight only quantization as an option (3795233d4c81c81bca28df592f4c303f1dbaeee6)
- fix(quant): add torchao to quantization option (371c49c8bf1df0569820dcd4aee19441247c159d)


Co-authored-by: Jerry Zhang <[email protected]>

hardware.yaml

@@ -19,6 +19,7 @@
     - awq
     - bnb
     - gptq
+    - torchao
   backends:
     - pytorch
 
@@ -31,6 +32,7 @@
     - awq
     - bnb
     - gptq
+    - torchao
   backends:
     - pytorch
 
@@ -45,4 +47,4 @@
   backends:
     - pytorch
     - openvino
-    - onnxruntime
+    - onnxruntime

src/kernels.py

@@ -38,6 +38,7 @@ def get_quant_df(llm_perf_df):
     exllamav2_df = copy_df[(copy_df["Quantization 🗜️"] == "GPTQ.4bit+ExllamaV2")]
     gemm_df = copy_df[(copy_df["Quantization 🗜️"] == "AWQ.4bit+GEMM")]
     gemv_df = copy_df[(copy_df["Quantization 🗜️"] == "AWQ.4bit+GEMV")]
+    torchao_df = copy_df[(copy_df["Quantization 🗜️"] == "torchao.4bit")]
     # merge the three dataframes
     exllamav1_df = pd.merge(
         vanilla_df,
@@ -63,8 +64,14 @@ def get_quant_df(llm_perf_df):
         on=["Model 🤗"],
         suffixes=["", " Custom Kernel"],
     )
+    torchao_df = pd.merge(
+        vanilla_df,
+        torchao_df,
+        on=["Model 🤗"],
+        suffixes=["", " Custom Kernel"],
+    )
     # concat the two dataframes row-wise
-    quant_df = pd.concat([exllamav1_df, exllamav2_df, gemm_df, gemv_df])
+    quant_df = pd.concat([exllamav1_df, exllamav2_df, gemm_df, gemv_df, torchao_df])
     # compute speedups
     quant_df["Prefill Speedup (%)"] = (
         (quant_df["Prefill (s)"] / quant_df["Prefill (s) Custom Kernel"]) * 100

src/panel.py

@@ -26,7 +26,7 @@ def create_control_panel(
     if hardware_provider == "nvidia":
         backends = ["pytorch"]
         attention_implementations = ["Eager", "SDPA", "FAv2"]
-        quantizations = ["Unquantized", "BnB.4bit", "BnB.8bit", "AWQ.4bit", "GPTQ.4bit"]
+        quantizations = ["Unquantized", "BnB.4bit", "BnB.8bit", "AWQ.4bit", "GPTQ.4bit", "torchao.4bit"]
         kernels = [
             "No Kernel",
             "GPTQ.ExllamaV1",

src/utils.py

@@ -70,6 +70,11 @@ def process_quantizations(x):
         and x["config.backend.quantization_config.bits"] == 4
     ):
         return "AWQ.4bit"
+    elif (
+            x["config.backend.quantization_scheme"] == "torchao"
+            and x["config.backend.quantization_config.quant_type"] == "int4_weight_only"
+    ):
+        return "torchao.4bit"
     else:
         return "Unquantized"