Fix failures and incorrect configurations in aimet_onnx AMP

Signed-off-by: Michael Tuttle <quic_mtuttle@quicinc.com>
Co-authored-by: Michael Tuttle <quic_mtuttle@quicinc.com>

Author: 2 people

TrainingExtensions/common/src/python/aimet_common/amp/mixed_precision_algo.py

@@ -275,8 +275,8 @@ def _count_and_get_quantizers_flipped(self):
                              (candidate[CandAttr.parameter][CandParam.bitwdith],
                               candidate[CandAttr.parameter][CandParam.data_type])))
 
-        percentage_act_quantizers_flipped = (count_act_quantizers_flipped * 100) / total_act_quantizers
-        percentage_param_quantizers_flipped = (count_param_quantizers_flipped * 100) / total_param_quantizers
+        percentage_act_quantizers_flipped = (count_act_quantizers_flipped * 100) / total_act_quantizers if total_act_quantizers else 0
+        percentage_param_quantizers_flipped = (count_param_quantizers_flipped * 100) / total_param_quantizers if total_param_quantizers else 0
         percentage_quantizers_flipped = ((count_param_quantizers_flipped + count_act_quantizers_flipped) * 100) / \
                                         (total_param_quantizers + total_act_quantizers)
         total_quantizers = total_param_quantizers + total_act_quantizers

TrainingExtensions/onnx/src/python/aimet_onnx/amp/mixed_precision_algo.py

@@ -539,6 +539,9 @@ def _optimize_mp_profile_and_evaluate_model(self):
         """
         Uses OpGraph if available to optimize the mixed precision profile in the sim object
         """
+        # Apply exception rules logic to enforce a valid quantizer configuration
+        self._sim._apply_exception_rules() # pylint: disable = protected-access
+
         # Recompute quantizer encodings
         self._sim.compute_encodings(self.algo_params.forward_pass_callback,
                                     self.algo_params.forward_pass_callback_args)

TrainingExtensions/onnx/src/python/aimet_onnx/amp/utils.py

@@ -42,7 +42,7 @@
 # Import AIMET specific modules
 from aimet_common.amp.utils import CANDIDATE_WITH_DTYPE, get_effective_bitwidth
 from aimet_common.cost_calculator import CostCalculator
-from aimet_onnx.meta.connectedgraph import ConnectedGraph
+from aimet_onnx.meta.connectedgraph import ConnectedGraph, WEIGHT_INDEX
 from aimet_onnx.amp.quantizer_groups import QuantizerGroup
 from aimet_onnx import utils
 from aimet_onnx.quantsim import QuantizationSimModel
@@ -107,7 +107,8 @@ def _get_weight_shape(op):
             if len(layer.output_shape) == 2:
                 # Append 1, 1 to Linear layer's shape
                 layer.output_shape = list(layer.output_shape) + [1, 1]
-            layer.weight_shape = _get_weight_shape(ops[node.name])
+            # If _get_weight_shape returns None, weight index is an activation
+            layer.weight_shape = _get_weight_shape(ops[node.name]) or activation_shapes[node.input[WEIGHT_INDEX]]
             op_database[node.name] = layer
 
     return op_database

TrainingExtensions/onnx/test/python/models/models_for_tests.py

@@ -2704,6 +2704,35 @@ def conv_with_weight_identity_input():
     onnx.checker.check_model(model, True)
     return model
 
+def dynamic_conv_model():
+    model = helper.make_model(
+        graph=helper.make_graph(
+            name="DynamicConvModel",
+            inputs=[helper.make_tensor_value_info('x', TensorProto.FLOAT, shape=[10, 10, 32, 32]),
+                    helper.make_tensor_value_info('y', TensorProto.FLOAT, shape=[10, 10, 1, 1])],
+            outputs=[helper.make_tensor_value_info('model_output', TensorProto.FLOAT, shape=[10, 10, 32, 32])],
+            initializer=[
+                numpy_helper.from_array(np.random.randn(10, 10, 1, 1).astype('float32'), name='add.input'),
+            ],
+            nodes=[
+                helper.make_node(
+                    "Add",
+                    inputs=["y", "add.input"],
+                    outputs=["dynamic_conv.weight"],
+                    name="add"
+                ),
+                helper.make_node(
+                    "Conv",
+                    inputs=["x", "dynamic_conv.weight"],
+                    outputs=["model_output"],
+                    name="conv"
+                )
+            ]
+        )
+    )
+    onnx.checker.check_model(model, True)
+    return model
+
 
 def squeezenet1_0(tmpdir):
     import torchvision

