|
| 1 | +# Copyright (c) Qualcomm Innovation Center, Inc. |
| 2 | +# All rights reserved |
| 3 | +# |
| 4 | +# This source code is licensed under the BSD-style license found in the |
| 5 | +# LICENSE file in the root directory of this source tree. |
| 6 | + |
| 7 | +"""Partition-level annotation for conv + batchnorm under QAT. |
| 8 | +
|
| 9 | +The per-node annotators in ``annotators/`` treat conv and batchnorm as unrelated |
| 10 | +ops. That is correct for PTQ, where batchnorm is folded into the conv before the |
| 11 | +graph is quantized, but it breaks QAT. |
| 12 | +
|
| 13 | +``prepare_qat_pt2e`` annotates the graph *before* torchao's ``_fuse_conv_bn_qat`` |
| 14 | +rewrites each conv+bn into the QAT folding pattern:: |
| 15 | +
|
| 16 | + scale = bn_weight / sqrt(bn_running_var + eps) |
| 17 | + y = conv(x, w * scale.reshape(w_shape)) / scale.reshape(b_shape) |
| 18 | + out = batch_norm(y, ...) |
| 19 | +
|
| 20 | +torchao does that ordering deliberately -- "Perform fusion after annotate to |
| 21 | +avoid quantizing ops in the new subgraph" -- but the rewrite copies each matched |
| 22 | +node's ``meta`` onto its replacement and re-points ``SharedQuantizationSpec`` |
| 23 | +chains at the new nodes, so a qspec on the *conv output* ends up on the scale |
| 24 | +arithmetic. ``convert_pt2e``'s ``_fold_conv_bn_qat`` matches a pattern carrying |
| 25 | +q/dq at exactly three points -- conv input, scaled weight, bn output -- so the |
| 26 | +extra q/dq makes the subgraph matcher miss and batchnorm is never folded. Float |
| 27 | +batchnorm then reaches the QNN partitioner, which has no node visitor for the |
| 28 | +training variant, drops it along with its FP16 scale arithmetic, and shatters |
| 29 | +the graph into one partition per conv. A partition left with only constant |
| 30 | +inputs fails context-binary serialization with the opaque |
| 31 | +``No graph inputs present for graph [0]`` / ``0x7532``. |
| 32 | +
|
| 33 | +Annotating the partition as a unit -- qspecs on the conv's inputs, the output |
| 34 | +qspec on the *last* node of the partition, nothing on the conv output -- gives |
| 35 | +the fold the shape it expects. This mirrors ``_do_annotate_conv_bn`` in the |
| 36 | +XNNPACK quantizer and the ``Conv*Bn*Quantizer`` family in BoltNN; QNN was the |
| 37 | +only one of the three without it. |
| 38 | +""" |
| 39 | + |
| 40 | +import operator |
| 41 | +from typing import Callable, List, Optional, Set |
| 42 | + |
| 43 | +import torch |
| 44 | +from torch.fx import GraphModule, Node |
| 45 | +from torchao.quantization.pt2e.quantizer import QuantizationAnnotation |
| 46 | + |
| 47 | +from .qconfig import QuantizationConfig |
| 48 | +from .rules import _is_annotated, _mark_nodes_as_annotated, Q_ANNOTATION_KEY |
| 49 | + |
| 50 | +CONV_TARGETS = ( |
| 51 | + torch.ops.aten.conv1d.default, |
| 52 | + torch.ops.aten.conv2d.default, |
| 53 | + torch.ops.aten.conv2d.padding, |
| 54 | + torch.ops.aten.convolution.default, |
| 55 | + torch.ops.aten.conv_transpose1d.default, |
| 56 | + torch.ops.aten.conv_transpose2d.input, |
| 57 | +) |
| 58 | + |
| 59 | +# ``batch_norm.default`` returns a bare tensor; every other variant returns a |
| 60 | +# tuple whose element 0 is the normalized output. |
| 61 | +SINGLE_OUTPUT_BN_TARGETS = (torch.ops.aten.batch_norm.default,) |
| 62 | + |
| 63 | +BN_TARGETS = SINGLE_OUTPUT_BN_TARGETS + ( |
| 64 | + torch.ops.aten._native_batch_norm_legit.default, |
| 65 | + torch.ops.aten._native_batch_norm_legit_functional.default, |
| 66 | + torch.ops.aten._native_batch_norm_legit_no_training.default, |
| 67 | + torch.ops.aten.cudnn_batch_norm.default, |
| 68 | +) |
| 69 | + |
| 70 | + |
| 71 | +def _sole_user(node: Node) -> Optional[Node]: |
| 72 | + return next(iter(node.users)) if len(node.users) == 1 else None |
| 73 | + |
| 74 | + |
| 75 | +def _bn_output(bn: Node) -> Optional[Node]: |
| 76 | + """The node carrying the batchnorm's normalized output.""" |
| 77 | + if bn.target in SINGLE_OUTPUT_BN_TARGETS: |
| 78 | + return bn |
| 79 | + for user in bn.users: |
| 80 | + if user.target is operator.getitem and user.args[1] == 0: |
| 81 | + return user |
| 82 | + return None |
| 83 | + |
| 84 | + |
| 85 | +def _materialize_conv_bias(gm: GraphModule, conv: Node) -> Optional[Node]: |
| 86 | + """Give a bias-less conv a zero bias so the fold produces a *quantized* one. |
| 87 | +
|
| 88 | + When the conv has no bias, ``fold_bn_weights_into_conv_node`` materializes the |
| 89 | + folded bias as a fresh fp32 parameter and explicitly leaves it unquantized |
| 90 | + ("NOTE: here we assume the bias of conv is not quantized!"). QNN rejects a |
| 91 | + per-channel-quantized conv carrying a float bias, so the conv drops out of the |
| 92 | + partition and its orphaned weight dequant has to become a standalone |
| 93 | + Dequantize -- which QNN does not support per-channel. |
| 94 | +
|
| 95 | + Adding a zero bias is a numerical no-op that gives the annotator a node to |
| 96 | + attach the derived bias qspec to. The fold then takes its |
| 97 | + ``conv_bias_node is not None`` branch and writes the folded values into that |
| 98 | + already-quantized node instead of inventing an unquantized one. |
| 99 | + """ |
| 100 | + args = list(conv.args) |
| 101 | + while len(args) < 3: |
| 102 | + args.append(None) |
| 103 | + if isinstance(args[2], Node): |
| 104 | + return args[2] |
| 105 | + |
| 106 | + val = conv.meta.get("val") |
| 107 | + if val is None: |
| 108 | + return None |
| 109 | + |
| 110 | + zeros = torch.zeros(val.shape[1], dtype=val.dtype, device=val.device) |
| 111 | + name = f"{conv.name}_zero_bias" |
| 112 | + gm.register_parameter(name, torch.nn.Parameter(zeros, requires_grad=False)) |
| 113 | + with gm.graph.inserting_before(conv): |
| 114 | + bias = gm.graph.get_attr(name) |
| 115 | + bias.meta["val"] = val.fake_mode.from_tensor(zeros, static_shapes=True) |
| 116 | + |
| 117 | + args[2] = bias |
| 118 | + conv.args = tuple(args) |
| 119 | + return bias |
| 120 | + |
| 121 | + |
| 122 | +def _annotate_partition( |
| 123 | + gm: GraphModule, |
| 124 | + conv: Node, |
| 125 | + partition: List[Node], |
| 126 | + output_node: Node, |
| 127 | + quantization_config: QuantizationConfig, |
| 128 | +) -> None: |
| 129 | + input_qspec_map = {} |
| 130 | + |
| 131 | + act = conv.args[0] |
| 132 | + if isinstance(act, Node) and quantization_config.input_activation is not None: |
| 133 | + input_qspec_map[act] = quantization_config.input_activation |
| 134 | + |
| 135 | + weight = conv.args[1] |
| 136 | + assert isinstance(weight, Node) |
| 137 | + input_qspec_map[weight] = quantization_config.weight |
| 138 | + |
| 139 | + bias = _materialize_conv_bias(gm, conv) |
| 140 | + if bias is not None and quantization_config.bias is not None: |
| 141 | + input_qspec_map[bias] = ( |
| 142 | + quantization_config.bias(conv) |
| 143 | + if callable(quantization_config.bias) |
| 144 | + else quantization_config.bias |
| 145 | + ) |
| 146 | + |
| 147 | + # Deliberately no output_qspec on the conv -- an observer there lands inside |
| 148 | + # the QAT fusion and stops `_fold_conv_bn_qat` from matching. |
| 149 | + conv.meta[Q_ANNOTATION_KEY] = QuantizationAnnotation( |
| 150 | + input_qspec_map=input_qspec_map, |
| 151 | + _annotated=True, |
| 152 | + ) |
| 153 | + _mark_nodes_as_annotated(partition) |
| 154 | + output_node.meta[Q_ANNOTATION_KEY] = QuantizationAnnotation( |
| 155 | + output_qspec=quantization_config.output_activation, |
| 156 | + _annotated=True, |
| 157 | + ) |
| 158 | + |
| 159 | + |
| 160 | +def annotate_conv_bn_partitions( |
| 161 | + gm: GraphModule, |
| 162 | + get_quant_config: Callable[[Node], Optional[QuantizationConfig]], |
| 163 | + discard_nodes: Set[str], |
| 164 | +) -> int: |
| 165 | + """Annotate every conv+bn chain as a single partition. |
| 166 | +
|
| 167 | + Must run before the per-node annotation pass; the nodes it claims are marked |
| 168 | + annotated so the per-node annotators skip them. |
| 169 | +
|
| 170 | + A trailing relu is deliberately NOT absorbed into the partition. XNNPACK's |
| 171 | + equivalent puts the output qspec on the relu because it fuses conv+relu into |
| 172 | + one backend op, but QNN lowers conv and relu to separate ops and rejects a |
| 173 | + conv whose output is unquantized (`0x232 != 0x408`, FLOAT_32 meeting |
| 174 | + UFIXED_POINT_8). Keeping the qspec on the bn output -- which the fold |
| 175 | + transfers onto the conv output -- leaves the relu to the per-node annotator, |
| 176 | + which gives it its own quantized input and output. |
| 177 | +
|
| 178 | + Returns the number of partitions annotated. |
| 179 | + """ |
| 180 | + count = 0 |
| 181 | + for conv in list(gm.graph.nodes): |
| 182 | + if conv.op != "call_function" or conv.target not in CONV_TARGETS: |
| 183 | + continue |
| 184 | + if conv.name in discard_nodes: |
| 185 | + continue |
| 186 | + |
| 187 | + bn = _sole_user(conv) |
| 188 | + if bn is None or bn.target not in BN_TARGETS: |
| 189 | + continue |
| 190 | + bn_out = _bn_output(bn) |
| 191 | + if bn_out is None: |
| 192 | + continue |
| 193 | + |
| 194 | + partition = [conv, bn] if bn_out is bn else [conv, bn, bn_out] |
| 195 | + output_node = bn_out |
| 196 | + |
| 197 | + if _is_annotated(partition): |
| 198 | + continue |
| 199 | + quantization_config = get_quant_config(conv) |
| 200 | + if quantization_config is None: |
| 201 | + continue |
| 202 | + |
| 203 | + _annotate_partition(gm, conv, partition, output_node, quantization_config) |
| 204 | + count += 1 |
| 205 | + |
| 206 | + if count: |
| 207 | + gm.recompile() |
| 208 | + return count |
0 commit comments