up

Author: metascroy

.lintrunner.toml

@@ -397,7 +397,6 @@ exclude_patterns = [
     "extension/llm/export/builder.py",
     "extension/llm/export/quantizer_lib.py",
     "exir/tests/test_memory_planning.py",
-    "backends/transforms/duplicate_dynamic_quant_chain.py",
     "exir/backend/test/demos/test_xnnpack_qnnpack.py",
 ]
 

backends/transforms/duplicate_dynamic_quant_chain.py

@@ -9,14 +9,12 @@
 
 import torch
 
-from torch.ao.quantization.pt2e.utils import (
-    _filter_sym_size_users,
-    _is_valid_annotation,
-)
-
 from torch.fx.node import map_arg
 from torch.fx.passes.infra.pass_base import PassBase, PassResult
 
+from torchao.quantization.pt2e.quantizer import is_valid_annotation
+from torchao.quantization.pt2e.utils import _filter_sym_size_users
+
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.WARNING)
@@ -129,7 +127,7 @@ def _maybe_duplicate_dynamic_quantize_chain(
     dq_node_users = list(dq_node.users.copy())
     for user in dq_node_users:
         annotation = user.meta.get("quantization_annotation", None)
-        if not _is_valid_annotation(annotation):
+        if not is_valid_annotation(annotation):
             return
         with gm.graph.inserting_after(dq_node):
             new_node = gm.graph.node_copy(dq_node)

extension/llm/export/builder.py

@@ -13,7 +13,7 @@
 import contextlib
 import logging
 from enum import Enum
-from typing import Any, Callable, Dict, List, Optional, Tuple
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 from unittest.mock import patch
 
 import torch
@@ -35,11 +35,15 @@
 
 from executorch.extension.llm.export.export_passes import RemoveRedundantTransposes
 from pytorch_tokenizers import get_tokenizer
-from torch.ao.quantization.quantizer import Quantizer
-from torch.ao.quantization.quantizer.composable_quantizer import ComposableQuantizer
+from torch.ao.quantization.quantizer import TorchQuantizer
+from torch.ao.quantization.quantizer.composable_quantizer import (
+    TorchComposableQuantizer,
+)
+
 from torch.export import export_for_training, ExportedProgram
 from torch.nn.attention import SDPBackend
 from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e
+from torchao.quantization.pt2e.quantizer import ComposableQuantizer, Quantizer
 from torchao.utils import unwrap_tensor_subclass
 
 FORMAT = "[%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s"
@@ -350,7 +354,9 @@ def calibrate_template(
             print(f"{task}: {res}")
         logging.info("Calibration finish...")
 
-    def pt2e_quantize(self, quantizers: Optional[List[Quantizer]]) -> "LLMEdgeManager":
+    def pt2e_quantize(
+        self, quantizers: Optional[List[Union[Quantizer, TorchQuantizer]]]
+    ) -> "LLMEdgeManager":
         """
         Quantize the model via pt2e flow and retrieve LLMEdgeManager including the quantized model.
         Args:
@@ -367,7 +373,12 @@ def pt2e_quantize(self, quantizers: Optional[List[Quantizer]]) -> "LLMEdgeManage
             with torch.nn.attention.sdpa_kernel([SDPBackend.MATH]), torch.no_grad():
                 if self.verbose:
                     logging.info(f"Applied quantizers: {quantizers}")
-                composed_quantizer = ComposableQuantizer(quantizers)
+
+                if any(isinstance(q, Quantizer) for q in quantizers):
+                    composed_quantizer = ComposableQuantizer(quantizers)
+                else:
+                    composed_quantizer = TorchComposableQuantizer(quantizers)
+
                 assert (
                     self.pre_autograd_graph_module is not None
                 ), "Please run export() first"

extension/llm/export/quantizer_lib.py

@@ -154,7 +154,7 @@ def get_qnn_quantizer(
             QnnQuantizer,
             QuantDtype,
         )
-        from torch.ao.quantization.observer import MinMaxObserver
+        from torchao.quantization.pt2e import MinMaxObserver
 
     except ImportError:
         raise ImportError(