Repository: quic/aimet (language: python)

Issue:
Add utility to clip weights to 7f7f



Discussion / hints:

Fix patch:
diff --git a/TrainingExtensions/torch/src/python/aimet_torch/v2/experimental/__init__.py b/TrainingExtensions/torch/src/python/aimet_torch/v2/experimental/__init__.py
new file mode 100644
index 0000000000..1687c73850
--- /dev/null
+++ b/TrainingExtensions/torch/src/python/aimet_torch/v2/experimental/__init__.py
@@ -0,0 +1,39 @@
+# -*- mode: python -*-
+# =============================================================================
+#  @@-COPYRIGHT-START-@@
+#
+#  Copyright (c) 2024, Qualcomm Innovation Center, Inc. All rights reserved.
+#
+#  Redistribution and use in source and binary forms, with or without
+#  modification, are permitted provided that the following conditions are met:
+#
+#  1. Redistributions of source code must retain the above copyright notice,
+#     this list of conditions and the following disclaimer.
+#
+#  2. Redistributions in binary form must reproduce the above copyright notice,
+#     this list of conditions and the following disclaimer in the documentation
+#     and/or other materials provided with the distribution.
+#
+#  3. Neither the name of the copyright holder nor the names of its contributors
+#     may be used to endorse or promote products derived from this software
+#     without specific prior written permission.
+#
+#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+#  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+#  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+#  ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+#  LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+#  CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+#  SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+#  INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+#  CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+#  ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+#  POSSIBILITY OF SUCH DAMAGE.
+#
+#  SPDX-License-Identifier: BSD-3-Clause
+#
+#  @@-COPYRIGHT-END-@@
+# =============================================================================
+
+# pylint: disable=missing-docstring
+from .quantsim_utils import *  # pylint: disable=import-error
diff --git a/TrainingExtensions/torch/src/python/aimet_torch/v2/experimental.py b/TrainingExtensions/torch/src/python/aimet_torch/v2/experimental/quantsim_utils.py
similarity index 80%
rename from TrainingExtensions/torch/src/python/aimet_torch/v2/experimental.py
rename to TrainingExtensions/torch/src/python/aimet_torch/v2/experimental/quantsim_utils.py
index f7d784762c..dd000f0ea2 100644
--- a/TrainingExtensions/torch/src/python/aimet_torch/v2/experimental.py
+++ b/TrainingExtensions/torch/src/python/aimet_torch/v2/experimental/quantsim_utils.py
@@ -34,18 +34,20 @@
 #
 #  @@-COPYRIGHT-END-@@
 # =============================================================================
-""" Module for experimental features """
+""" Experimental quantsim utilities """
 
 from typing import overload, Callable, Type
-
 import torch
 
+from aimet_common.utils import AimetLogger
 from aimet_common.connected_graph.product import Product
-from aimet_torch.meta.connectedgraph import Op
-from aimet_torch import utils
 import aimet_torch.elementwise_ops as aimet_ops
+from aimet_torch.meta.connectedgraph import Op
+from aimet_torch.v2.quantization.affine.quantizer import AffineQuantizerBase
 from aimet_torch.v2.quantsim import QuantizationSimModel
+from aimet_torch import utils
 
