Support pre-quantized models in aimet pass

olive/olive_config.json

@@ -602,7 +602,7 @@
         }
     },
     "extra_dependencies": {
-        "aimet-onnx": [ "aimet-onnx>=2.10.0" ],
+        "aimet-onnx": [ "aimet-onnx>=2.12.0" ],
         "auto-opt": [ "optimum" ],
         "azureml": [ "azure-ai-ml>=1.11.1", "azure-identity" ],
         "bnb": [ "bitsandbytes", "triton" ],

olive/passes/onnx/aimet_quantization.py

@@ -62,11 +62,15 @@ class QuantScheme(StrEnumBase):
     TF_ENHANCED = "tf_enhanced"
 
 
-def _has_quantization_nodes(model: onnx.ModelProto):
-    quantize_op_types = {"QuantizeLinear", "DequantizeLinear", "DynamicQuantizeLinear", "MatMulNBits"}
+def _has_qdq_nodes(model: onnx.ModelProto):
+    quantize_op_types = {"QuantizeLinear", "DequantizeLinear"}
     return any(node.op_type in quantize_op_types for node in model.graph.node)
 
 
+def _has_dynamic_quantization(model: onnx.ModelProto):
+    return any(node.op_type == "DynamicQuantizeLinear" for node in model.graph.node)
+
+
 def _disable_quantizer(sim, tensor_name: str):
     quantizer = sim.qc_quantize_op_dict.get(tensor_name)
     if quantizer and not quantizer.is_encoding_frozen():
@@ -201,11 +205,17 @@ def _run_for_config(
 
         onnx_model = onnx.load(model.model_path)
 
-        if _has_quantization_nodes(onnx_model):
-            raise NotImplementedError("AIMET Quantization does not support pre-quantized models")
+        if _has_dynamic_quantization(onnx_model):
+            raise NotImplementedError("AIMET Quantization does not support dynamically quantized models.")
 
         with tempfile.TemporaryDirectory(prefix="olive_tmp") as tmp_dir:
-            sim = aimet_onnx.QuantizationSimModel(
+            # pylint:disable = protected-access
+            sim_initializer = (
+                aimet_onnx.QuantizationSimModel
+                if not _has_qdq_nodes(onnx_model)
+                else aimet_onnx.QuantizationSimModel._from_onnx_qdq
+            )
+            sim = sim_initializer(
                 onnx_model,
                 param_type=param_type,
                 activation_type=act_type,

test/passes/onnx/test_aimet_quantization.py

@@ -236,7 +236,7 @@ def test_aimet_quantization_excludes_op_types(tmp_path, op_types, disabled_quant
 
 @pytest.mark.skipif(not IS_LINUX, reason="Only run on linux")
 @pytest.mark.skipif(CUDA_AVAILABLE, reason="Only run on cpu tests")
-def test_aimet_quantization_raises_error_with_prequantized_model(tmp_path):
+def test_aimet_quantization_preserves_quantization_in_prequantized_model(tmp_path):
     input_model = dummy_quantized_onnx_model(tmp_path / "dummy_model.onnx")
     config = {
         "data_config": DataConfig(
@@ -249,8 +249,19 @@ def test_aimet_quantization_raises_error_with_prequantized_model(tmp_path):
     }
     p = create_pass_from_dict(AimetQuantization, config, disable_search=True)
 
-    with pytest.raises(NotImplementedError):
-        p.run(input_model, tmp_path)
+    out = p.run(input_model, tmp_path)
+
+    model = onnx.load(out.model_path)
+
+    tensor_to_quantizer = {
+        node.input[0]: node for node in model.graph.node if node.op_type in ("QuantizeLinear", "DequantizeLinear")
+    }
+
+    weight_quantizer = tensor_to_quantizer["weight_dq"]
+    weight_scale = [t for t in model.graph.initializer if t.name == weight_quantizer.input[1]]
+    weight_scale = onnx.numpy_helper.to_array(weight_scale[0])
+    assert weight_scale == np.array(0.1).astype(np.float32)
+    assert "input" in tensor_to_quantizer
 
 
 @pytest.mark.skipif(not IS_LINUX, reason="Only run on linux")