[quant][graphmode][fx] Add additional_qat_module_mapping

torch/quantization/fx/quantize.py

@@ -230,8 +230,11 @@ def __init__(self):
         self.patterns = None
 
 
-    def _qat_swap_modules(self, root):
-        convert(root, mapping=get_default_qat_module_mappings(), inplace=True, remove_qconfig=False)
+    def _qat_swap_modules(self, root, additional_qat_module_mapping):
+        all_mappings = get_default_qat_module_mappings().copy()
+        for k, v in additional_qat_module_mapping.items():
+            all_mappings[k] = v
+        convert(root, mapping=all_mappings, inplace=True, remove_qconfig=False)
 
     def _generate_qconfig_map(self,
                               root,
@@ -330,7 +333,8 @@ def _prepare(self, model, qconfig_dict, inplace, prepare_custom_config_dict, is_
         # TODO: support regex as well
         propagate_qconfig_(model, flattened_qconfig_dict)
         if model.training:
-            self._qat_swap_modules(model)
+            additional_qat_module_mapping = prepare_custom_config_dict.get("additioanl_qat_module_mapping", {})
+            self._qat_swap_modules(model, additional_qat_module_mapping)
 
         self.modules = dict(model.named_modules())
 

torch/quantization/quantize_fx.py

@@ -184,7 +184,12 @@ def prepare_fx(model, qconfig_dict, inplace=False, prepare_custom_config_dict=No
         # the module classes that are not symbolically traceable
         "non_traceable_module_class": [
            NonTraceableModule
-        ]
+        ],
+
+        # Additioanl module mapping for qat
+        "additional_qat_module_mapping": {
+           FloatModule: QATModule
+        },
       }