+logger = AimetLogger.get_area_logger(AimetLogger.LogAreas.Quant)
 
 _MATH_INVARIANT_OPS = (
     aimet_ops.Reshape,
@@ -162,3 +164,27 @@ def _set_src_qtzr(x: Product, consumer: Op, qtzr):
 
         for input in op.inputs:
             _set_src_qtzr(input, consumer=op, qtzr=qtzr)
+
+def clip_weights_to_7f7f(sim: 'QuantizationSimModel'):
+    """
+    Clip sim model weights which are 16 bit symmetric to have a max of 0x7f7f when quantized.
+
+    :param sim: Quantsim model to clip weights for
+    """
+    affected_layers = []
+    for name, quant_layer in sim.named_qmodules():
+        # pylint: disable=too-many-boolean-expressions
+        if 'weight' in quant_layer.param_quantizers and \
+                quant_layer.param_quantizers['weight'] is not None and \
+                quant_layer.param_quantizers['weight'].bitwidth == 16 and \
+                isinstance(quant_layer.param_quantizers['weight'], AffineQuantizerBase) and \
+                quant_layer.param_quantizers['weight'].symmetric and \
+                quant_layer.param_quantizers['weight'].is_initialized():
+            clipped_weight = torch.minimum(quant_layer.weight,
+                                           quant_layer.param_quantizers['weight'].get_scale() * 0x7f7f)
+            with torch.no_grad():
+                quant_layer.weight.copy_(clipped_weight)
+
+            affected_layers.append(name)
+    logger_str = f'Clipping weights of the following layers to 0x7f7f max quantized value: {affected_layers}'
+    logger.debug(logger_str)


Test patch:
diff --git a/TrainingExtensions/torch/test/python/experimental/v2/quantsim/test_quantsim_utils.py b/TrainingExtensions/torch/test/python/experimental/v2/quantsim/test_quantsim_utils.py
new file mode 100644
index 0000000000..6b9d62cd44
--- /dev/null
+++ b/TrainingExtensions/torch/test/python/experimental/v2/quantsim/test_quantsim_utils.py
@@ -0,0 +1,99 @@
+# -*- mode: python -*-
+# =============================================================================
+#  @@-COPYRIGHT-START-@@
+#
+#  Copyright (c) 2024, Qualcomm Innovation Center, Inc. All rights reserved.
+#
+#  Redistribution and use in source and binary forms, with or without
+#  modification, are permitted provided that the following conditions are met:
+#
+#  1. Redistributions of source code must retain the above copyright notice,
+#     this list of conditions and the following disclaimer.
+#
+#  2. Redistributions in binary form must reproduce the above copyright notice,
+#     this list of conditions and the following disclaimer in the documentation
+#     and/or other materials provided with the distribution.
+#
+#  3. Neither the name of the copyright holder nor the names of its contributors
+#     may be used to endorse or promote products derived from this software
+#     without specific prior written permission.
+#
+#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+#  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+#  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+#  ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+#  LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+#  CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+#  SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+#  INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+#  CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+#  ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+#  POSSIBILITY OF SUCH DAMAGE.
+#
+#  SPDX-License-Identifier: BSD-3-Clause
+#
+#  @@-COPYRIGHT-END-@@
+# =============================================================================
+""" Test experimental utilities for QuantizationSimModel """
+
+import json
+import os
+import torch
+import tempfile
+from aimet_torch.examples.test_models import SingleResidualWithAvgPool
+from aimet_torch.v2.quantsim import QuantizationSimModel
+from aimet_torch.v2.quantization.affine.backends.torch_builtins import quantize
+from aimet_torch.v2.experimental import clip_weights_to_7f7f
+
+def test_clip_weights_to_7f7f():
+    torch.manual_seed(0)
+    model = SingleResidualWithAvgPool().eval()
+    dummy_input = torch.randn(1, 3, 32, 32)
+
+    # Force all weights to positive to guarantee max quantized value will be > 32639
+    for module in model.modules():
+        if hasattr(module, 'weight'):
+            with torch.no_grad():
+                module.weight.copy_(torch.abs(module.weight))
+
+    quantsim_config = {
+            "defaults": {
+                "ops": {
+                    "is_output_quantized": "True",
+                    "is_symmetric": "False"
+                },
+                "params": {
+                    "is_quantized": "True",
+                    "is_symmetric": "True"
+                },
+                "per_channel_quantization": "True",
+            },
+            "params": {},
+            "op_type": {},
+            "supergroups": [],
+            "model_input": {},
+            "model_output": {}
+        }
+
+    with tempfile.TemporaryDirectory() as tempdir:
+        with open(os.path.join(tempdir, 'quantsim_config.json'), 'w') as f:
+            json.dump(quantsim_config, f)
+
+        qsim = QuantizationSimModel(model, dummy_input, config_file=os.path.join(tempdir, 'quantsim_config.json'), default_param_bw=16)
+    qsim.compute_encodings(lambda m, _: m(dummy_input), None)
+
+    affected_quant_layers = []
+    for _, quant_layer in qsim.named_qmodules():
+        if 'weight' in quant_layer.param_quantizers and quant_layer.param_quantizers['weight'] is not None:
+            encoding = quant_layer.param_quantizers['weight'].get_encoding()
+            quantized_weight = quantize(quant_layer.weight, encoding.scale, encoding.offset, -32768, 32767)
+            assert torch.equal(torch.max(quantized_weight), torch.tensor(32767))
+            affected_quant_layers.append(quant_layer)
+        assert affected_quant_layers
+
+    clip_weights_to_7f7f(qsim)
+
+    for quant_layer in affected_quant_layers:
+        encoding = quant_layer.param_quantizers['weight'].get_encoding()
+        quantized_weight = quantize(quant_layer.weight, encoding.scale, encoding.offset, -32768, 32767)
+        assert torch.equal(torch.max(quantized_weight), torch.tensor(32639))