Merge pull request #401 from puneetmatharu/patch-quantization-in-answer-questions

Patch quantization in PyTorch examples

Author: jondea

ML-Frameworks/pytorch-aarch64/examples/answer_questions.py

@@ -19,6 +19,16 @@
 import random
 import torch
 from transformers import AutoTokenizer, AutoModelForQuestionAnswering
+from torchao.quantization.quant_api import (
+    Int8DynamicActivationIntxWeightConfig,
+    quantize_,
+)
+from torchao.dtypes.uintx.packed_linear_int8_dynamic_activation_intx_weight_layout import (
+    PackedLinearInt8DynamicActivationIntxWeightLayout,
+    Target,
+)
+from torchao.quantization.granularity import PerAxis
+from torchao.quantization.quant_primitives import MappingType
 
 from utils import nlp
 
@@ -43,7 +53,7 @@ def main():
     parser.add_argument("--bert-large", action='store_true',
         help="Use BERT large instead of DistilBERT")
     parser.add_argument("--quantize", action='store_true',
-        help="Quantize the model to int8 using dynamic quantization")
+        help="Quantize the model to int4 using dynamic quantization")
     parser.add_argument("--warmup", action='store_true',
         help="Run warmup")
 
@@ -127,10 +137,20 @@ def main():
     model = AutoModelForQuestionAnswering.from_pretrained(model_hf_path)
 
     if args["quantize"]:
-        model = torch.ao.quantization.quantize_dynamic(
+        layout = PackedLinearInt8DynamicActivationIntxWeightLayout(target=Target.ATEN)
+        quantize_(
             model,
-            {torch.nn.Linear},
-            dtype=torch.qint8)
+            Int8DynamicActivationIntxWeightConfig(
+                weight_scale_dtype=torch.float32,
+                weight_granularity=PerAxis(0),
+                weight_mapping_type=MappingType.SYMMETRIC_NO_CLIPPING_ERR,
+                layout=layout,
+                weight_dtype=torch.int4,
+                intx_packing_format="opaque_aten_kleidiai",
+                version=2,
+            ),
+            filter_fn=lambda m, _: isinstance(m, torch.nn.Linear),
+        )
 
     encoding = token.encode_plus(
         question,

ML-Frameworks/pytorch-aarch64/examples/llama_vision_instruct.py

@@ -29,7 +29,7 @@
     PackedLinearInt8DynamicActivationIntxWeightLayout,
     Target,
 )
-from torchao.quantization.granularity import PerGroup, PerAxis
+from torchao.quantization.granularity import PerAxis
 from torchao.quantization.quant_primitives import MappingType
 import numpy as np
 import os
@@ -53,7 +53,7 @@ def main(args):
                 layout=layout,
                 weight_dtype=torch.int4,
                 intx_packing_format="opaque_aten_kleidiai",
-                version=1,
+                version=2,
             ),
         )
 

ML-Frameworks/pytorch-aarch64/examples/quantized_linear.py

@@ -1,5 +1,5 @@
 # *******************************************************************************
-# Copyright 2024 Arm Limited and affiliates.
+# Copyright 2024-2025 Arm Limited and affiliates.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -16,10 +16,19 @@
 # *******************************************************************************
 
 import sys
-import os
 
 import torch
 import torch.nn as nn
+from torchao.quantization.quant_api import (
+    Int8DynamicActivationIntxWeightConfig,
+    quantize_,
+)
+from torchao.dtypes.uintx.packed_linear_int8_dynamic_activation_intx_weight_layout import (
+    PackedLinearInt8DynamicActivationIntxWeightLayout,
+    Target,
+)
+from torchao.quantization.granularity import PerAxis
+from torchao.quantization.quant_primitives import MappingType
 
 import time
 
@@ -53,10 +62,19 @@ def forward(self, x):
         model(data)
         fp32_runtimes.append(time.time() - t0)
 
-model = torch.ao.quantization.quantize_dynamic(
+quantize_(
     model,
-    {torch.nn.Linear},
-    dtype=torch.qint8)
+    Int8DynamicActivationIntxWeightConfig(
+        weight_scale_dtype=torch.float32,
+        weight_granularity=PerAxis(0),
+        weight_mapping_type=MappingType.SYMMETRIC_NO_CLIPPING_ERR,
+        layout=PackedLinearInt8DynamicActivationIntxWeightLayout(target=Target.ATEN),
+        weight_dtype=torch.int4,
+        intx_packing_format="opaque_aten_kleidiai",
+        version=2,
+    ),
+    filter_fn=lambda m, _: isinstance(m, torch.nn.Linear),
+)
 
 # Quantized
 runtimes = []