Skip to content

feat: Add export override options to winml eval command #1141

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
xieofxie wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat: Add export override options to winml eval command #1141

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions docs/commands/eval.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,10 @@ $ winml eval [options]
| `--precision` | | `TEXT` | `auto` | Precision used when building the model from a HuggingFace ID. One of `auto`, `fp32`, `fp16`, `int8`, `int16`, or a mixed `w{x}a{y}` spec (e.g., `w8a16`). `fp16`/`fp32` skip quantization. **Ignored** when `-m` is a pre-built `.onnx` file — the precision is already baked in. |
| `--device` | | choice | `auto` | Target device. Choices: `auto`, `npu`, `gpu`, `cpu`. `auto` selects the best available device. Combined with `--precision`, this drives the build when `-m` is a HuggingFace ID. |
| `--ep` / `--execution-provider` | | `TEXT` | — | Target ONNX Runtime execution provider when finer control than `--device` is needed. Full names (e.g., `QNNExecutionProvider`, `OpenVINOExecutionProvider`, `VitisAIExecutionProvider`) and aliases (`qnn`, `ov`/`openvino`, `vitis`/`vitisai`) are accepted. |
| `--shape-config` | | `PATH` | — | JSON shape overrides used while auto-generating a Hugging Face export config, for example `{"height": 480, "width": 480}`. Applies only when `-m` is a HuggingFace ID that eval builds; **ignored for pre-built `.onnx` inputs**. |
| `--input-specs` | | `PATH` | — | JSON input tensor specs to merge into the Hugging Face export config. Symbolic string dimensions infer dynamic axes. **Ignored for pre-built `.onnx` inputs**. |
| `--export-config` | | `PATH` | — | JSON ONNX export config overrides (opset version, constant folding, etc.) to merge into the Hugging Face export config. **Ignored for pre-built `.onnx` inputs**. |
| `--dynamic-axes` | | `PATH` | — | JSON dynamic axes mapping for Hugging Face ONNX export, for example `{"input_ids": {"0": "batch", "1": "sequence"}}`. **Ignored for pre-built `.onnx` inputs**. |
| `--dataset` | | `TEXT` | task default | HuggingFace dataset path (e.g., `imagenet-1k`, `nyu-mll/glue`). If omitted, a default dataset is selected based on the task. |
| `--dataset-name` | | `TEXT` | — | Dataset configuration name for multi-config datasets. |
| `--dataset-revision` | | `TEXT` | — | Git revision (branch, tag, or commit) of the dataset to load. Use `refs/convert/parquet` for HF datasets that are only served via the parquet mirror. |
Expand Down Expand Up @@ -130,6 +134,7 @@ $ winml eval -m encoder=encoder.onnx -m decoder=decoder.onnx --model-id microsof
- **Some dataset requires Hub credentials for gated datasets.** Some datasets (e.g., `imagenet-1k`) require a HuggingFace account with accepted terms of use. Log in with `huggingface-cli login` before running eval on gated data.
- **`--shuffle` is on by default.** The random 100-sample slice changes between runs unless you pass `--no-shuffle`. Use `--no-shuffle` when comparing two model variants to ensure they see identical samples.
- **`--streaming` skips the local cache.** Streaming mode avoids downloading the full split but prevents random shuffling on large datasets. For reproducible evaluation, download the split once and omit `--streaming`.
- **Export overrides only apply when eval builds from a HuggingFace ID.** `--shape-config`, `--input-specs`, `--export-config`, and `--dynamic-axes` shape the ONNX export that `eval` generates when `-m` is a HuggingFace model ID. When `-m` is a pre-built `.onnx` file, there is no export step, so these flags are ignored and the command prints a warning.
- **Column names vary across datasets.** If the evaluator raises a missing-column error, run `winml eval --schema --task <task>` to inspect the expected schema and use `--column` to remap dataset field names to the expected names.

