Let quantizers add custom meta data to GraphModule

test/quantization/pt2e/test_quantize_pt2e.py

@@ -1578,6 +1578,38 @@ def forward(self, x):
         }
         self.checkGraphModuleNodes(m, expected_node_occurrence=node_occurrence)
 
+    def test_custom_meta_transform_for_annotation(self):
+        class TestQuantizer(Quantizer):
+            def transform_for_annotation(
+                self, m: torch.fx.GraphModule
+            ) -> torch.fx.GraphModule:
+                # Make a copy of the graph to ensure that we are using the
+                # return value of this function.
+                graph = torch.fx.Graph()
+                graph.graph_copy(m.graph, {})
+                model = torch.fx.GraphModule(m, graph)
+                model.meta["custom"] = {"_test_data": True}
+                return model
+
+            def annotate(self, model: torch.fx.GraphModule) -> torch.fx.GraphModule:
+                return model
+
+            def validate(self, model: torch.fx.GraphModule) -> None:
+                pass
+
+        class M(torch.nn.Module):
+            def forward(self, x):
+                return x + 3
+
+        m = M().eval()
+        quantizer = TestQuantizer()
+        example_inputs = (torch.randn(1, 2, 3, 3),)
+        m = torch.export.export(m, example_inputs, strict=True).module()
+        prepared = prepare_pt2e(m, quantizer)
+
+        custom_meta = prepared.meta.get("custom", {}).get("_test_data", False)
+        self.assertTrue(custom_meta)
+
     def test_composable_quantizer_transform_for_annotation(self):
         class TestQuantizer1(Quantizer):
             def transform_for_annotation(

torchao/quantization/pt2e/quantize_pt2e.py

@@ -116,13 +116,17 @@ def calibrate(model, data_loader):
     model = quantizer.transform_for_annotation(model)
     quantizer.annotate(model)
     quantizer.validate(model)
+    # Store the 'custom' meta data if any added by quantizer
+    annotated_meta = model.meta.get("custom", {})
     model = prepare(
         model,
         node_name_to_scope,
         is_qat=False,
         obs_or_fq_callback=quantizer.prepare_obs_or_fq_callback,
     )
     model.meta.update(original_graph_meta)
+    # Update the 'custom' meta data from quantizer
+    model.meta.update({"custom": annotated_meta})
     model = _disallow_eval_train(model)
     return model