update custom annotations for QAT (#13747)

Summary:

Some custom annotations were PTQ only, updating them to support QAT, as it is needed for QAT of static transformer

Reviewed By: sxu

Differential Revision: D80190466

Author: navsud

backends/qualcomm/quantizer/custom_annotation.py

@@ -92,7 +92,9 @@ def annotate_mimi_decoder(gm: torch.fx.GraphModule):
             break
 
 
-def annotate_linear_16a8w_in_affine_layer(gm: torch.fx.GraphModule) -> None:
+def annotate_linear_16a8w_in_affine_layer(
+    gm: torch.fx.GraphModule, is_qat: bool = False
+) -> None:
     def annotate_conv2d(node: Node, quantization_config: QuantizationConfig) -> None:
         input_qspec_map = {}
         input_act = node.args[0]
@@ -108,9 +110,14 @@ def annotate_conv2d(node: Node, quantization_config: QuantizationConfig) -> None
             _annotated=True,
         )
 
-    quantization_config_16a8w_per_channel = get_ptq_per_channel_quant_config(
-        torch.uint16, weight_dtype=torch.int8, act_observer=MinMaxObserver
-    )
+    if is_qat:
+        quantization_config_16a8w_per_channel = get_qat_per_channel_quant_config(
+            torch.uint16, weight_dtype=torch.int8, act_observer=MinMaxObserver
+        )
+    else:
+        quantization_config_16a8w_per_channel = get_ptq_per_channel_quant_config(
+            torch.uint16, weight_dtype=torch.int8, act_observer=MinMaxObserver
+        )
     for node in gm.graph.nodes:
         if node.op == "call_function" and node.target == torch.ops.aten.conv2d.default:
             if "nn_module_stack" in node.meta: