|
| 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 | +import types |
| 7 | + |
| 8 | +import torch |
| 9 | +import torchao |
| 10 | +from torchao.quantization.pt2e import PerChannelMinMaxObserver |
| 11 | +from executorch.backends.qualcomm.quantizer.observers.per_block_param_observer import ( |
| 12 | + PerBlockParamObserver, |
| 13 | +) |
| 14 | +from executorch.exir.pass_base import ExportPass, PassResult |
| 15 | + |
| 16 | +class SeqMseModule(torch.nn.Module): |
| 17 | + """ |
| 18 | + Args: |
| 19 | + nominal_weight: Tensor |
| 20 | + nominal parameters from operator |
| 21 | + nominal_bias: Tensor |
| 22 | + nominal parameters from operator |
| 23 | + operator: fx.Node |
| 24 | + operator to be executed |
| 25 | + observer: UniformQuantizationObserverBase |
| 26 | + parameter observer (specific for weight) |
| 27 | + num_candidates: int |
| 28 | + grids to search minimal mse loss |
| 29 | + """ |
| 30 | + |
| 31 | + def __init__( |
| 32 | + self, |
| 33 | + nominal_weight, |
| 34 | + nominal_bias, |
| 35 | + operator, |
| 36 | + observer, |
| 37 | + num_candidates, |
| 38 | + ): |
| 39 | + super().__init__() |
| 40 | + self.nominal_weight = nominal_weight |
| 41 | + self.nominal_bias = nominal_bias |
| 42 | + self.observer = observer |
| 43 | + step = 1 / num_candidates |
| 44 | + self.steps = torch.arange(start=step, end=1+step, step=step).tolist() |
| 45 | + self.operator = self._make_operator(operator) |
| 46 | + self.best_candidate = 1.0 |
| 47 | + |
| 48 | + def _make_operator(self, aten_op): |
| 49 | + if aten_op.target == torch.ops.aten.conv2d.default: |
| 50 | + stride = [1, 1] if len(aten_op.args) < 4 else aten_op.args[3] |
| 51 | + padding = [0, 0] if len(aten_op.args) < 5 else aten_op.args[4] |
| 52 | + dilation = [1, 1] if len(aten_op.args) < 6 else aten_op.args[5] |
| 53 | + groups = 1 if len(aten_op.args) < 7 else aten_op.args[6] |
| 54 | + has_bias = self.nominal_bias is not None |
| 55 | + module = torch.nn.Conv2d( |
| 56 | + in_channels=self.nominal_weight.shape[1], |
| 57 | + out_channels=self.nominal_weight.shape[0], |
| 58 | + kernel_size=self.nominal_weight.shape[-2:], |
| 59 | + stride=stride, |
| 60 | + padding=padding, |
| 61 | + dilation=dilation, |
| 62 | + groups=groups, |
| 63 | + bias=has_bias |
| 64 | + ) |
| 65 | + module.weight.data = self.nominal_weight |
| 66 | + if has_bias: |
| 67 | + module.bias.data = self.nominal_bias |
| 68 | + return module |
| 69 | + else: |
| 70 | + raise NotImplementedError(f"target of {aten_op.target} is not implemented") |
| 71 | + |
| 72 | + def _per_block_qdq(self, scale, zero_point): |
| 73 | + return torchao.quantization.quant_primitives._fake_quantize_affine( |
| 74 | + input=self.nominal_weight, |
| 75 | + block_size=self.observer.block_size, |
| 76 | + scale=scale, |
| 77 | + zero_point=zero_point, |
| 78 | + quant_dtype=self.observer.dtype, |
| 79 | + quant_min=self.observer.quant_min, |
| 80 | + quant_max=self.observer.quant_max, |
| 81 | + ) |
| 82 | + |
| 83 | + def _per_channel_qdq(self, scale, zero_point): |
| 84 | + return torch.fake_quantize_per_channel_affine( |
| 85 | + input=self.nominal_weight, |
| 86 | + scale=scale, |
| 87 | + zero_point=zero_point, |
| 88 | + axis=0, |
| 89 | + quant_min=self.observer.quant_min, |
| 90 | + quant_max=self.observer.quant_max, |
| 91 | + ) |
| 92 | + |
| 93 | + def _fake_quant(self, scale, zero_point): |
| 94 | + dispatcher = { |
| 95 | + PerChannelMinMaxObserver: self._per_channel_qdq, |
| 96 | + PerBlockParamObserver: self._per_block_qdq, |
| 97 | + } |
| 98 | + return dispatcher[type(self.observer)](scale, zero_point) |
| 99 | + |
| 100 | + def _find_best_candidate(self, nominal_input, nominal_output): |
| 101 | + # calculate current baseline |
| 102 | + scale, zero_point = self.observer.calculate_qparams() |
| 103 | + zero_point = zero_point.to(torch.int32) |
| 104 | + self.operator.weight.data = self._fake_quant(scale, zero_point) |
| 105 | + candidate, current_loss = 1, torch.nn.functional.mse_loss( |
| 106 | + self.operator(nominal_input), nominal_output |
| 107 | + ).item() |
| 108 | + for step in self.steps: |
| 109 | + self.operator.weight.data = self._fake_quant(scale * step, zero_point) |
| 110 | + loss = torch.nn.functional.mse_loss( |
| 111 | + self.operator(nominal_input), nominal_output |
| 112 | + ).item() |
| 113 | + if loss < current_loss: |
| 114 | + candidate, current_loss = step, loss |
| 115 | + return candidate |
| 116 | + |
| 117 | + def forward(self, nominal_input, nominal_output): |
| 118 | + self.best_candidate = self._find_best_candidate( |
| 119 | + nominal_input=nominal_input, nominal_output=nominal_output |
| 120 | + ) |
| 121 | + |
| 122 | + |
| 123 | +class InsertSeqMse(ExportPass): |
| 124 | + """ |
| 125 | + Insert Seq Mse Observer to find the best quant config for certain node's weight. |
| 126 | + """ |
| 127 | + |
| 128 | + seq_mse_ops = { |
| 129 | + torch.ops.aten.conv2d.default |
| 130 | + } |
| 131 | + |
| 132 | + def __init__(self, num_candidates=1000): |
| 133 | + super(InsertSeqMse, self).__init__() |
| 134 | + self.num_candidates = num_candidates |
| 135 | + |
| 136 | + def _insert_seq_mse(self, graph_module: torch.fx.GraphModule) -> torch.fx.GraphModule: |
| 137 | + count = 0 |
| 138 | + for node in graph_module.graph.nodes: |
| 139 | + if node.target in self.seq_mse_ops: |
| 140 | + # extract observer |
| 141 | + weight_node_obs = node.args[1] |
| 142 | + observer = getattr(graph_module, weight_node_obs.name) |
| 143 | + # extract parameters |
| 144 | + weight_node = weight_node_obs.args[0] |
| 145 | + weight_tensor = graph_module.get_parameter( |
| 146 | + weight_node.target |
| 147 | + ).detach() |
| 148 | + bias_tensor = None |
| 149 | + if len(node.args) > 2 and node.args[2] is not None: |
| 150 | + bias_tensor = graph_module.get_parameter( |
| 151 | + node.args[2].args[0].target |
| 152 | + ).detach() |
| 153 | + |
| 154 | + with graph_module.graph.inserting_after(node): |
| 155 | + seq_mse_mod = SeqMseModule( |
| 156 | + nominal_weight=weight_tensor, |
| 157 | + nominal_bias=bias_tensor, |
| 158 | + operator=node, |
| 159 | + observer=observer, |
| 160 | + num_candidates=self.num_candidates, |
| 161 | + ) |
| 162 | + module_name = f"seq_mse_{count}" |
| 163 | + count += 1 |
| 164 | + setattr(graph_module, module_name, seq_mse_mod) |
| 165 | + input_nodes = (node.args[0], node) |
| 166 | + graph_module.graph.create_node("call_module", module_name, input_nodes, {}) |
| 167 | + |
| 168 | + def call(self, graph_module: torch.fx.GraphModule): |
| 169 | + self._insert_seq_mse(graph_module) |
| 170 | + graph_module.recompile() |
| 171 | + return PassResult(graph_module, True) |
| 172 | + |
| 173 | + |
| 174 | +class RemoveSeqMse(ExportPass): |
| 175 | + """ |
| 176 | + Remove Seq Mse before invoking convert_pt2e and update final quantization encoding. |
| 177 | + """ |
| 178 | + def __init__(self): |
| 179 | + super(RemoveSeqMse, self).__init__() |
| 180 | + |
| 181 | + def _remove_seq_mse(self, graph_module: torch.fx.GraphModule) -> torch.fx.GraphModule: |
| 182 | + node_to_erase = [] |
| 183 | + for node in graph_module.graph.nodes: |
| 184 | + if node.op == "call_module" and "seq_mse" in node.name: |
| 185 | + # extract SeqMse module |
| 186 | + module = getattr(graph_module, node.target) |
| 187 | + # rewrite observer method for pre-calculated scale |
| 188 | + scale, zero_point = module.observer.calculate_qparams() |
| 189 | + module.observer.updated_encoding = ( |
| 190 | + scale * module.best_candidate, zero_point |
| 191 | + ) |
| 192 | + module.observer.calculate_qparams = ( |
| 193 | + types.MethodType(lambda s: s.updated_encoding, module.observer) |
| 194 | + ) |
| 195 | + node_to_erase.append(node) |
| 196 | + |
| 197 | + for node in node_to_erase: |
| 198 | + graph_module.graph.erase_node(node) |
| 199 | + |
| 200 | + def call(self, graph_module: torch.fx.GraphModule): |
| 201 | + self._remove_seq_mse(graph_module) |
| 202 | + graph_module.recompile() |
| 203 | + return PassResult(graph_module, True) |
0 commit comments