TrainingExtensions/onnx/test/python/test_mixed_precision.py

@@ -49,12 +49,15 @@
 from aimet_onnx.quantsim import QuantizationSimModel
 from aimet_onnx.amp.mixed_precision_algo import GreedyMixedPrecisionAlgo, _compute_sqnr, EvalCallbackFactory
 from aimet_onnx.defs import DataLoader
+from aimet_onnx.utils import make_dummy_input
+from aimet_onnx.mixed_precision import choose_mixed_precision
 
 from aimet_common.defs import QuantizationDataType, CallbackFunc
 from aimet_common.amp.mixed_precision_algo import interpolation_search, brute_force_search, binary_search
 from aimet_common.amp.utils import calculate_starting_bit_ops, AMPSearchAlgo
 
 from .models.test_models import single_residual_model
+from .models import models_for_tests
 
 INPUT_SHAPE = (1, 3, 32, 32)
 
@@ -585,6 +588,69 @@ def test_respect_frozen_encodings(self, sim, forward_pass_callback, eval_callbac
         assert quantizer.bitwidth == 4
 
 
+    @pytest.mark.parametrize("model", (
+            single_residual_model().model, 
+            models_for_tests.dynamic_matmul_model(10),
+            models_for_tests.matmul_with_constant_first_input(),
+            models_for_tests.weight_matmul_model(),
+            models_for_tests.dynamic_conv_model(),
+            models_for_tests.mobilenetv2().model,
+            models_for_tests.depthwise_transposed_conv_model().model,
+            models_for_tests.model_with_split_matmul(),
+            models_for_tests.hierarchical_model().model,
+            ))
+    def test_choose_mixed_precision(self, model, tmpdir):
+        np.random.seed(0)
+
+        sim = QuantizationSimModel(model, default_activation_bw=8, default_param_bw=8, config_file="htp_v73")
+        enabled_quantizers = {q for q in sim.qc_quantize_op_dict.values() if q.enabled}
+        total_bits = 16 * len(enabled_quantizers)
+
+        forward_callback = CallbackFunc(lambda sess, _: sess.run(None, make_dummy_input(model)), None)
+
+        def phase_2_callback(sess, _):
+            bits = sum(q.bitwidth if q.enabled else 16 for q in enabled_quantizers)
+            return bits / total_bits
+
+        # Define dummy eval callbacks
+        eval_callback_phase1 = CallbackFunc(lambda sess, _: np.random.rand())
+        eval_callback_phase2 = CallbackFunc(phase_2_callback, None)
+
+        candidates = [((16, QuantizationDataType.float), (16, QuantizationDataType.float)),
+                      ((16, QuantizationDataType.int), (8, QuantizationDataType.int)),
+                      ((8, QuantizationDataType.int), (8, QuantizationDataType.int))]
+
+        # Apply mixed precision
+        choose_mixed_precision(sim, candidates, eval_callback_phase1, eval_callback_phase2, 0.4, tmpdir, True,
+                               forward_callback)
+        
+        # Assert that no param quantizers are in int16 (not a valid candidate)
+        for name in sim.param_names:
+            quantizer = sim.qc_quantize_op_dict[name]
+            assert not (quantizer.bitwidth == 16 and quantizer.data_type == QuantizationDataType.int)
+
+        # Assert that the final result meets the accuracy metric
+        assert sum(q.bitwidth for q in enabled_quantizers) <= total_bits
+        assert sum(q.bitwidth for q in enabled_quantizers) >= total_bits * 0.6
+
+        # Assert that the final mixed-precision profile obeys config file's exception rules
+        for op in sim.connected_graph.ordered_ops:
+            if not op.type in ("MatMul", "Gemm"):
+                continue
+
+            q1, q2 = sim._get_closest_enabled_quantizer(op.inputs[0]), sim._get_closest_enabled_quantizer(op.inputs[1])
+            if not q1 or not q2:
+                continue
+
+            # Config requires symmetric second input for 16-bit matmul
+            if q2.bitwidth == 16 and not q2.data_type == QuantizationDataType.float:
+                assert q2.use_symmetric_encodings
+
+            # 8 x 16 MatMul is not a valid combination
+            if q1.bitwidth == 8:
+                assert q2.bitwidth == 8
+
+
 class TestAMPv2:
     def test_compute_sqnr(self):
         """ Verify _compute_sqnr() method """