Update callsite for pt2e quant

Summary:
We just removed pt2e quant from pytorch/pytorch in pytorch/pytorch#169151
This PR updated the pt2e quantization callsites to torchao

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:

Author: jerryzh168

coremltools/converters/mil/frontend/torch/test/test_torch_export_quantization.py

@@ -29,10 +29,10 @@
     )
 
 from torch.export import export_for_training
-from torch.ao.quantization.quantize_pt2e import convert_pt2e, prepare_pt2e, prepare_qat_pt2e
-from torch.ao.quantization.quantizer.xnnpack_quantizer import (
-    XNNPACKQuantizer,
-    get_symmetric_quantization_config,
+from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e, prepare_qat_pt2e
+from executorch.backends.xnnpack.quantizer.xnnpack_quantizer import (
+   get_symmetric_quantization_config,
+   XNNPACKQuantizer,
 )
 
 import coremltools as ct

coremltools/optimize/torch/quantization/_annotation_config.py

@@ -8,7 +8,7 @@
 import torch as _torch
 import torch.ao.quantization as _aoquant
 from attr import define as _define
-from torch.ao.quantization.quantizer.quantizer import (
+from torchao.quantization.pt2e.quantizer import (
     QuantizationSpec as _TorchQuantizationSpec,
 )
 

coremltools/optimize/torch/quantization/_coreml_quantizer.py

@@ -9,8 +9,8 @@
 from typing import Optional as _Optional
 
 import torch as _torch
-from torch.ao.quantization.quantizer.quantizer import Quantizer as _TorchQuantizer
-from torch.ao.quantization.quantizer.xnnpack_quantizer import _get_module_name_filter
+from torchao.quantization.pt2e.quantizer import Quantizer as _TorchQuantizer
+from torchao.quantization.pt2e.quantizer import _get_module_name_filter
 from torch.fx import Node as _Node
 
 import coremltools.optimize.torch.quantization._coreml_quantizer_utils as _annotation_utils
@@ -519,7 +519,7 @@ class CoreMLQuantizer(_TorchQuantizer):
     """
     Annotates all recognized patterns using ``config``.
 
-    Extends py:class:`torch.ao.quantization.quantizer.quantizer.Quantizer` to
+    Extends py:class:`torchao.quantization.pt2e.quantizer.Quantizer` to
     add support for quantization patterns supported by Core ML.
 
     Use it in conjunction with PyTorch 2.0 ``prepare_pt2e`` and ``prepare_qat_pt2e`` APIs
@@ -532,7 +532,7 @@ class CoreMLQuantizer(_TorchQuantizer):
 
                 import torch.nn as nn
                 from torch.export import export_for_training
-                from torch.ao.quantization.quantize_pt2e import convert_pt2e, prepare_qat_pt2e
+                from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_qat_pt2e
 
                 from coremltools.optimize.torch.quantization._coreml_quantizer import CoreMLQuantizer
 

coremltools/optimize/torch/quantization/_coreml_quantizer_utils.py

@@ -17,26 +17,22 @@
 if _IS_TORCH_OLDER_THAN_2_4:
     from torch.ao.quantization.pt2e.utils import get_aten_graph_module
 else:
-    from torch.ao.quantization.pt2e.utils import _get_aten_graph_module_for_pattern
+    from torchao.quantization.pt2e.utils import _get_aten_graph_module_for_pattern
 
-from torch.ao.quantization.quantizer.quantizer import (
+from torchao.quantization.pt2e.quantizer import (
     FixedQParamsQuantizationSpec as _FixedQParamsQuantizationSpec,
 )
-from torch.ao.quantization.quantizer.quantizer import (
+from torchao.quantization.pt2e.quantizer import (
     QuantizationAnnotation as _QuantizationAnnotation,
 )
-from torch.ao.quantization.quantizer.quantizer import QuantizationSpec as _TorchQuantizationSpec
-from torch.ao.quantization.quantizer.quantizer import (
+from torchao.quantization.pt2e.quantizer import QuantizationSpec as _TorchQuantizationSpec
+from torchao.quantization.pt2e.quantizer import (
     QuantizationSpecBase as _TorchQuantizationSpecBase,
 )
-from torch.ao.quantization.quantizer.quantizer import (
+from torchao.quantization.pt2e.quantizer import (
     SharedQuantizationSpec as _SharedQuantizationSpec,
 )
-from torch.ao.quantization.quantizer.xnnpack_quantizer import _get_module_name_filter
-from torch.ao.quantization.quantizer.xnnpack_quantizer_utils import (
-    _is_annotated,
-    _mark_nodes_as_annotated,
-)
+from torchao.quantization.pt2e.quantizer import get_module_name_filter
 from torch.fx import Node as _Node
 from torch.fx.passes.utils.matcher_with_name_node_map_utils import (
     SubgraphMatcherWithNameNodeMap as _SubgraphMatcherWithNameNodeMap,
@@ -49,6 +45,27 @@
     AnnotationConfig as _AnnotationConfig,
 )
 
+def _is_annotated(nodes: list[_Node]):
+    """
+    Given a list of nodes (that represents an operator pattern),
+    check if any of the node is annotated, return True if any of the node
+    is annotated, otherwise return False
+    """
+    for node in nodes:
+        if annotated or (
+            "quantization_annotation" in node.meta
+            and node.meta["quantization_annotation"]._annotated
+        ):
+            return True
+    return False
+
+def _mark_nodes_as_annotated(nodes: list[_Node]):
+    for node in nodes:
+        if node is not None:
+            if "quantization_annotation" not in node.meta:
+                node.meta["quantization_annotation"] = _QuantizationAnnotation()
+            node.meta["quantization_annotation"]._annotated = True
+
 # All activations recognized for conv-act/conv-bn-act patterns
 _supported_activations = (
     _F.relu,
@@ -185,7 +202,7 @@ def get_not_object_type_or_name_filter(
     type in ``tp_list``.
     """
     object_type_filters = [get_object_type_filter(tp) for tp in tp_list]
-    module_name_list_filters = [_get_module_name_filter(m) for m in module_name_list]
+    module_name_list_filters = [get_module_name_filter(m) for m in module_name_list]
 
     def not_object_type_or_name_filter(n: _Node) -> bool:
         return not any(f(n) for f in object_type_filters + module_name_list_filters)

coremltools/test/optimize/torch/quantization/test_coreml_quantizer.py

@@ -20,7 +20,7 @@
 from coremltools._deps import _HAS_TORCH_EXPORT_API
 if _HAS_TORCH_EXPORT_API:
     from torch.export import export_for_training
-    from torch.ao.quantization.quantize_pt2e import (
+    from torchao.quantization.pt2e.quantize_pt2e import (
         convert_pt2e,
         prepare_pt2e,
         prepare_qat_pt2e,

reqs/pytorch.pip

@@ -13,4 +13,4 @@ torchsr==1.0.4; platform_machine == "arm64"
 # TODO (rdar://141476729) support a more recent timm
 timm==0.6.13; platform_machine == "arm64"
 
-torchao==0.10.0; platform_machine == "arm64" and python_version >= '3.10'
+torchao==0.15.0; platform_machine == "arm64" and python_version >= '3.10'