Enable AIMET in olive quantize CLI (microsoft#2187)

## Describe your changes

Adds support for AIMET algorithms to quantize CLI

## Checklist before requesting a review
- [ ] Add unit tests for this change.
- [ ] Make sure all tests can pass.
- [ ] Update documents if necessary.
- [ ] Lint and apply fixes to your code by running `lintrunner -a`
- [ ] Is this a user-facing change? If yes, give a description of this
change to be included in the release notes.

## (Optional) Issue link

Signed-off-by: Michael Tuttle <mtuttle@qti.qualcomm.com>

Author: michaelgtuttle

olive/cli/quantize.py

@@ -36,6 +36,7 @@ class ImplName(StrEnumBase):
     QUAROT = "quarot"
     AWQ = "awq"
     AUTOGPTQ = "autogptq"
+    AIMET = "aimet"
 
 
 class QuantizeCommand(BaseOliveCLICommand):
@@ -165,6 +166,12 @@ def _get_passes_dict(self, pass_list):
                 "bits": precision_bits_from_precision(self.args.precision),
             },
             "MatMulNBitsToQDQ": {},
+            "AimetQuantization": {
+                "precision": self.args.precision,
+                "activation_type": self.args.act_precision,
+                "data_config": "default_data_config",
+                "techniques": [{"name": self.args.algorithm}],
+            },
         }
 
         passes_dict = {}
@@ -251,4 +258,5 @@ def run(self):
     {"impl_name": ImplName.OLIVE, "pass_type": "OnnxHqqQuantization"},
     {"impl_name": ImplName.OLIVE, "pass_type": "OnnxBlockWiseRtnQuantization"},
     {"impl_name": ImplName.INC, "pass_type": "IncStaticQuantization"},
+    {"impl_name": ImplName.AIMET, "pass_type": "AimetQuantization"},
 ]

olive/constants.py

@@ -66,6 +66,9 @@ class QuantAlgorithm(CaseInsensitiveEnum):
     RTN = "rtn"
     SPINQUANT = "spinquant"
     QUAROT = "quarot"
+    LPBQ = "lpbq"
+    SEQMSE = "seqmse"
+    ADAROUND = "adaround"
 
 
 class QuantEncoding(StrEnumBase):

olive/olive_config.json

@@ -13,8 +13,8 @@
             "supported_providers": [ "*" ],
             "supported_accelerators": [ "*" ],
             "supported_precisions": [ "int4", "int8", "int16" ],
-            "supported_algorithms": [  ],
-            "supported_quantization_encodings": [  ],
+            "supported_algorithms": [ "lpbq", "seqmse", "adaround" ],
+            "supported_quantization_encodings": [ "qdq" ],
             "extra_dependencies": [ "aimet-onnx" ],
             "dataset": "dataset_required"
         },

test/cli/test_cli.py

@@ -324,35 +324,42 @@ def test_shared_cache_delete_all_with_confirmation(mock_AzureContainerClientFact
     mock_factory_instance.delete_all.assert_called_once()
 
 
-@pytest.mark.parametrize("algorithm_name", ["awq", "gptq"])
+@pytest.mark.parametrize("algorithm_name", ["awq", "gptq", "lpbq", "seqmse", "adaround"])
 @patch("olive.workflows.run")
 @patch("huggingface_hub.repo_exists")
 def test_quantize_command(mock_repo_exists, mock_run, algorithm_name, tmp_path):
+    from test.utils import ONNX_MODEL_PATH
+
     # setup
     output_dir = tmp_path / "output_dir"
 
     # setup
     command_args = [
         "quantize",
-        "-m",
-        "dummy_model",
         "--algorithm",
         algorithm_name,
         "-o",
         str(output_dir),
     ]
 
+    model_name = "dummy_model"
     if algorithm_name == "gptq":
         command_args += ["-d", "dummy_dataset"]
         command_args += ["--implementation", "autogptq"]
     if algorithm_name == "awq":
         command_args += ["--implementation", "awq"]
+    if algorithm_name in {"lpbq", "seqmse", "adaround"}:
+        model_name = str(ONNX_MODEL_PATH)
+        command_args += ["-d", "dummy_dataset"]
+        command_args += ["--implementation", "aimet"]
+
+    command_args += ["-m", model_name]
 
     # execute
     cli_main(command_args)
 
     config = mock_run.call_args[0][0]
-    assert config["input_model"]["model_path"] == "dummy_model"
+    assert config["input_model"]["model_path"] == model_name
     assert mock_run.call_count == 1