From 4465ea4c0735ce068aa795c06b3f7ca443d1fb95 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Xavier=20Dupr=C3=A9?= Date: Tue, 2 Jun 2026 17:25:07 +0200 Subject: [PATCH 1/8] enable mobius in fast tests --- test/cli/test_cli_test_model_smoke.py | 221 +++++++++++++++++++++----- 1 file changed, 178 insertions(+), 43 deletions(-) diff --git a/test/cli/test_cli_test_model_smoke.py b/test/cli/test_cli_test_model_smoke.py index b28130658..d5cac958a 100644 --- a/test/cli/test_cli_test_model_smoke.py +++ b/test/cli/test_cli_test_model_smoke.py @@ -3,6 +3,7 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- import argparse +import importlib.util import json import sys import tempfile @@ -28,6 +29,13 @@ ) MAX_ARTIFACT_SIZE_BYTES = 1024 * 1024 +# Supported exporters for the discrepancy test +EXPORTER_MODEL_BUILDER = "model_builder" +EXPORTER_MOBIUS = "mobius" +EXPORTER_TORCH = "torch_export" + +_HAS_MOBIUS = importlib.util.find_spec("mobius") is not None + def _save_local_tiny_llama(model_path: Path): model = LlamaForCausalLM( @@ -133,6 +141,7 @@ def _run_documented_test_model_smoke_flow(tmp_path: Path, model_id: str): class TestCliTestModelSmoke(unittest.TestCase): model_ids = DEFAULT_MODEL_IDS + exporters = (EXPORTER_MODEL_BUILDER,) workdir = None def test_documented_test_model_smoke_flow(self): @@ -173,7 +182,7 @@ def _assert_smoke_flows(self, tmp_path: Path): self._assert_file_size_below_limit(run_output_dir / "model.onnx.data") def test_model_discrepancy(self): - """Verify that OnnxDiscrepancyCheck runs successfully when auto-injected via --test.""" + """Verify that OnnxDiscrepancyCheck runs successfully with the configured exporter.""" if self.workdir is None: with tempfile.TemporaryDirectory() as temp_dir: self._assert_discrepancy(Path(temp_dir)) @@ -183,48 +192,165 @@ def test_model_discrepancy(self): self._assert_discrepancy(workdir) def _assert_discrepancy(self, tmp_path: Path): - for model_id in self.model_ids: - with self.subTest(model_id=model_id): - model_name = model_id.replace("/", "--") - model_path = tmp_path / "models" / f"{model_name}-disc" - config_output_dir = tmp_path / f"{model_name}-disc-cfg" - test_model_dir = tmp_path / f"{model_name}-disc-test-model" - run_output_dir = tmp_path / f"{model_name}-disc-run" - - _save_local_tiny_llama(model_path) - _run_cli_main( - [ - "optimize", - "-m", - str(model_path), - "--device", - "cpu", - "--provider", - "CPUExecutionProvider", - "--precision", - "int4", - "--output_path", - str(config_output_dir), - "--dry_run", - ] - ) - - config_path = config_output_dir / "config.json" - assert config_path.exists() - _set_offline_gptq_data_config(config_path) - - # Run with --test; OnnxDiscrepancyCheck is auto-injected and reports discrepancy metrics (fails only if thresholds are configured) - _run_cli_main( - [ - "run", - "--config", - str(config_path), - "--test", - str(test_model_dir), - "--output_path", - str(run_output_dir), - ] - ) + for exporter in self.exporters: + if exporter == EXPORTER_MOBIUS and not _HAS_MOBIUS: + continue + for model_id in self.model_ids: + with self.subTest(model_id=model_id, exporter=exporter): + if exporter == EXPORTER_MODEL_BUILDER: + self._assert_discrepancy_model_builder(tmp_path, model_id) + elif exporter == EXPORTER_MOBIUS: + self._assert_discrepancy_mobius(tmp_path, model_id) + elif exporter == EXPORTER_TORCH: + self._assert_discrepancy_torch_export(tmp_path, model_id) + + def _assert_discrepancy_model_builder(self, tmp_path: Path, model_id: str): + model_name = model_id.replace("/", "--") + model_path = tmp_path / "models" / f"{model_name}-disc" + config_output_dir = tmp_path / f"{model_name}-disc-cfg" + test_model_dir = tmp_path / f"{model_name}-disc-test-model" + run_output_dir = tmp_path / f"{model_name}-disc-run" + + _save_local_tiny_llama(model_path) + _run_cli_main( + [ + "optimize", + "-m", + str(model_path), + "--device", + "cpu", + "--provider", + "CPUExecutionProvider", + "--precision", + "int4", + "--output_path", + str(config_output_dir), + "--dry_run", + ] + ) + + config_path = config_output_dir / "config.json" + assert config_path.exists() + _set_offline_gptq_data_config(config_path) + + # Run with --test; OnnxDiscrepancyCheck is auto-injected and reports discrepancy metrics + _run_cli_main( + [ + "run", + "--config", + str(config_path), + "--test", + str(test_model_dir), + "--output_path", + str(run_output_dir), + ] + ) + + def _assert_discrepancy_mobius(self, tmp_path: Path, model_id: str): + model_name = model_id.replace("/", "--") + model_path = tmp_path / "models" / f"{model_name}-mobius-disc" + run_output_dir = tmp_path / f"{model_name}-mobius-disc-run" + + _save_local_tiny_llama(model_path) + + run_config = { + "input_model": { + "type": "HfModel", + "model_path": str(model_path), + }, + "systems": { + "local_system": { + "type": "LocalSystem", + "accelerators": [{"device": "cpu", "execution_providers": ["CPUExecutionProvider"]}], + } + }, + "output_dir": str(run_output_dir), + "host": "local_system", + "target": "local_system", + "passes": { + "mobius_builder": { + "type": "MobiusBuilder", + "precision": "fp32", + "runtime": "none", + }, + "discrepancy_check": { + "type": "OnnxDiscrepancyCheck", + "reference_model_path": str(model_path), + }, + }, + } + config_path = tmp_path / f"{model_name}-mobius-disc-config.json" + config_path.write_text(json.dumps(run_config, indent=2)) + + _run_cli_main( + [ + "run", + "--config", + str(config_path), + "--output_path", + str(run_output_dir), + ] + ) + + def _assert_discrepancy_torch_export(self, tmp_path: Path, model_id: str): + import torch + + from olive.hardware.accelerator import AcceleratorSpec, Device + from olive.hardware.constants import ExecutionProvider + from olive.model import ONNXModelHandler + from olive.passes.olive_pass import create_pass_from_dict + from olive.passes.onnx.discrepancy_check import OnnxDiscrepancyCheck + + model_name = model_id.replace("/", "--") + model_path = tmp_path / "models" / f"{model_name}-torch-disc" + onnx_dir = tmp_path / f"{model_name}-torch-disc-onnx" + onnx_dir.mkdir(parents=True, exist_ok=True) + + _save_local_tiny_llama(model_path) + + from transformers import AutoModelForCausalLM + + ref_model = AutoModelForCausalLM.from_pretrained(model_path) + ref_model.eval() + + class _LogitsOnlyWrapper(torch.nn.Module): + def __init__(self, model): + super().__init__() + self.model = model + + def forward(self, input_ids, attention_mask): + return self.model(input_ids=input_ids, attention_mask=attention_mask, use_cache=False).logits + + wrapper = _LogitsOnlyWrapper(ref_model) + wrapper.eval() + + dummy_input_ids = torch.randint(0, 32, (1, 4)) + dummy_attention_mask = torch.ones(1, 4, dtype=torch.int64) + onnx_path = onnx_dir / "model.onnx" + + torch.onnx.export( + wrapper, + (dummy_input_ids, dummy_attention_mask), + str(onnx_path), + input_names=["input_ids", "attention_mask"], + output_names=["logits"], + dynamo=False, + ) + assert onnx_path.exists() + + onnx_model = ONNXModelHandler(model_path=str(onnx_dir), onnx_file_name="model.onnx") + accelerator_spec = AcceleratorSpec( + accelerator_type=Device.CPU, execution_provider=ExecutionProvider.CPUExecutionProvider + ) + discrepancy_pass = create_pass_from_dict( + OnnxDiscrepancyCheck, + {"reference_model_path": str(model_path)}, + disable_search=True, + accelerator_spec=accelerator_spec, + ) + output_path = tmp_path / f"{model_name}-torch-disc-output" + result = discrepancy_pass.run(onnx_model, output_path) + assert isinstance(result, ONNXModelHandler) def _assert_file_size_below_limit(self, path: Path): assert path.exists() @@ -239,6 +365,13 @@ def _parse_args(): parser = argparse.ArgumentParser(add_help=False) parser.add_argument("--workdir") parser.add_argument("--model-id", dest="model_ids", action="append") + parser.add_argument( + "--exporter", + dest="exporters", + action="append", + choices=[EXPORTER_MODEL_BUILDER, EXPORTER_MOBIUS, EXPORTER_TORCH], + help="Exporter(s) to test for discrepancy check (can be repeated).", + ) return parser.parse_known_args() @@ -248,4 +381,6 @@ def _parse_args(): TestCliTestModelSmoke.workdir = Path(parsed_args.workdir) if parsed_args.model_ids: TestCliTestModelSmoke.model_ids = tuple(parsed_args.model_ids) + if parsed_args.exporters: + TestCliTestModelSmoke.exporters = tuple(parsed_args.exporters) unittest.main(argv=[__file__, *remaining]) From 1272150b23fc44ab45a005889987dba511271ed0 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 2 Jun 2026 16:17:06 +0000 Subject: [PATCH 2/8] Add missing --test arg in mobius smoke run --- test/cli/test_cli_test_model_smoke.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/test/cli/test_cli_test_model_smoke.py b/test/cli/test_cli_test_model_smoke.py index d5cac958a..9d86d73df 100644 --- a/test/cli/test_cli_test_model_smoke.py +++ b/test/cli/test_cli_test_model_smoke.py @@ -249,6 +249,7 @@ def _assert_discrepancy_model_builder(self, tmp_path: Path, model_id: str): def _assert_discrepancy_mobius(self, tmp_path: Path, model_id: str): model_name = model_id.replace("/", "--") model_path = tmp_path / "models" / f"{model_name}-mobius-disc" + test_model_dir = tmp_path / f"{model_name}-mobius-disc-test-model" run_output_dir = tmp_path / f"{model_name}-mobius-disc-run" _save_local_tiny_llama(model_path) @@ -287,6 +288,8 @@ def _assert_discrepancy_mobius(self, tmp_path: Path, model_id: str): "run", "--config", str(config_path), + "--test", + str(test_model_dir), "--output_path", str(run_output_dir), ] From a8d6539574f603c653649d1866631302d7b20787 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 2 Jun 2026 16:51:36 +0000 Subject: [PATCH 3/8] Refactor shared discrepancy run invocation in smoke tests --- test/cli/test_cli_test_model_smoke.py | 38 +++++++++++---------------- 1 file changed, 16 insertions(+), 22 deletions(-) diff --git a/test/cli/test_cli_test_model_smoke.py b/test/cli/test_cli_test_model_smoke.py index 9d86d73df..f10ba9ce3 100644 --- a/test/cli/test_cli_test_model_smoke.py +++ b/test/cli/test_cli_test_model_smoke.py @@ -204,6 +204,20 @@ def _assert_discrepancy(self, tmp_path: Path): elif exporter == EXPORTER_TORCH: self._assert_discrepancy_torch_export(tmp_path, model_id) + @staticmethod + def _run_discrepancy_with_test(config_path: Path, test_model_dir: Path, run_output_dir: Path): + _run_cli_main( + [ + "run", + "--config", + str(config_path), + "--test", + str(test_model_dir), + "--output_path", + str(run_output_dir), + ] + ) + def _assert_discrepancy_model_builder(self, tmp_path: Path, model_id: str): model_name = model_id.replace("/", "--") model_path = tmp_path / "models" / f"{model_name}-disc" @@ -234,17 +248,7 @@ def _assert_discrepancy_model_builder(self, tmp_path: Path, model_id: str): _set_offline_gptq_data_config(config_path) # Run with --test; OnnxDiscrepancyCheck is auto-injected and reports discrepancy metrics - _run_cli_main( - [ - "run", - "--config", - str(config_path), - "--test", - str(test_model_dir), - "--output_path", - str(run_output_dir), - ] - ) + self._run_discrepancy_with_test(config_path, test_model_dir, run_output_dir) def _assert_discrepancy_mobius(self, tmp_path: Path, model_id: str): model_name = model_id.replace("/", "--") @@ -283,17 +287,7 @@ def _assert_discrepancy_mobius(self, tmp_path: Path, model_id: str): config_path = tmp_path / f"{model_name}-mobius-disc-config.json" config_path.write_text(json.dumps(run_config, indent=2)) - _run_cli_main( - [ - "run", - "--config", - str(config_path), - "--test", - str(test_model_dir), - "--output_path", - str(run_output_dir), - ] - ) + self._run_discrepancy_with_test(config_path, test_model_dir, run_output_dir) def _assert_discrepancy_torch_export(self, tmp_path: Path, model_id: str): import torch From bd2a22c724a5c355a08e31f8f4507db2086307a7 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 2 Jun 2026 16:54:01 +0000 Subject: [PATCH 4/8] Use CLI --test flow for torch discrepancy smoke path --- test/cli/test_cli_test_model_smoke.py | 79 +++++++++------------------ 1 file changed, 26 insertions(+), 53 deletions(-) diff --git a/test/cli/test_cli_test_model_smoke.py b/test/cli/test_cli_test_model_smoke.py index f10ba9ce3..7b951a881 100644 --- a/test/cli/test_cli_test_model_smoke.py +++ b/test/cli/test_cli_test_model_smoke.py @@ -290,64 +290,37 @@ def _assert_discrepancy_mobius(self, tmp_path: Path, model_id: str): self._run_discrepancy_with_test(config_path, test_model_dir, run_output_dir) def _assert_discrepancy_torch_export(self, tmp_path: Path, model_id: str): - import torch - - from olive.hardware.accelerator import AcceleratorSpec, Device - from olive.hardware.constants import ExecutionProvider - from olive.model import ONNXModelHandler - from olive.passes.olive_pass import create_pass_from_dict - from olive.passes.onnx.discrepancy_check import OnnxDiscrepancyCheck - model_name = model_id.replace("/", "--") model_path = tmp_path / "models" / f"{model_name}-torch-disc" - onnx_dir = tmp_path / f"{model_name}-torch-disc-onnx" - onnx_dir.mkdir(parents=True, exist_ok=True) + test_model_dir = tmp_path / f"{model_name}-torch-disc-test-model" + run_output_dir = tmp_path / f"{model_name}-torch-disc-run" _save_local_tiny_llama(model_path) - from transformers import AutoModelForCausalLM - - ref_model = AutoModelForCausalLM.from_pretrained(model_path) - ref_model.eval() - - class _LogitsOnlyWrapper(torch.nn.Module): - def __init__(self, model): - super().__init__() - self.model = model - - def forward(self, input_ids, attention_mask): - return self.model(input_ids=input_ids, attention_mask=attention_mask, use_cache=False).logits - - wrapper = _LogitsOnlyWrapper(ref_model) - wrapper.eval() - - dummy_input_ids = torch.randint(0, 32, (1, 4)) - dummy_attention_mask = torch.ones(1, 4, dtype=torch.int64) - onnx_path = onnx_dir / "model.onnx" - - torch.onnx.export( - wrapper, - (dummy_input_ids, dummy_attention_mask), - str(onnx_path), - input_names=["input_ids", "attention_mask"], - output_names=["logits"], - dynamo=False, - ) - assert onnx_path.exists() - - onnx_model = ONNXModelHandler(model_path=str(onnx_dir), onnx_file_name="model.onnx") - accelerator_spec = AcceleratorSpec( - accelerator_type=Device.CPU, execution_provider=ExecutionProvider.CPUExecutionProvider - ) - discrepancy_pass = create_pass_from_dict( - OnnxDiscrepancyCheck, - {"reference_model_path": str(model_path)}, - disable_search=True, - accelerator_spec=accelerator_spec, - ) - output_path = tmp_path / f"{model_name}-torch-disc-output" - result = discrepancy_pass.run(onnx_model, output_path) - assert isinstance(result, ONNXModelHandler) + run_config = { + "input_model": { + "type": "HfModel", + "model_path": str(model_path), + }, + "systems": { + "local_system": { + "type": "LocalSystem", + "accelerators": [{"device": "cpu", "execution_providers": ["CPUExecutionProvider"]}], + } + }, + "output_dir": str(run_output_dir), + "host": "local_system", + "target": "local_system", + "passes": { + "conversion": { + "type": "OnnxConversion", + "use_dynamo_exporter": False, + } + }, + } + config_path = tmp_path / f"{model_name}-torch-disc-config.json" + config_path.write_text(json.dumps(run_config, indent=2)) + self._run_discrepancy_with_test(config_path, test_model_dir, run_output_dir) def _assert_file_size_below_limit(self, path: Path): assert path.exists() From 6b59f34d47b747e0645ea0eccf5cc8644d563bff Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Xavier=20Dupr=C3=A9?= Date: Tue, 2 Jun 2026 19:03:09 +0200 Subject: [PATCH 5/8] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- test/cli/test_cli_test_model_smoke.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/cli/test_cli_test_model_smoke.py b/test/cli/test_cli_test_model_smoke.py index 7b951a881..164891fb8 100644 --- a/test/cli/test_cli_test_model_smoke.py +++ b/test/cli/test_cli_test_model_smoke.py @@ -194,7 +194,7 @@ def test_model_discrepancy(self): def _assert_discrepancy(self, tmp_path: Path): for exporter in self.exporters: if exporter == EXPORTER_MOBIUS and not _HAS_MOBIUS: - continue + self.fail("Requested exporter 'mobius' but mobius-ai is not installed. Install mobius-ai or remove '--exporter mobius'.") for model_id in self.model_ids: with self.subTest(model_id=model_id, exporter=exporter): if exporter == EXPORTER_MODEL_BUILDER: From 8aa0114cfdc943427ffca557325381ec88df2cbc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Xavier=20Dupr=C3=A9?= Date: Tue, 2 Jun 2026 19:03:20 +0200 Subject: [PATCH 6/8] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- test/cli/test_cli_test_model_smoke.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/test/cli/test_cli_test_model_smoke.py b/test/cli/test_cli_test_model_smoke.py index 164891fb8..7db28d69c 100644 --- a/test/cli/test_cli_test_model_smoke.py +++ b/test/cli/test_cli_test_model_smoke.py @@ -203,6 +203,8 @@ def _assert_discrepancy(self, tmp_path: Path): self._assert_discrepancy_mobius(tmp_path, model_id) elif exporter == EXPORTER_TORCH: self._assert_discrepancy_torch_export(tmp_path, model_id) + else: + self.fail(f"Unknown exporter: {exporter!r}") @staticmethod def _run_discrepancy_with_test(config_path: Path, test_model_dir: Path, run_output_dir: Path): From caf66446639fb92554cad288a073e0f4067688d7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Xavier=20Dupr=C3=A9?= Date: Tue, 2 Jun 2026 19:03:35 +0200 Subject: [PATCH 7/8] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- test/cli/test_cli_test_model_smoke.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/cli/test_cli_test_model_smoke.py b/test/cli/test_cli_test_model_smoke.py index 7db28d69c..9bb2ac6c3 100644 --- a/test/cli/test_cli_test_model_smoke.py +++ b/test/cli/test_cli_test_model_smoke.py @@ -282,8 +282,8 @@ def _assert_discrepancy_mobius(self, tmp_path: Path, model_id: str): }, "discrepancy_check": { "type": "OnnxDiscrepancyCheck", - "reference_model_path": str(model_path), - }, + "reference_model_path": str(test_model_dir), + } }, } config_path = tmp_path / f"{model_name}-mobius-disc-config.json" From ca869b1e59ed52dcdc3c79c2b6ecdd67cc6b485d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Xavier=20Dupr=C3=A9?= Date: Wed, 3 Jun 2026 10:48:43 +0200 Subject: [PATCH 8/8] lint --- test/cli/test_cli_test_model_smoke.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/test/cli/test_cli_test_model_smoke.py b/test/cli/test_cli_test_model_smoke.py index 9bb2ac6c3..c9d1bdc7e 100644 --- a/test/cli/test_cli_test_model_smoke.py +++ b/test/cli/test_cli_test_model_smoke.py @@ -194,7 +194,9 @@ def test_model_discrepancy(self): def _assert_discrepancy(self, tmp_path: Path): for exporter in self.exporters: if exporter == EXPORTER_MOBIUS and not _HAS_MOBIUS: - self.fail("Requested exporter 'mobius' but mobius-ai is not installed. Install mobius-ai or remove '--exporter mobius'.") + self.fail( + "Requested exporter 'mobius' but mobius-ai is not installed. Install mobius-ai or remove '--exporter mobius'." + ) for model_id in self.model_ids: with self.subTest(model_id=model_id, exporter=exporter): if exporter == EXPORTER_MODEL_BUILDER: @@ -283,7 +285,7 @@ def _assert_discrepancy_mobius(self, tmp_path: Path, model_id: str): "discrepancy_check": { "type": "OnnxDiscrepancyCheck", "reference_model_path": str(test_model_dir), - } + }, }, } config_path = tmp_path / f"{model_name}-mobius-disc-config.json"