From a4ad838ecb83d103ee2a338dda5746a72492262b Mon Sep 17 00:00:00 2001 From: Shiyi Zheng Date: Wed, 26 Aug 2026 15:10:35 +0800 Subject: [PATCH 1/2] fix(build): honor no-optimize in CLI pipelines --- .../cpu/sentence-similarity_fp16_config.json | 79 ++++++++ .../cpu/sentence-similarity_fp32_config.json | 57 ++++++ src/winml/modelkit/commands/build.py | 54 ++--- tests/unit/commands/test_build.py | 187 +++++++++++++++++- 4 files changed, 348 insertions(+), 29 deletions(-) create mode 100644 examples/recipes/dell-research-harvard_lt-wikidata-comp-multi/cpu/cpu/sentence-similarity_fp16_config.json create mode 100644 examples/recipes/dell-research-harvard_lt-wikidata-comp-multi/cpu/cpu/sentence-similarity_fp32_config.json diff --git a/examples/recipes/dell-research-harvard_lt-wikidata-comp-multi/cpu/cpu/sentence-similarity_fp16_config.json b/examples/recipes/dell-research-harvard_lt-wikidata-comp-multi/cpu/cpu/sentence-similarity_fp16_config.json new file mode 100644 index 000000000..8de2655c2 --- /dev/null +++ b/examples/recipes/dell-research-harvard_lt-wikidata-comp-multi/cpu/cpu/sentence-similarity_fp16_config.json @@ -0,0 +1,79 @@ +{ + "export": { + "opset_version": 17, + "batch_size": 1, + "export_params": true, + "do_constant_folding": true, + "verbose": false, + "dynamo": false, + "enable_hierarchy_tags": true, + "clean_onnx": false, + "hierarchy_tag_format": "full", + "input_tensors": [ + { + "name": "input_ids", + "dtype": "int32", + "shape": [ + 1, + 512 + ], + "value_range": [ + 0, + 250002 + ] + }, + { + "name": "attention_mask", + "dtype": "int32", + "shape": [ + 1, + 512 + ], + "value_range": [ + 0, + 2 + ] + } + ], + "output_tensors": [ + { + "name": "last_hidden_state" + } + ], + "compatibility": { + "transformers_attention": "eager" + } + }, + "optim": { + "clamp_constant_values": true + }, + "quant": { + "mode": "fp16", + "samples": 10, + "calibration_method": "minmax", + "weight_type": "uint8", + "activation_type": "uint8", + "per_channel": false, + "symmetric": false, + "weight_symmetric": null, + "activation_symmetric": null, + "save_calibration": false, + "distribution": "uniform", + "seed": null, + "calibration_load_path": null, + "calibration_save_path": null, + "op_types_to_quantize": null, + "nodes_to_exclude": null, + "task": "sentence-similarity", + "model_id": "dell-research-harvard/lt-wikidata-comp-multi", + "model_type": "xlm-roberta", + "fp16_keep_io_types": true, + "fp16_op_block_list": null + }, + "compile": null, + "loader": { + "task": "sentence-similarity", + "model_class": "AutoModel", + "model_type": "xlm-roberta" + } +} \ No newline at end of file diff --git a/examples/recipes/dell-research-harvard_lt-wikidata-comp-multi/cpu/cpu/sentence-similarity_fp32_config.json b/examples/recipes/dell-research-harvard_lt-wikidata-comp-multi/cpu/cpu/sentence-similarity_fp32_config.json new file mode 100644 index 000000000..a1b99dbc2 --- /dev/null +++ b/examples/recipes/dell-research-harvard_lt-wikidata-comp-multi/cpu/cpu/sentence-similarity_fp32_config.json @@ -0,0 +1,57 @@ +{ + "export": { + "opset_version": 17, + "batch_size": 1, + "export_params": true, + "do_constant_folding": true, + "verbose": false, + "dynamo": false, + "enable_hierarchy_tags": true, + "clean_onnx": false, + "hierarchy_tag_format": "full", + "input_tensors": [ + { + "name": "input_ids", + "dtype": "int32", + "shape": [ + 1, + 512 + ], + "value_range": [ + 0, + 250002 + ] + }, + { + "name": "attention_mask", + "dtype": "int32", + "shape": [ + 1, + 512 + ], + "value_range": [ + 0, + 2 + ] + } + ], + "output_tensors": [ + { + "name": "last_hidden_state" + } + ], + "compatibility": { + "transformers_attention": "eager" + } + }, + "optim": { + "clamp_constant_values": true + }, + "quant": null, + "compile": null, + "loader": { + "task": "sentence-similarity", + "model_class": "AutoModel", + "model_type": "xlm-roberta" + } +} \ No newline at end of file diff --git a/src/winml/modelkit/commands/build.py b/src/winml/modelkit/commands/build.py index a56d9f52e..52712541f 100644 --- a/src/winml/modelkit/commands/build.py +++ b/src/winml/modelkit/commands/build.py @@ -2068,6 +2068,7 @@ def _build_hf_pipeline( max_iters: int = extra_kwargs.pop("hack_max_optim_iterations", 3) allow_unsupported_nodes: bool = extra_kwargs.pop("allow_unsupported_nodes", False) + skip_optimize = extra_kwargs.pop("skip_optimize", False) or config.skip_optimize model_label = model_id or "random-init" # ── Validate + setup ───────────────────────────────────────── @@ -2142,18 +2143,19 @@ def _name(base: str) -> str: stage_timings.append(("Export", _export_elapsed)) # ── Optimize stage ─────────────────────────────────────────── - current_path, _ = _run_optimize_stage( - config=config, - model_path=current_path, - optimized_path=optimized_path, - ep=ep, - device=device, - max_iters=max_iters, - stage_timings=stage_timings, - show_io_first=False, - analyze_output_path=analyze_result_path, - allow_unsupported_nodes=allow_unsupported_nodes, - ) + if not skip_optimize: + current_path, _ = _run_optimize_stage( + config=config, + model_path=current_path, + optimized_path=optimized_path, + ep=ep, + device=device, + max_iters=max_iters, + stage_timings=stage_timings, + show_io_first=False, + analyze_output_path=analyze_result_path, + allow_unsupported_nodes=allow_unsupported_nodes, + ) # Persist config after autoconf config_path.write_text(json.dumps(config.to_dict(), indent=2)) @@ -2209,6 +2211,7 @@ def _build_onnx_pipeline( max_iters: int = extra_kwargs.pop("hack_max_optim_iterations", 3) allow_unsupported_nodes: bool = extra_kwargs.pop("allow_unsupported_nodes", False) + skip_optimize = extra_kwargs.pop("skip_optimize", False) # ── Validate + setup ───────────────────────────────────────── if not onnx_path.exists(): @@ -2251,21 +2254,22 @@ def _build_onnx_pipeline( # config before any stage reads it, otherwise the optimize stage will still # run on integer ops and the quantize stage may try to re-quantize. ensure_pre_quantized_stamped(config, current_path) + skip_optimize = skip_optimize or config.skip_optimize # ── Optimize stage (first stage for ONNX — show I/O here) ──── - current_path, _ = _run_optimize_stage( - config=config, - model_path=current_path, - optimized_path=optimized_path, - ep=ep, - device=device, - max_iters=max_iters, - stage_timings=stage_timings, - show_io_first=True, - analyze_output_path=analyze_result_path, - allow_unsupported_nodes=allow_unsupported_nodes, - skip_optimize=config.skip_optimize, - ) + if not skip_optimize: + current_path, _ = _run_optimize_stage( + config=config, + model_path=current_path, + optimized_path=optimized_path, + ep=ep, + device=device, + max_iters=max_iters, + stage_timings=stage_timings, + show_io_first=True, + analyze_output_path=analyze_result_path, + allow_unsupported_nodes=allow_unsupported_nodes, + ) config_path.write_text(json.dumps(config.to_dict(), indent=2)) diff --git a/tests/unit/commands/test_build.py b/tests/unit/commands/test_build.py index 37f6bde79..b92f1289e 100644 --- a/tests/unit/commands/test_build.py +++ b/tests/unit/commands/test_build.py @@ -2021,6 +2021,185 @@ def test_quant_model_type_carried_into_quantize_stage( mock_quantize.assert_called_once() +class TestBuildPipelineSkipOptimize: + """Concrete CLI pipelines must consume the no-optimize build control.""" + + @pytest.mark.parametrize("quant_mode", [None, "fp16"]) + @patch("winml.modelkit.onnx.copy_onnx_model") + @patch("winml.modelkit.commands.build._run_compile_stage") + @patch("winml.modelkit.commands.build._run_quantize_stage") + @patch("winml.modelkit.commands.build._run_optimize_stage") + @patch("winml.modelkit.utils.console.StageLive") + @patch("winml.modelkit.export.export_onnx") + @patch("winml.modelkit.build.hf._load_model") + def test_hf_skip_optimize_preserves_quantize( + self, + mock_load_model: MagicMock, + mock_export_onnx: MagicMock, + mock_stage_live: MagicMock, + mock_optimize: MagicMock, + mock_quantize: MagicMock, + mock_compile: MagicMock, + mock_copy: MagicMock, + quant_mode: str | None, + tmp_path: Path, + ) -> None: + from winml.modelkit.commands.build import _build_hf_pipeline + + mock_stage_live.return_value.__enter__ = MagicMock(return_value=MagicMock()) + mock_stage_live.return_value.__exit__ = MagicMock(return_value=False) + output_dir = tmp_path / "out" + export_path = output_dir / "export.onnx" + quantized_path = output_dir / "quantized.onnx" + mock_quantize.return_value = quantized_path if quant_mode else export_path + mock_compile.side_effect = lambda **kwargs: kwargs["current_path"] + + config = MagicMock() + config.skip_optimize = False + config.loader.model_type = "xlm-roberta" + config.loader.task = "sentence-similarity" + config.loader.model_class = "AutoModel" + config.export = MagicMock() + config.quant = None if quant_mode is None else MagicMock(mode=quant_mode) + if config.quant is not None: + config.quant.model_type = None + config.to_dict.return_value = {} + + result = _build_hf_pipeline( + config=config, + model_id="example/model", + output_dir=output_dir, + rebuild=True, + cache_key=None, + ep="cpu", + device="cpu", + extra_kwargs={"skip_optimize": True}, + ) + + assert result is not None + assert [name for name, _ in result] == ["Export"] + mock_optimize.assert_not_called() + assert mock_quantize.call_args.kwargs["current_path"] == export_path + expected_final_source = quantized_path if quant_mode else export_path + mock_copy.assert_called_once_with(expected_final_source, output_dir / "model.onnx") + + @pytest.mark.parametrize("quant_mode", [None, "fp16"]) + @patch("winml.modelkit.onnx.copy_onnx_model") + @patch("winml.modelkit.commands.build._run_compile_stage") + @patch("winml.modelkit.commands.build._run_quantize_stage") + @patch("winml.modelkit.commands.build._run_optimize_stage") + def test_onnx_skip_optimize_preserves_quantize( + self, + mock_optimize: MagicMock, + mock_quantize: MagicMock, + mock_compile: MagicMock, + mock_copy: MagicMock, + quant_mode: str | None, + tmp_path: Path, + ) -> None: + from winml.modelkit.commands.build import _build_onnx_pipeline + + onnx_path = tmp_path / "input.onnx" + onnx_path.write_bytes(b"fake-onnx-data") + output_dir = tmp_path / "out" + copied_input = output_dir / onnx_path.name + quantized_path = output_dir / "input_quantized.onnx" + mock_quantize.return_value = quantized_path if quant_mode else copied_input + mock_compile.side_effect = lambda **kwargs: kwargs["current_path"] + + config = MagicMock() + config.skip_optimize = False + config.quant = None if quant_mode is None else MagicMock(mode=quant_mode) + config.to_dict.return_value = {} + + with patch("winml.modelkit.build.common.ensure_pre_quantized_stamped"): + result = _build_onnx_pipeline( + config=config, + onnx_path=onnx_path, + output_dir=output_dir, + rebuild=True, + ep="cpu", + device="cpu", + extra_kwargs={"skip_optimize": True}, + ) + + assert result == [] + mock_optimize.assert_not_called() + assert mock_quantize.call_args.kwargs["current_path"] == copied_input + expected_final_source = quantized_path if quant_mode else copied_input + assert mock_copy.call_args_list[-1].args == ( + expected_final_source, + output_dir / "model.onnx", + ) + + @pytest.mark.parametrize("pipeline", ["hf", "onnx"]) + @patch("winml.modelkit.commands.build._run_compile_stage") + @patch("winml.modelkit.commands.build._run_quantize_stage") + @patch("winml.modelkit.commands.build._run_optimize_stage") + def test_default_pipeline_still_optimizes( + self, + mock_optimize: MagicMock, + mock_quantize: MagicMock, + mock_compile: MagicMock, + pipeline: str, + tmp_path: Path, + ) -> None: + from winml.modelkit.commands.build import _build_hf_pipeline, _build_onnx_pipeline + + output_dir = tmp_path / "out" + optimized_name = "optimized.onnx" if pipeline == "hf" else "input_optimized.onnx" + optimized_path = output_dir / optimized_name + mock_optimize.return_value = (optimized_path, 0.1) + mock_quantize.side_effect = lambda **kwargs: kwargs["current_path"] + mock_compile.side_effect = lambda **kwargs: kwargs["current_path"] + config = MagicMock() + config.skip_optimize = False + config.quant = None + config.to_dict.return_value = {} + + if pipeline == "hf": + config.loader.model_type = "xlm-roberta" + config.loader.task = "sentence-similarity" + config.loader.model_class = "AutoModel" + config.export = MagicMock() + with ( + patch("winml.modelkit.build.hf._load_model"), + patch("winml.modelkit.export.export_onnx"), + patch("winml.modelkit.utils.console.StageLive") as mock_stage_live, + patch("winml.modelkit.onnx.copy_onnx_model"), + ): + mock_stage_live.return_value.__enter__ = MagicMock(return_value=MagicMock()) + mock_stage_live.return_value.__exit__ = MagicMock(return_value=False) + _build_hf_pipeline( + config=config, + model_id="example/model", + output_dir=output_dir, + rebuild=True, + cache_key=None, + ep="cpu", + device="cpu", + extra_kwargs={}, + ) + else: + onnx_path = tmp_path / "input.onnx" + onnx_path.write_bytes(b"fake-onnx-data") + with ( + patch("winml.modelkit.build.common.ensure_pre_quantized_stamped"), + patch("winml.modelkit.onnx.copy_onnx_model"), + ): + _build_onnx_pipeline( + config=config, + onnx_path=onnx_path, + output_dir=output_dir, + rebuild=True, + ep="cpu", + device="cpu", + extra_kwargs={}, + ) + + mock_optimize.assert_called_once() + + class TestBuildEpResolution: """--ep forwarding into config generation + the compile EP-availability gate.""" @@ -2330,11 +2509,11 @@ def test_quantize_stage_runs_when_raw_model_only_skips_optimize( assert timings and timings[0][0] == "Quantize" mock_quantize.assert_called_once() - def test_pre_quantized_stamp_runs_before_optimize(self, tmp_path: Path) -> None: - """_build_onnx_pipeline must stamp config before optimize/quantize stages. + def test_pre_quantized_stamp_runs_before_stage_dispatch(self, tmp_path: Path) -> None: + """_build_onnx_pipeline must stamp config before optimize/quantize dispatch. This ensures pre-quantized ONNX inputs can set skip_optimize and clear - quant before stage dispatch, preventing optimize/quantize double-work. + quant before stage dispatch, bypassing optimize/quantize double-work. """ from winml.modelkit.commands.build import _build_onnx_pipeline @@ -2382,7 +2561,7 @@ def _stamp(cfg: MagicMock, _path: Path) -> None: assert result is not None mock_stamp.assert_called_once() - assert mock_opt.call_args.kwargs["skip_optimize"] is True + mock_opt.assert_not_called() assert mock_quant.call_args.kwargs["config"].quant is None From 3cd22269dfbfaac8baef63cb221c44c9e803af91 Mon Sep 17 00:00:00 2001 From: Shiyi Zheng Date: Fri, 28 Aug 2026 13:57:37 +0800 Subject: [PATCH 2/2] test(build): prove compile dataflow with no optimize --- tests/unit/commands/test_build.py | 20 ++++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) diff --git a/tests/unit/commands/test_build.py b/tests/unit/commands/test_build.py index b92f1289e..36fbfe392 100644 --- a/tests/unit/commands/test_build.py +++ b/tests/unit/commands/test_build.py @@ -2051,8 +2051,9 @@ def test_hf_skip_optimize_preserves_quantize( output_dir = tmp_path / "out" export_path = output_dir / "export.onnx" quantized_path = output_dir / "quantized.onnx" + compiled_path = output_dir / "compiled.onnx" mock_quantize.return_value = quantized_path if quant_mode else export_path - mock_compile.side_effect = lambda **kwargs: kwargs["current_path"] + mock_compile.return_value = compiled_path config = MagicMock() config.skip_optimize = False @@ -2061,6 +2062,7 @@ def test_hf_skip_optimize_preserves_quantize( config.loader.model_class = "AutoModel" config.export = MagicMock() config.quant = None if quant_mode is None else MagicMock(mode=quant_mode) + config.compile = MagicMock() if config.quant is not None: config.quant.model_type = None config.to_dict.return_value = {} @@ -2080,8 +2082,10 @@ def test_hf_skip_optimize_preserves_quantize( assert [name for name, _ in result] == ["Export"] mock_optimize.assert_not_called() assert mock_quantize.call_args.kwargs["current_path"] == export_path - expected_final_source = quantized_path if quant_mode else export_path - mock_copy.assert_called_once_with(expected_final_source, output_dir / "model.onnx") + expected_compile_source = quantized_path if quant_mode else export_path + mock_compile.assert_called_once() + assert mock_compile.call_args.kwargs["current_path"] == expected_compile_source + mock_copy.assert_called_once_with(compiled_path, output_dir / "model.onnx") @pytest.mark.parametrize("quant_mode", [None, "fp16"]) @patch("winml.modelkit.onnx.copy_onnx_model") @@ -2104,12 +2108,14 @@ def test_onnx_skip_optimize_preserves_quantize( output_dir = tmp_path / "out" copied_input = output_dir / onnx_path.name quantized_path = output_dir / "input_quantized.onnx" + compiled_path = output_dir / "input_compiled.onnx" mock_quantize.return_value = quantized_path if quant_mode else copied_input - mock_compile.side_effect = lambda **kwargs: kwargs["current_path"] + mock_compile.return_value = compiled_path config = MagicMock() config.skip_optimize = False config.quant = None if quant_mode is None else MagicMock(mode=quant_mode) + config.compile = MagicMock() config.to_dict.return_value = {} with patch("winml.modelkit.build.common.ensure_pre_quantized_stamped"): @@ -2126,9 +2132,11 @@ def test_onnx_skip_optimize_preserves_quantize( assert result == [] mock_optimize.assert_not_called() assert mock_quantize.call_args.kwargs["current_path"] == copied_input - expected_final_source = quantized_path if quant_mode else copied_input + expected_compile_source = quantized_path if quant_mode else copied_input + mock_compile.assert_called_once() + assert mock_compile.call_args.kwargs["current_path"] == expected_compile_source assert mock_copy.call_args_list[-1].args == ( - expected_final_source, + compiled_path, output_dir / "model.onnx", )