## See also
Expand Down
78 changes: 78 additions & 0 deletions src/winml/modelkit/commands/eval.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -88,6 +88,32 @@
optional_message="Applied during model build. Ignored for pre-built ONNX inputs."
)
@cli_utils.max_optim_iterations_option(optional_message="Ignored for pre-built ONNX inputs.")
@cli_utils.shape_config_option(
param_name="shape_config_path",
help_text=(
"JSON with shape overrides for auto-generated HuggingFace export configs. "
"Ignored for pre-built ONNX inputs."
),
)
@cli_utils.input_specs_option(

@zhenchaoni zhenchaoni Jul 24, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

An optional suggestion: as we are having more and more options, do you think they can be wrapped into a single "--additional-options" for all those optional command line inputs?

help_text=(
"JSON file with input specifications for HuggingFace export. "
"Ignored for pre-built ONNX inputs."
),
)
@cli_utils.export_config_option(
help_text=(
"ONNX export configuration JSON for HuggingFace model builds "
"(opset_version, do_constant_folding, etc.). Ignored for pre-built ONNX inputs."
),
)
@cli_utils.dynamic_axes_option(
help_text=(
"JSON dynamic axes mapping for HuggingFace ONNX export "
'(e.g., {"input_ids": {"0": "batch", "1": "sequence"}}). '
"Ignored for pre-built ONNX inputs."
),
)
@click.option(
"--samples",
type=int,
Expand Down Expand Up @@ -180,6 +206,10 @@ def eval(
optimize: bool,
analyze: bool,
max_optim_iterations: int | None,
shape_config_path: Path | None,
input_specs: Path | None,
export_config: Path | None,
dynamic_axes: Path | None,
ep: EPNameOrAlias | None,
samples: int,
split: str,
Expand Down Expand Up @@ -241,6 +271,7 @@ def eval(

# ── 2. Resolve in place ──
_resolve_model(cfg, model, model_id)
_apply_export_overrides(cfg, shape_config_path, input_specs, export_config, dynamic_axes)
_resolve_device(cfg)
_resolve_label_mapping(cfg)
_run_dataset_script(cfg, trust_remote_code)
Expand Down Expand Up @@ -368,6 +399,53 @@ def _resolve_model(
cfg.model_id = resolved_id


def _apply_export_overrides(
cfg: WinMLEvaluationConfig,
shape_config_path: Path | None,
input_specs: Path | None,
export_config: Path | None,
dynamic_axes: Path | None,
) -> None:
"""Parse the HuggingFace export CLI overrides onto *cfg* (in place).

``--shape-config``/``--input-specs``/``--export-config``/``--dynamic-axes``
only affect the HF-build path (``model_path is None``), where eval exports
and builds the model. A pre-built ONNX input is consumed as-is (no export
step), so any export/shape overrides are dropped with a warning — mirroring
the ``--precision``/build-flag warnings and winml perf's ONNX path.
Requires ``cfg.model_path`` to already be resolved (call after
:func:`_resolve_model`).
"""
export_flags = (
("--shape-config", shape_config_path),
("--input-specs", input_specs),
("--export-config", export_config),
("--dynamic-axes", dynamic_axes),
)
provided = [flag for flag, value in export_flags if value is not None]
if not provided:
return

if cfg.model_path is not None:
logger.warning(
"%s ignored for pre-built ONNX inputs "
"(no export runs; these apply only when building from a model ID).",
", ".join(provided),
)
return

if shape_config_path is not None:
cfg.shape_config = cli_utils.load_json_object(shape_config_path, "--shape-config")

export_overrides = cli_utils.load_export_overrides(
export_config=export_config,
input_specs=input_specs,
dynamic_axes=dynamic_axes,
)
if export_overrides:
cfg.export_overrides = export_overrides


def _resolve_device(cfg: WinMLEvaluationConfig) -> None:
"""Resolve ``'auto'`` → concrete device string on *cfg* in place."""
if cfg.device and cfg.device.lower() != "auto":
Expand Down
39 changes: 39 additions & 0 deletions src/winml/modelkit/eval/config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -79,6 +79,25 @@ def to_dict(self) -> dict[str, Any]:
return result


def _serialize_export_overrides(overrides: dict[str, Any]) -> dict[str, Any]:
"""Render export overrides JSON-safe for :meth:`WinMLEvaluationConfig.to_dict`.

``--input-specs`` is parsed into ``InputTensorSpec`` objects (via
``load_export_overrides``), which are not JSON serializable; convert them to
plain dicts. Every other override key (``dynamic_axes`` mapping, opset ints,
bool flags) is already JSON-safe and passed through unchanged.
"""
serialized: dict[str, Any] = {}
for key, value in overrides.items():
if key == "input_tensors" and isinstance(value, list):
serialized[key] = [
spec.to_dict() if hasattr(spec, "to_dict") else spec for spec in value
]
else:
serialized[key] = value
return serialized


@dataclass
class WinMLEvaluationConfig:
"""Configuration for model evaluation.
Expand All @@ -92,6 +111,13 @@ class WinMLEvaluationConfig:
device: Target device for inference.
ep: Explicit execution provider (e.g., "qnn", "dml"). Overrides
device-to-provider mapping when provided.
shape_config: Shape overrides for the auto-generated HuggingFace export
config. Only used when building from ``model_id`` (ignored for
pre-built ONNX inputs).
export_overrides: Sparse ONNX export overrides
(``--input-specs``/``--export-config``/``--dynamic-axes``) merged
under the build config's ``export`` section. Only used when building
from ``model_id`` (ignored for pre-built ONNX inputs).
dataset: Dataset configuration.
output_path: Path to write JSON results.
mode: Evaluation mode (see :data:`EvalMode`).
Expand Down Expand Up @@ -122,6 +148,13 @@ class WinMLEvaluationConfig:
optimize: bool = True
analyze: bool = True
max_optim_iterations: int | None = None
# HuggingFace export overrides, applied only when building from model_id
# (ignored for pre-built ONNX inputs). Shared semantics with winml
# build/perf: ``shape_config`` is passed straight to from_pretrained, while
# ``export_overrides`` (--input-specs/--export-config/--dynamic-axes) is
# merged under the build config's ``export`` section.
shape_config: dict | None = None
export_overrides: dict[str, Any] | None = None
dataset: DatasetConfig = field(default_factory=DatasetConfig)
output_path: Path | None = field(default=None, metadata={"cli_name": "output"})
mode: EvalMode = "onnx"
Expand Down Expand Up @@ -153,6 +186,10 @@ def to_dict(self) -> dict:
result["analyze"] = self.analyze
if self.max_optim_iterations is not None:
result["max_optim_iterations"] = self.max_optim_iterations
if self.shape_config:
result["shape_config"] = self.shape_config
if self.export_overrides:
result["export_overrides"] = _serialize_export_overrides(self.export_overrides)
result["dataset"] = self.dataset.to_dict()
if self.output_path is not None:
result["output_path"] = str(self.output_path)
Expand Down Expand Up @@ -190,6 +227,8 @@ def from_dict(cls, data: dict) -> WinMLEvaluationConfig:
optimize=data.get("optimize", True),
analyze=data.get("analyze", True),
max_optim_iterations=data.get("max_optim_iterations"),
shape_config=data.get("shape_config"),
export_overrides=data.get("export_overrides"),
dataset=dataset,
output_path=(Path(data["output_path"]) if data.get("output_path") else None),
mode=data.get("mode", "onnx"),
Expand Down
20 changes: 18 additions & 2 deletions src/winml/modelkit/eval/evaluate.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -253,7 +253,9 @@ def _load_model(config: WinMLEvaluationConfig) -> WinMLPreTrainedModel | WinMLCo

if config.model_path is not None:
# Pre-built ONNX: precision is already baked into the model and is
# ignored here (mirrors winml perf's ONNX path).
# ignored here (mirrors winml perf's ONNX path). Export overrides /
# shape_config are HuggingFace-export concepts and never reach here —
# the CLI warns and drops them for pre-built ONNX inputs.
from transformers import AutoConfig

hf_config = AutoConfig.from_pretrained(config.model_id)
Expand All @@ -270,14 +272,28 @@ def _load_model(config: WinMLEvaluationConfig) -> WinMLPreTrainedModel | WinMLCo
model.config = hf_config
return model

# HuggingFace build path — export overrides (--input-specs/--export-config/
# --dynamic-axes) are merged under the build config's ``export`` section as a
# sparse dict so from_pretrained routes them through merge_export_overrides
# (patching auto-resolved input_tensors by name / re-deriving dynamic_axes).
# Passing a dict rather than a WinMLBuildConfig avoids clobbering the
# auto-resolved export config with default fields. Mirrors winml build/perf.
build_override: Any = quant_override
if config.export_overrides:

@zhenchaoni zhenchaoni Jul 24, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Updated todo list

Here's a PR-ready comment:

Review: input option precedence

I checked the implementation against the intended precedence order:

CLI explicit > config-file explicit > CLI default > config-class default

The four new flags (--shape-config, --input-specs, --export-config, --dynamic-axes) bypass the standard collect_cli_overrides/merge_config machinery (they bind to shape_config_path/input_specs/export_config/dynamic_axes, which aren't WinMLEvaluationConfig field names) and are handled entirely by the new _apply_export_overrides.

✅ What's correct

  • CLI default > config-class default — the options default to None; _apply_export_overrides returns early when no flag is provided, so CLI defaults never clobber anything.
  • config-file explicit > CLI default — the config-file eval section can carry shape_config/export_overrides, and they survive when the CLI omits the flags. They then flow into from_pretrained(config={"export": ...}) → generate_hf_build_config, which routes them through merge_export_overrides (patch input_tensors by name, re-derive dynamic axes), correctly preserving the auto-generated/Optimum defaults tier.
  • CLI explicit > config-file explicit for shape_config — single-value override, fine.

❌ Precedence bug: export_overrides is replaced wholesale, not merged

In _apply_export_overrides:

export_overrides = cli_utils.load_export_overrides(
    export_config=export_config, input_specs=input_specs, dynamic_axes=dynamic_axes,
)
if export_overrides:
    cfg.export_overrides = export_overrides   # wholesale replace

load_export_overrides returns a sparse dict containing only the CLI-provided sub-keys, so assigning it discards any config-file export_overrides sub-keys the CLI didn't set.

Repro: config file sets export_overrides = {"opset_version": 17, "dynamic_axes": {...}}, then run with only --input-specs foo.json. Result: cfg.export_overrides = {"input_tensors": [...]} — the config-file opset_version and dynamic_axes are lost, even though the CLI never set them. Per the intended ordering, those config-file values should win (config-file explicit > CLI default).

This is also a parity gap with build, which deliberately avoids exactly this by calling merge_export_overrides(cfg, export_overrides) on the loaded config "instead of merge_config replacing the list wholesale".

Suggested fix — merge instead of replace:

if export_overrides:
    merged = dict(cfg.export_overrides or {})
    merged.update(export_overrides)
    cfg.export_overrides = merged

Related concerns (not strictly precedence)

  1. Config-file input_tensors will crash. Config-file export_overrides is merged as a raw dict via merge_config(cfg, eval_data) (not through from_dict/load_input_tensor_specs), so its input_tensors stay plain dicts. Downstream _patch_input_tensors accesses patch.name/patch.dtype, which raises AttributeError on a dict — so config-file-supplied input_tensors is effectively unusable.
  2. Pre-built ONNX path: the warn-and-skip only covers CLI flags. If the config file carries shape_config/export_overrides alongside a pre-built .onnx, they stay on cfg (ignored downstream, but with no warning) — minor inconsistency.

Overall the ordering is correct for shape_config and for the config-file/default tiers, but export_overrides needs field-level merging between the CLI and config-file layers to satisfy the intended precedence.

override_dict: dict[str, Any] = {"export": config.export_overrides}
if not config.quant:
override_dict["quant"] = None
build_override = override_dict

return WinMLAutoModel.from_pretrained(
config.model_id,
task=config.task,
device=config.device,
precision=config.precision,
ep=config.ep,
allow_unsupported_nodes=config.allow_unsupported_nodes,
config=quant_override,
config=build_override,
shape_config=config.shape_config,
**pipeline_kwargs,
)

Expand Down
Loading
Loading