Update on "[Quantization] Add "quantization_tag" as metadata to fx proxy"

kimishpatel · kimishpatel · commit a2c8f01836d4 · 2023-11-01T08:45:35.000-07:00
Summary: In order to make sure that quantization_tag is preserved through second stage export, this PR adds it as a special metadata that should be preserved. Since quantization in export path will work on top of pre dispatch graph, subsequent post dispatch op decomposition, will decompose ops that quant workflow tagged. In order to make sure that the patterns identified by quantizer, remains identifiable, even after decompositions are applied, we must preserve "quantization_tag". This enables backend delegates, that quantized a model for specific backend, to be able to identify "quantized" patterns. Test Plan: metadata porting tests Reviewers: Subscribers: Tasks: Tags: Differential Revision: [D49056259](https://our.internmc.facebook.com/intern/diff/D49056259) [ghstack-poisoned]
diff --git a/test/quantization/pt2e/test_metadata_porting.py b/test/quantization/pt2e/test_metadata_porting.py
@@ -437,7 +437,7 @@ def annotate(self, gm: torch.fx.GraphModule) -> torch.fx.GraphModule:
                     is_per_channel=True
                 )
                 OP_TO_ANNOTATOR["linear"](gm, quantization_config)
-                OP_TO_ANNOTATOR["conv2d"](gm, quantization_config)
+                OP_TO_ANNOTATOR["conv"](gm, quantization_config)
                 OP_TO_ANNOTATOR["adaptive_avg_pool2d"](gm, quantization_config)
 
             def validate(self, model: torch.fx.GraphModule) -> None:
diff --git a/torch/fx/proxy.py b/torch/fx/proxy.py
@@ -84,7 +84,7 @@ def __exit__(self, *args):
         return
 
 
-_COPY_META_FIELDS = ["nn_module_stack", "source_fn_stack", "original_aten", "recompute", "from_node", "quantizatcon_tag"]
+_COPY_META_FIELDS = ["nn_module_stack", "source_fn_stack", "original_aten", "recompute", "from_node", "quantization_tag"]
 
 
 @compatibility(is_backward_compatible=True)
