From af73c6493d845e75713651827be5dc3cbdea5777 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Tue, 21 Jul 2026 11:42:18 +0800 Subject: [PATCH 01/15] Add dynamo export support and make it the default exporter Flip `winml export` to torch's dynamo-based ONNX exporter by default (--dynamo/--no-dynamo, default True). Add DynamoMetadataTagger that reads torch's native pkg.torch.onnx.class_hierarchy / name_scopes node metadata into the existing winml.hierarchy.* tag contract, so module-hierarchy tags survive without a separate TorchScript trace. - config + CLI default flipped to dynamo=True; stale "not supported" warning removed - DynamoMetadataTagger drop-in tagger (same interface as ONNXNodeTagger) - HTP exporter skips the redundant TorchScript trace under dynamo and sets verbose=False to avoid a UnicodeEncodeError on Windows consoles - TorchScript-path-specific tests pin dynamo=False; docs updated --- docs/commands/export.md | 11 +- docs/concepts/load-and-export.md | 2 +- docs/reference/index.md | 2 +- src/winml/modelkit/commands/export.py | 18 +- src/winml/modelkit/core/onnx_node_tagger.py | 161 ++++++++++++++ src/winml/modelkit/export/config.py | 4 +- src/winml/modelkit/export/htp/exporter.py | 58 +++-- tests/e2e/test_export_e2e.py | 27 +++ .../unit/commands/test_boolean_flag_pairs.py | 2 +- .../commands/test_config_value_priority.py | 2 +- tests/unit/commands/test_export.py | 57 ++++- .../unit/core/test_dynamo_metadata_tagger.py | 201 ++++++++++++++++++ tests/unit/export/test_pytorch_export.py | 9 + 13 files changed, 508 insertions(+), 46 deletions(-) create mode 100644 tests/unit/core/test_dynamo_metadata_tagger.py diff --git a/docs/commands/export.md b/docs/commands/export.md index ff2cf42ec..0e7c80562 100644 --- a/docs/commands/export.md +++ b/docs/commands/export.md @@ -22,7 +22,7 @@ $ winml export [options] | `--output` | `-o` | path | *(required)* | Output ONNX file path (e.g., `model.onnx`). | | `--with-report/--no-with-report` | | flag | `false` | Generate full export reports: Markdown, JSON, and a console tree. | | `--hierarchy/--no-hierarchy` | | flag | `true` | Preserve `hierarchy_tag` metadata in ONNX nodes (use `--no-hierarchy` for a clean ONNX file). | -| `--dynamo/--no-dynamo` | | flag | `false` | Enable PyTorch 2.9+ dynamo export for richer node metadata. (Experimental — currently logs a warning.) | +| `--dynamo/--no-dynamo` | | flag | `true` | Use PyTorch's TorchDynamo ONNX exporter (default) for richer per-node module metadata. Pass `--no-dynamo` for the legacy TorchScript exporter (QNN-safe). | | `--torch-module` | | string | `None` | Comma-separated list of `torch.nn` module types to include in hierarchy (e.g., `LayerNorm,Embedding`). (Experimental — currently logs a warning.) | | `--input-specs` | | path | `None` | JSON file with explicit input tensor specifications. Auto-generated when omitted. | | `--task` | `-t` | string | `None` | Override auto-detected Hugging Face task (e.g., `image-feature-extraction`). | @@ -114,9 +114,12 @@ winml export -m microsoft/resnet-50 -o resnet50_clean.onnx --no-hierarchy - **Dynamic dimensions can reduce QNN optimization coverage.** Static batch and static shapes remain the default because some QNN fusions require them. Use `--dynamic-axes` only when downstream runtime scenarios need variable sizes. -- **`--dynamo` and `--torch-module` are experimental.** Both flags emit a - warning and have no effect in the current release. Do not rely on them in - automated pipelines yet. +- **Dynamo is the default exporter.** `winml export` uses PyTorch's TorchDynamo + ONNX exporter, which records rich per-node module metadata that drives the + hierarchy tags. Pass `--no-dynamo` to select the legacy TorchScript exporter, + which remains the QNN-safe choice for hand exports targeting static shapes. +- **`--torch-module` is experimental.** The flag emits a warning and has no + effect in the current release. Do not rely on it in automated pipelines yet. - **Output directory must be writable.** The command creates parent directories automatically, but will fail with a permission error on read-only paths. - **Model weights are downloaded to the Hugging Face cache.** Set `HF_HOME` or diff --git a/docs/concepts/load-and-export.md b/docs/concepts/load-and-export.md index b195adcd6..1e7467a85 100644 --- a/docs/concepts/load-and-export.md +++ b/docs/concepts/load-and-export.md @@ -16,7 +16,7 @@ Some community models host custom Python code in their repositories. The loader ## Exporting to ONNX -`winml export` converts the loaded model to ONNX. The conversion uses TorchScript tracing by default, which follows actual execution paths and tends to produce compact, inference-oriented graphs. A `--dynamo` flag exists for the PyTorch 2.x dynamo exporter; however, **Note:** the `--dynamo` flag is reserved for the PyTorch 2.x dynamo exporter but is **not yet functional** in the current release — passing it logs a warning and the flag is ignored. +`winml export` converts the loaded model to ONNX. By default it uses PyTorch's TorchDynamo ONNX exporter (`torch.onnx.export(dynamo=True)`), which records rich per-node module metadata that is used to derive the `winml.hierarchy.*` node tags. Pass `--no-dynamo` to fall back to the legacy TorchScript exporter, which follows actual execution paths, tends to produce compact inference-oriented graphs, and remains the QNN-safe choice for static-shape exports. By default the exporter runs an eight-step process that includes hierarchy tracing and tag injection. The result is an ONNX file enriched with structural metadata that powers downstream features such as per-module benchmarking, inspector views, and optimizer scoping. diff --git a/docs/reference/index.md b/docs/reference/index.md index 086f12cf0..568ad9ede 100644 --- a/docs/reference/index.md +++ b/docs/reference/index.md @@ -54,7 +54,7 @@ stages based on the target device and precision. | `export_params` | `bool` | `true` | Include model parameters in ONNX. | | `do_constant_folding` | `bool` | `true` | Fold constants during export. | | `verbose` | `bool` | `false` | Verbose export logging. | -| `dynamo` | `bool` | `false` | Use PyTorch 2.x Dynamo exporter. | +| `dynamo` | `bool` | `true` | Use PyTorch's TorchDynamo ONNX exporter (default); set `false` for the legacy TorchScript exporter. | | `enable_hierarchy_tags` | `bool` | `true` | Add module hierarchy tags to ONNX nodes. | | `clean_onnx` | `bool` | `false` | Strip hierarchy tags after export. | | `hierarchy_tag_format` | `"full" \| "module_only"` | `"full"` | Tag detail level. | diff --git a/src/winml/modelkit/commands/export.py b/src/winml/modelkit/commands/export.py index 15e04f7b2..ad45dae82 100644 --- a/src/winml/modelkit/commands/export.py +++ b/src/winml/modelkit/commands/export.py @@ -116,9 +116,10 @@ def _warn_partial_composite(completed: list[Path]) -> None: @click.option( "--dynamo/--no-dynamo", "dynamo", - default=False, + default=True, show_default=True, - help="Enable PyTorch 2.9+ dynamo export for rich node metadata", + help="Use PyTorch's TorchDynamo ONNX exporter (default). " + "Pass --no-dynamo for the legacy TorchScript exporter (QNN-safe).", ) @click.option( "--torch-module", @@ -206,8 +207,8 @@ def export( # Clean ONNX output (no hierarchy metadata, for optimization) winml export -m prajjwal1/bert-tiny -o model.onnx --clean-onnx - # Use PyTorch dynamo export (for rich node metadata) - winml export -m prajjwal1/bert-tiny -o model.onnx --dynamo + # Use the legacy TorchScript exporter (dynamo is the default) + winml export -m prajjwal1/bert-tiny -o model.onnx --no-dynamo # Include torch.nn modules in hierarchy winml export -m prajjwal1/bert-tiny -o model.onnx --torch-module LayerNorm,Embedding @@ -323,15 +324,6 @@ def export( "TODO: Add torch_module support to export_onnx() and WinMLExportConfig.", torch_module, ) - if dynamo: - console.print( - "[yellow]Warning:[/yellow] --dynamo is not yet supported in export_onnx(). " - "export_onnx() defaults to dynamo=False for QNN compatibility." - ) - logger.warning( - "dynamo=True is not supported by export_onnx(). " - "TODO: Add dynamo support to WinMLExportConfig if needed." - ) def _run_component_export(component_task: str | None, out_path: Path) -> None: """Resolve I/O, build config, load, and export one (model, task) to ``out_path``.""" diff --git a/src/winml/modelkit/core/onnx_node_tagger.py b/src/winml/modelkit/core/onnx_node_tagger.py index 98cfbbc75..faa47a53e 100644 --- a/src/winml/modelkit/core/onnx_node_tagger.py +++ b/src/winml/modelkit/core/onnx_node_tagger.py @@ -18,6 +18,7 @@ from __future__ import annotations +import ast from collections import defaultdict from typing import TYPE_CHECKING, ClassVar @@ -311,3 +312,163 @@ def create_node_tagger_from_hierarchy( Configured ONNXNodeTagger instance """ return ONNXNodeTagger(hierarchy_data, enable_operation_fallback) + + +class DynamoMetadataTagger: + """ONNX node tagger that reads PyTorch dynamo-native module metadata. + + The dynamo ONNX exporter (``torch.onnx.export(dynamo=True)``) records the + originating module hierarchy directly on each node's ``metadata_props``: + + - ``pkg.torch.onnx.name_scopes``: ``repr()`` of a list of cumulative module + paths, e.g. ``['', 'blocks.0', 'blocks.0.lin', 'linear']``. The first + entry is the root (``""``); the last entry is the node's own name. + - ``pkg.torch.onnx.class_hierarchy``: ``repr()`` of a parallel list of + fully-qualified class names, e.g. + ``['pkg.Net', 'pkg.Blk', 'torch.nn.modules.linear.Linear', + 'aten.linear.default']``. It is aligned element-for-element with + ``name_scopes``; the last entry is the aten op target. + + This tagger converts that metadata into the same ``/Root/Child.N/Leaf`` tag + contract produced by :class:`ONNXNodeTagger`, so downstream consumers + (inspector, per-module benchmarking, optimizer scoping) are unchanged. + + NO HARDCODED LOGIC: only generic dynamo metadata keys are read; no model, + architecture, or operator names are special-cased. + """ + + NAME_SCOPES_KEY: ClassVar[str] = "pkg.torch.onnx.name_scopes" + CLASS_HIERARCHY_KEY: ClassVar[str] = "pkg.torch.onnx.class_hierarchy" + UNKNOWN_ROOT_TAG: ClassVar[str] = "/UnknownModel" + + def __init__(self) -> None: + """Initialize the dynamo metadata tagger.""" + # Derived from node metadata during tag_all_nodes; kept for parity with + # ONNXNodeTagger and as the guaranteed non-empty fallback tag. + self.model_root_tag = self.UNKNOWN_ROOT_TAG + + @staticmethod + def _node_metadata(node: onnx.NodeProto) -> dict[str, str]: + """Collect a node's metadata_props into a plain dict.""" + return {prop.key: prop.value for prop in node.metadata_props} + + @staticmethod + def _parse_list(raw: str | None) -> list[str]: + """Parse a ``repr(list)`` metadata value into a list of strings. + + Returns an empty list when the value is missing or malformed so callers + degrade to the root fallback instead of raising. + """ + if not raw: + return [] + try: + value = ast.literal_eval(raw) + except (ValueError, SyntaxError): + return [] + if not isinstance(value, (list, tuple)): + return [] + return [str(item) for item in value] + + @staticmethod + def _short_class(fq_class: str) -> str: + """Reduce a fully-qualified class name to its final component. + + ``torch.nn.modules.linear.Linear`` -> ``Linear``. + """ + return fq_class.rsplit(".", 1)[-1] if fq_class else "" + + def _tag_from_metadata(self, class_hierarchy: list[str], name_scopes: list[str]) -> str | None: + """Build a ``/Root/Child.N/Leaf`` tag from aligned dynamo metadata lists. + + The last element of each list is the node's own op target / node name + (not a module level) and is dropped. For every remaining level the + class supplies the segment name; when the cumulative scope's final + component is a pure index (a ``ModuleList``/``Sequential`` child) it is + folded in as ``Class.N`` to mirror the TorchScript tagger's indexed + module tags. + """ + classes = class_hierarchy[:-1] + scopes = name_scopes[:-1] + + segments: list[str] = [] + for cls, scope in zip(classes, scopes, strict=False): + short = self._short_class(cls) + if not short: + continue + local = scope.rsplit(".", 1)[-1] if scope else "" + if local.isdigit(): + segments.append(f"{short}.{local}") + else: + segments.append(short) + + if not segments: + return None + return "/" + "/".join(segments) + + def _extract_model_root_tag(self, onnx_model: onnx.ModelProto) -> str: + """Derive the model root tag from the first usable class hierarchy.""" + for node in onnx_model.graph.node: + metadata = self._node_metadata(node) + classes = self._parse_list(metadata.get(self.CLASS_HIERARCHY_KEY)) + if len(classes) >= 2: + root = self._short_class(classes[0]) + if root: + return f"/{root}" + return self.UNKNOWN_ROOT_TAG + + def tag_all_nodes(self, onnx_model: onnx.ModelProto) -> dict[str, str]: + """Tag all ONNX nodes from dynamo metadata. + + Returns: + Dictionary mapping node names to hierarchy tags (NO EMPTY TAGS). + """ + self.model_root_tag = self._extract_model_root_tag(onnx_model) + + tagged_nodes: dict[str, str] = {} + for node in onnx_model.graph.node: + node_name = node.name or f"{node.op_type}_{id(node)}" + metadata = self._node_metadata(node) + classes = self._parse_list(metadata.get(self.CLASS_HIERARCHY_KEY)) + scopes = self._parse_list(metadata.get(self.NAME_SCOPES_KEY)) + tag = self._tag_from_metadata(classes, scopes) if classes and scopes else None + tagged_nodes[node_name] = tag or self.model_root_tag + + # Verify no empty tags (same contract as ONNXNodeTagger). + for node_name, tag in tagged_nodes.items(): + assert tag, f"Empty tag generated for node {node_name}" + assert tag.strip(), f"Whitespace-only tag generated for node {node_name}" + assert tag.startswith("/"), f"Invalid tag format: {tag}" + + return tagged_nodes + + def get_tagging_statistics(self, onnx_model: onnx.ModelProto) -> dict[str, int]: + """Get statistics about the tagging process. + + Keys mirror :meth:`ONNXNodeTagger.get_tagging_statistics` so the shared + console/report/metadata writers render consistent output regardless of + which tagger produced the tags. + """ + total_nodes = len(onnx_model.graph.node) + metadata_matches = 0 + unique_tags: set[str] = set() + + for node in onnx_model.graph.node: + metadata = self._node_metadata(node) + classes = self._parse_list(metadata.get(self.CLASS_HIERARCHY_KEY)) + scopes = self._parse_list(metadata.get(self.NAME_SCOPES_KEY)) + tag = self._tag_from_metadata(classes, scopes) if classes and scopes else None + if tag: + metadata_matches += 1 + unique_tags.add(tag) + + root_fallbacks = total_nodes - metadata_matches + return { + "total_nodes": total_nodes, + "root_nodes": root_fallbacks, + "scoped_nodes": metadata_matches, + "unique_scopes": len(unique_tags), + "direct_matches": metadata_matches, + "parent_matches": 0, + "operation_matches": 0, + "root_fallbacks": root_fallbacks, + } diff --git a/src/winml/modelkit/export/config.py b/src/winml/modelkit/export/config.py index 34ac3ba8c..f73b73caa 100644 --- a/src/winml/modelkit/export/config.py +++ b/src/winml/modelkit/export/config.py @@ -177,7 +177,7 @@ class WinMLExportConfig: export_params: bool = True do_constant_folding: bool = True verbose: bool = False - dynamo: bool = False # Use TorchScript exporter by default (dynamo=True for PyTorch 2.x) + dynamo: bool = True # TorchDynamo exporter by default; False = legacy TorchScript path # Phase 2: Hierarchy Preservation Options enable_hierarchy_tags: bool = True # Enable HTP hierarchy tagging by default @@ -403,7 +403,7 @@ def from_dict(cls, data: dict[str, Any]) -> WinMLExportConfig: export_params=data.get("export_params", True), do_constant_folding=data.get("do_constant_folding", True), verbose=data.get("verbose", False), - dynamo=data.get("dynamo", False), + dynamo=data.get("dynamo", True), enable_hierarchy_tags=data.get("enable_hierarchy_tags", True), clean_onnx=data.get("clean_onnx", False), hierarchy_tag_format=data.get("hierarchy_tag_format", "full"), diff --git a/src/winml/modelkit/export/htp/exporter.py b/src/winml/modelkit/export/htp/exporter.py index c61a162d3..211e394f2 100644 --- a/src/winml/modelkit/export/htp/exporter.py +++ b/src/winml/modelkit/export/htp/exporter.py @@ -30,7 +30,11 @@ import torch.nn as nn from rich.console import Console -from ...core.onnx_node_tagger import ONNXNodeTagger, create_node_tagger_from_hierarchy +from ...core.onnx_node_tagger import ( + DynamoMetadataTagger, + ONNXNodeTagger, + create_node_tagger_from_hierarchy, +) from ...core.onnx_utils import infer_output_names from .base_writer import ExportStep from .hierarchy import TracingHierarchyBuilder @@ -141,10 +145,12 @@ def __init__( # Core components self._hierarchy_builder: TracingHierarchyBuilder | None = None - self._node_tagger: ONNXNodeTagger | None = None + self._node_tagger: ONNXNodeTagger | DynamoMetadataTagger | None = None self._hierarchy_data: dict[str, Any] = {} self._tagged_nodes: dict[str, str] = {} self._tagging_stats: dict[str, Any] = {} + # Whether to source hierarchy from dynamo node metadata (set in export()). + self._use_dynamo_hierarchy: bool = False # Export statistics self._export_stats = HTPConfig.DEFAULT_EXPORT_STATS.copy() @@ -183,6 +189,11 @@ def export( """ start_time = time.time() + # Dynamo records the module hierarchy natively on each ONNX node's + # metadata_props, so select the hierarchy source (and whether to run the + # TorchScript trace) from the exporter choice. + self._use_dynamo_hierarchy = bool(export_config.dynamo) + # Initialize export monitor self._monitor = HTPExportMonitor( output_path=output_path, @@ -229,13 +240,21 @@ def export( monitor.update(ExportStep.INPUT_GEN, **input_gen_data) # Step 3: Hierarchy Building - # Trace under the Optimum patcher so models that inject constant - # forward arguments at export time (e.g. ViTPose MoE's dataset_index) - # are traced with the same inputs they are exported with. The export - # in Step 4 re-enters the patcher; the contexts are sequential, not - # nested. - with self._get_optimum_patcher(model, task): - self._trace_model_hierarchy(model, inputs) + # + # Under dynamo the exporter records the module hierarchy natively on + # each ONNX node's metadata_props, so the TorchScript trace is + # redundant (and can fail for models that only export via dynamo). + # The hierarchy is recovered post-export from that metadata in Step 5. + if self._use_dynamo_hierarchy: + self._hierarchy_data = {} + else: + # Trace under the Optimum patcher so models that inject constant + # forward arguments at export time (e.g. ViTPose MoE's dataset_index) + # are traced with the same inputs they are exported with. The export + # in Step 4 re-enters the patcher; the contexts are sequential, not + # nested. + with self._get_optimum_patcher(model, task): + self._trace_model_hierarchy(model, inputs) execution_steps = ( self._hierarchy_builder.get_execution_summary().get("execution_steps", 0) @@ -477,6 +496,12 @@ def _convert_model_to_onnx( } # Always explicitly set dynamo — PyTorch 2.10+ defaults to True onnx_kwargs["dynamo"] = bool(export_config.dynamo) + if onnx_kwargs["dynamo"]: + # torch's dynamo exporter prints capture progress with emoji glyphs + # (e.g. ✅) that raise UnicodeEncodeError on non-UTF-8 consoles + # (Windows cp1252). winml drives its own progress UI, so silence + # torch's verbose printing to keep exports robust across terminals. + onnx_kwargs["verbose"] = False if input_names: onnx_kwargs["input_names"] = input_names if output_names: @@ -575,10 +600,17 @@ def _update_tag_stats(self, total_nodes: int) -> None: ) def _initialize_node_tagger(self, enable_operation_fallback: bool) -> None: - """Create node tagger internally.""" - self._node_tagger = create_node_tagger_from_hierarchy( - self._hierarchy_data, enable_operation_fallback=enable_operation_fallback - ) + """Create node tagger internally. + + Under dynamo the hierarchy comes from each node's dynamo metadata; the + legacy path builds a tagger from the TorchScript trace hierarchy. + """ + if self._use_dynamo_hierarchy: + self._node_tagger = DynamoMetadataTagger() + else: + self._node_tagger = create_node_tagger_from_hierarchy( + self._hierarchy_data, enable_operation_fallback=enable_operation_fallback + ) def _apply_hierarchy_tags(self, onnx_model: onnx.ModelProto) -> None: """Tag nodes internally.""" diff --git a/tests/e2e/test_export_e2e.py b/tests/e2e/test_export_e2e.py index 5f5b02cab..68757ffac 100644 --- a/tests/e2e/test_export_e2e.py +++ b/tests/e2e/test_export_e2e.py @@ -247,6 +247,22 @@ def test_minimal_resnet50(self, tmp_path: Path): _assert_all_nodes_have(model, "winml.hierarchy.tag") _assert_all_nodes_have(model, "winml.hierarchy.depth") + # Dynamo is the default exporter, so nodes carry torch's native module + # metadata and the hierarchy tags are derived from it — not collapsed to + # the model root (the pre-fix regression). Assert both facts generically, + # without referencing any architecture-specific names. + _assert_some_node_has(model, "pkg.torch.onnx.class_hierarchy") + depths = [ + int(prop.value) + for node in model.graph.node + for prop in node.metadata_props + if prop.key == "winml.hierarchy.depth" + ] + assert depths and max(depths) >= 2, ( + "expected dynamo-derived hierarchy tags deeper than the model root; " + f"got max depth {max(depths) if depths else 0}" + ) + class TestExportDinoV2: MODEL = "facebook/dinov2-base" @@ -334,6 +350,17 @@ def test_dynamo(self, tmp_path: Path): # Only rewritten nodes carry this key; "at least one" is the correct check. _assert_some_node_has(model, "pkg.onnxscript.rewriter.rule_name") + def test_no_dynamo_uses_torchscript_hierarchy(self, tmp_path: Path): + # Dynamo is the default, so --no-dynamo selects the legacy TorchScript + # exporter. It must still populate hierarchy tags (derived from node + # names via the module trace) and, unlike dynamo, emit no torch-native + # class_hierarchy metadata — proving the two paths stay distinct. + onnx_path = tmp_path / "model.onnx" + model = _assert_succeeds(_happy_args(onnx_path, "--no-dynamo"), onnx_path) + _assert_all_nodes_have(model, "winml.hierarchy.tag") + _assert_all_nodes_have(model, "winml.hierarchy.depth") + _assert_no_node_has(model, "pkg.torch.onnx.class_hierarchy") + def test_torch_module_warning(self, tmp_path: Path): # --torch-module is currently a no-op; the command must still succeed # but emit a warning identifying the option. diff --git a/tests/unit/commands/test_boolean_flag_pairs.py b/tests/unit/commands/test_boolean_flag_pairs.py index d428fa70e..a651e52ba 100644 --- a/tests/unit/commands/test_boolean_flag_pairs.py +++ b/tests/unit/commands/test_boolean_flag_pairs.py @@ -150,7 +150,6 @@ class TestDefaultValues: (compile, "embed", False), (eval_cmd, "streaming", False), (export, "with_report", False), - (export, "dynamo", False), (inspect, "hierarchy", False), (perf, "rebuild", False), (perf, "ignore_cache", False), @@ -172,6 +171,7 @@ class TestDefaultValues: (eval_cmd, "optimize", True), (eval_cmd, "analyze", True), (export, "hierarchy", True), + (export, "dynamo", True), ], ) def test_default_value(self, command, param_name: str, expected_default) -> None: diff --git a/tests/unit/commands/test_config_value_priority.py b/tests/unit/commands/test_config_value_priority.py index 1a7d8e1d6..51088d167 100644 --- a/tests/unit/commands/test_config_value_priority.py +++ b/tests/unit/commands/test_config_value_priority.py @@ -855,7 +855,7 @@ def test_empty_export_section_keeps_dynamo_cli_default(self, tmp_path: Path) -> Guards the same fix as ``enable_hierarchy_tags`` above. """ eff = _run_export([], {"export": {}}, tmp_path) - assert eff["dynamo"] is False # CLI default + assert eff["dynamo"] is True # CLI default (dynamo is the default exporter) class TestQuantizePriority: diff --git a/tests/unit/commands/test_export.py b/tests/unit/commands/test_export.py index 687e61cd8..ad592f9fd 100644 --- a/tests/unit/commands/test_export.py +++ b/tests/unit/commands/test_export.py @@ -612,30 +612,67 @@ def test_export_warns_on_torch_module( assert "not yet supported" in result.output or "Warning" in result.output - def test_export_warns_on_dynamo( + def test_export_dynamo_enabled_by_default( self, runner: CliRunner, mock_export_onnx: MagicMock, mock_load_hf_model: MagicMock, tmp_path: Path, ) -> None: - """Test --dynamo shows warning (not yet supported).""" + """Dynamo is the default exporter when neither flag is passed.""" from winml.modelkit.commands.export import export output_path = tmp_path / "model.onnx" result = runner.invoke( export, - [ - "--model", - "test-model", - "--output", - str(output_path), - "--dynamo", - ], + ["--model", "test-model", "--output", str(output_path)], obj={"debug": False}, ) - assert "not yet supported" in result.output or "Warning" in result.output + config = mock_export_onnx.call_args.kwargs["export_config"] + assert config.dynamo is True + assert "not yet supported" not in result.output + + def test_export_dynamo_flag_enables_dynamo( + self, + runner: CliRunner, + mock_export_onnx: MagicMock, + mock_load_hf_model: MagicMock, + tmp_path: Path, + ) -> None: + """--dynamo reaches WinMLExportConfig.dynamo and no longer warns.""" + from winml.modelkit.commands.export import export + + output_path = tmp_path / "model.onnx" + result = runner.invoke( + export, + ["--model", "test-model", "--output", str(output_path), "--dynamo"], + obj={"debug": False}, + ) + + config = mock_export_onnx.call_args.kwargs["export_config"] + assert config.dynamo is True + assert "not yet supported" not in result.output + + def test_export_no_dynamo_selects_torchscript( + self, + runner: CliRunner, + mock_export_onnx: MagicMock, + mock_load_hf_model: MagicMock, + tmp_path: Path, + ) -> None: + """--no-dynamo selects the legacy TorchScript exporter.""" + from winml.modelkit.commands.export import export + + output_path = tmp_path / "model.onnx" + runner.invoke( + export, + ["--model", "test-model", "--output", str(output_path), "--no-dynamo"], + obj={"debug": False}, + ) + + config = mock_export_onnx.call_args.kwargs["export_config"] + assert config.dynamo is False class TestExportErrorHandling: diff --git a/tests/unit/core/test_dynamo_metadata_tagger.py b/tests/unit/core/test_dynamo_metadata_tagger.py new file mode 100644 index 000000000..45fcccbfe --- /dev/null +++ b/tests/unit/core/test_dynamo_metadata_tagger.py @@ -0,0 +1,201 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Tests for DynamoMetadataTagger. + +The dynamo ONNX exporter (``torch.onnx.export(dynamo=True)``) records the +originating module hierarchy on each node's ``metadata_props`` as ``repr()`` +of two aligned Python lists: + +- ``pkg.torch.onnx.name_scopes``: cumulative module paths, first entry ``""`` + (root), last entry the node's own name. +- ``pkg.torch.onnx.class_hierarchy``: parallel fully-qualified class names, + last entry the aten op target. + +These tests build synthetic ONNX nodes carrying that exact metadata format and +assert the tagger reproduces the ``/Root/Child.N/Leaf`` hierarchy-tag contract. +""" + +from __future__ import annotations + +import onnx +from onnx import helper + +from winml.modelkit.core.onnx_node_tagger import DynamoMetadataTagger + + +NAME_SCOPES_KEY = "pkg.torch.onnx.name_scopes" +CLASS_HIERARCHY_KEY = "pkg.torch.onnx.class_hierarchy" + + +def _make_node( + op_type: str, + name: str, + *, + name_scopes: list[str] | None = None, + class_hierarchy: list[str] | None = None, + raw_name_scopes: str | None = None, + raw_class_hierarchy: str | None = None, +) -> onnx.NodeProto: + """Build an ONNX node with dynamo-style metadata_props. + + ``name_scopes``/``class_hierarchy`` are serialized with ``repr`` exactly as + torch does. ``raw_*`` overrides let a test inject malformed strings. + """ + node = helper.make_node(op_type, inputs=["x"], outputs=["y"], name=name) + if raw_name_scopes is not None: + node.metadata_props.append( + onnx.StringStringEntryProto(key=NAME_SCOPES_KEY, value=raw_name_scopes) + ) + elif name_scopes is not None: + node.metadata_props.append( + onnx.StringStringEntryProto(key=NAME_SCOPES_KEY, value=repr(name_scopes)) + ) + if raw_class_hierarchy is not None: + node.metadata_props.append( + onnx.StringStringEntryProto(key=CLASS_HIERARCHY_KEY, value=raw_class_hierarchy) + ) + elif class_hierarchy is not None: + node.metadata_props.append( + onnx.StringStringEntryProto(key=CLASS_HIERARCHY_KEY, value=repr(class_hierarchy)) + ) + return node + + +def _make_model(nodes: list[onnx.NodeProto]) -> onnx.ModelProto: + """Wrap nodes in a minimal ModelProto (tagger only iterates graph.node).""" + graph = helper.make_graph(nodes, "g", inputs=[], outputs=[]) + return helper.make_model(graph) + + +def _resnet_like_nodes() -> list[onnx.NodeProto]: + """Two nodes mirroring the empirically verified dynamo output shape.""" + return [ + _make_node( + "Gemm", + "node_gemm", + name_scopes=["", "blocks.0", "blocks.0.lin", "linear"], + class_hierarchy=[ + "pkg.Net", + "pkg.Blk", + "torch.nn.modules.linear.Linear", + "aten.linear.default", + ], + ), + _make_node( + "Relu", + "node_relu", + name_scopes=["", "blocks.0", "blocks.0.act", "relu"], + class_hierarchy=[ + "pkg.Net", + "pkg.Blk", + "torch.nn.modules.activation.ReLU", + "aten.relu.default", + ], + ), + ] + + +class TestDynamoMetadataTagger: + """Tag generation from dynamo node metadata.""" + + def test_tags_indexed_and_named_modules(self) -> None: + model = _make_model(_resnet_like_nodes()) + tags = DynamoMetadataTagger().tag_all_nodes(model) + + # Digit scope component ("blocks.0") folds into "Blk.0"; named module + # ("blocks.0.lin" -> Linear) uses the class short-name only. + assert tags["node_gemm"] == "/Net/Blk.0/Linear" + assert tags["node_relu"] == "/Net/Blk.0/ReLU" + + def test_model_root_tag_from_first_class(self) -> None: + tagger = DynamoMetadataTagger() + tagger.tag_all_nodes(_make_model(_resnet_like_nodes())) + assert tagger.model_root_tag == "/Net" + + def test_tags_are_never_empty(self) -> None: + model = _make_model(_resnet_like_nodes()) + tags = DynamoMetadataTagger().tag_all_nodes(model) + for tag in tags.values(): + assert tag + assert tag.strip() + assert tag.startswith("/") + + def test_missing_metadata_falls_back_to_root(self) -> None: + nodes = _resnet_like_nodes() + # A node with no dynamo metadata at all (e.g. a fused/optimized node). + nodes.append(_make_node("Add", "node_bare")) + model = _make_model(nodes) + + tags = DynamoMetadataTagger().tag_all_nodes(model) + # Falls back to the derived model root, never empty. + assert tags["node_bare"] == "/Net" + + def test_malformed_metadata_does_not_raise(self) -> None: + nodes = _resnet_like_nodes() + nodes.append( + _make_node( + "Add", + "node_bad", + raw_name_scopes="not-a-list", + raw_class_hierarchy="[unclosed", + ) + ) + model = _make_model(nodes) + + tags = DynamoMetadataTagger().tag_all_nodes(model) + assert tags["node_bad"] == "/Net" + + def test_no_metadata_anywhere_uses_unknown_root(self) -> None: + model = _make_model([_make_node("Add", "n0"), _make_node("Mul", "n1")]) + tagger = DynamoMetadataTagger() + tags = tagger.tag_all_nodes(model) + + assert tagger.model_root_tag == "/UnknownModel" + assert tags["n0"] == "/UnknownModel" + assert tags["n1"] == "/UnknownModel" + + def test_root_only_node_tags_to_root(self) -> None: + # A node whose only module level is the root model itself. + node = _make_node( + "Gemm", + "node_root", + name_scopes=["", "linear"], + class_hierarchy=["pkg.Net", "aten.linear.default"], + ) + tags = DynamoMetadataTagger().tag_all_nodes(_make_model([node])) + assert tags["node_root"] == "/Net" + + def test_unnamed_node_key_uses_optype_fallback(self) -> None: + node = _make_node( + "Gemm", + "", # unnamed + name_scopes=["", "blocks.0", "blocks.0.lin", "linear"], + class_hierarchy=[ + "pkg.Net", + "pkg.Blk", + "torch.nn.modules.linear.Linear", + "aten.linear.default", + ], + ) + tags = DynamoMetadataTagger().tag_all_nodes(_make_model([node])) + # Key mirrors exporter convention: "_". + (only_key,) = tags.keys() + assert only_key.startswith("Gemm_") + assert tags[only_key] == "/Net/Blk.0/Linear" + + def test_get_tagging_statistics(self) -> None: + nodes = _resnet_like_nodes() + nodes.append(_make_node("Add", "node_bare")) # root fallback + model = _make_model(nodes) + + stats = DynamoMetadataTagger().get_tagging_statistics(model) + assert stats["total_nodes"] == 3 + assert stats["direct_matches"] == 2 + assert stats["scoped_nodes"] == 2 + assert stats["root_fallbacks"] == 1 + assert stats["unique_scopes"] == 2 + # Keys required by the shared console/report/metadata writers exist. + for key in ("root_nodes", "parent_matches", "operation_matches"): + assert key in stats diff --git a/tests/unit/export/test_pytorch_export.py b/tests/unit/export/test_pytorch_export.py index 5b6aeae26..d74b4a3c6 100644 --- a/tests/unit/export/test_pytorch_export.py +++ b/tests/unit/export/test_pytorch_export.py @@ -318,8 +318,11 @@ def test_normalization_succeeds_and_shape_inferences(self, tmp_path) -> None: def test_failed_normalization_skips_shape_inference(self, tmp_path) -> None: """When normalization is mocked to return False, status is failed.""" model = TwoLayerNet() + # dynamo already emits value_info shapes during export; pin the legacy + # TorchScript exporter so this test exercises winml's own shape inference. config = WinMLExportConfig( input_tensors=[InputTensorSpec(name="x", dtype="float32", shape=(1, 10))], + dynamo=False, ) with patch( "winml.modelkit.export.pytorch._normalize_exported_model", @@ -335,8 +338,11 @@ def test_failed_normalization_skips_shape_inference(self, tmp_path) -> None: def test_normalize_false_skips_normalization(self, tmp_path) -> None: """When normalize=False, the helper isn't called and status is not_run.""" model = TwoLayerNet() + # dynamo already emits value_info shapes during export; pin the legacy + # TorchScript exporter so the "no shapes without normalization" invariant holds. config = WinMLExportConfig( input_tensors=[InputTensorSpec(name="x", dtype="float32", shape=(1, 10))], + dynamo=False, ) with patch( "winml.modelkit.export.pytorch._normalize_exported_model", @@ -486,8 +492,11 @@ def convert_with_per_tensor_sidecars( raw_sidecars.extend(get_external_data_files(out)) model_path = tmp_path / "model.onnx" + # This test forces per-tensor external-data sidecars to mimic the raw + # TorchScript exporter, so pin that exporter to test its cleanup path. config = WinMLExportConfig( input_tensors=[InputTensorSpec(name="x", dtype="float32", shape=(1, 256))], + dynamo=False, ) with patch.object( From 4e6d11fd71769af0c95f47574093168fb040e47d Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Wed, 22 Jul 2026 10:12:05 +0800 Subject: [PATCH 02/15] Address review: clarify QNN wording in docs, fix CodeQL double-import - test_dynamo_metadata_tagger: consolidate onnx imports to a single 'from onnx import ...' (CodeQL: module imported with both import and import-from) - docs/help: 'QNN-safe' reworded to name the real differentiator (opset/op decomposition, not shapes, since both exporters default to static shapes); add a static-shape export example --- docs/commands/export.md | 20 ++++++++++++++++--- docs/concepts/load-and-export.md | 2 +- src/winml/modelkit/commands/export.py | 3 ++- .../unit/core/test_dynamo_metadata_tagger.py | 17 ++++++++-------- 4 files changed, 28 insertions(+), 14 deletions(-) diff --git a/docs/commands/export.md b/docs/commands/export.md index 0e7c80562..218b7d2ef 100644 --- a/docs/commands/export.md +++ b/docs/commands/export.md @@ -22,7 +22,7 @@ $ winml export [options] | `--output` | `-o` | path | *(required)* | Output ONNX file path (e.g., `model.onnx`). | | `--with-report/--no-with-report` | | flag | `false` | Generate full export reports: Markdown, JSON, and a console tree. | | `--hierarchy/--no-hierarchy` | | flag | `true` | Preserve `hierarchy_tag` metadata in ONNX nodes (use `--no-hierarchy` for a clean ONNX file). | -| `--dynamo/--no-dynamo` | | flag | `true` | Use PyTorch's TorchDynamo ONNX exporter (default) for richer per-node module metadata. Pass `--no-dynamo` for the legacy TorchScript exporter (QNN-safe). | +| `--dynamo/--no-dynamo` | | flag | `true` | Use PyTorch's TorchDynamo ONNX exporter (default) for richer per-node module metadata. Pass `--no-dynamo` for the legacy TorchScript exporter, whose opset-17 op decomposition is the validated path for QNN/NPU compilation today. | | `--torch-module` | | string | `None` | Comma-separated list of `torch.nn` module types to include in hierarchy (e.g., `LayerNorm,Embedding`). (Experimental — currently logs a warning.) | | `--input-specs` | | path | `None` | JSON file with explicit input tensor specifications. Auto-generated when omitted. | | `--task` | `-t` | string | `None` | Override auto-detected Hugging Face task (e.g., `image-feature-extraction`). | @@ -79,6 +79,15 @@ winml export -m bert-base-uncased -o bert.onnx \ winml export -m bert-base-uncased -o bert.onnx --input-specs inputs.json ``` +```bash +# Export a fully static-shaped model (the default) for NPU/QNN compilation. +# The default dummy inputs already produce static shapes; use --shape-config to +# pin any symbolic dimensions to concrete sizes, or --input-specs to fully +# specify every input tensor. +winml export -m bert-base-uncased -o bert.onnx --shape-config shape_config.json +# shape_config.json: {"sequence_length": 128} +``` + ```bash # Export with dynamic batch and sequence dimensions winml export -m bert-base-uncased -o bert.onnx --dynamic-axes dynamic_axes.json @@ -116,8 +125,13 @@ winml export -m microsoft/resnet-50 -o resnet50_clean.onnx --no-hierarchy `--dynamic-axes` only when downstream runtime scenarios need variable sizes. - **Dynamo is the default exporter.** `winml export` uses PyTorch's TorchDynamo ONNX exporter, which records rich per-node module metadata that drives the - hierarchy tags. Pass `--no-dynamo` to select the legacy TorchScript exporter, - which remains the QNN-safe choice for hand exports targeting static shapes. + hierarchy tags. Pass `--no-dynamo` to select the legacy TorchScript exporter. + Shape staticness is independent of the exporter: both default to static shapes + and only emit dynamic axes when you ask for them. The QNN-relevant + difference is op decomposition -- the dynamo exporter defaults to a newer opset + and lowers some ops differently, which can reduce QNN fusion coverage + (e.g. MatMul + Add fusion). Prefer `--no-dynamo` for hand exports targeting + QNN/NPU compilation until the dynamo graph is validated for your model. - **`--torch-module` is experimental.** The flag emits a warning and has no effect in the current release. Do not rely on it in automated pipelines yet. - **Output directory must be writable.** The command creates parent directories diff --git a/docs/concepts/load-and-export.md b/docs/concepts/load-and-export.md index 1e7467a85..8a35c2e5c 100644 --- a/docs/concepts/load-and-export.md +++ b/docs/concepts/load-and-export.md @@ -16,7 +16,7 @@ Some community models host custom Python code in their repositories. The loader ## Exporting to ONNX -`winml export` converts the loaded model to ONNX. By default it uses PyTorch's TorchDynamo ONNX exporter (`torch.onnx.export(dynamo=True)`), which records rich per-node module metadata that is used to derive the `winml.hierarchy.*` node tags. Pass `--no-dynamo` to fall back to the legacy TorchScript exporter, which follows actual execution paths, tends to produce compact inference-oriented graphs, and remains the QNN-safe choice for static-shape exports. +`winml export` converts the loaded model to ONNX. By default it uses PyTorch's TorchDynamo ONNX exporter (`torch.onnx.export(dynamo=True)`), which records rich per-node module metadata that is used to derive the `winml.hierarchy.*` node tags. Pass `--no-dynamo` to fall back to the legacy TorchScript exporter, which follows actual execution paths and tends to produce compact inference-oriented graphs. Both exporters default to static shapes; the QNN-relevant difference is that the dynamo exporter defaults to a newer opset and decomposes some ops differently, so `--no-dynamo` remains the validated choice for QNN/NPU hand exports today. By default the exporter runs an eight-step process that includes hierarchy tracing and tag injection. The result is an ONNX file enriched with structural metadata that powers downstream features such as per-module benchmarking, inspector views, and optimizer scoping. diff --git a/src/winml/modelkit/commands/export.py b/src/winml/modelkit/commands/export.py index ad45dae82..3449ec858 100644 --- a/src/winml/modelkit/commands/export.py +++ b/src/winml/modelkit/commands/export.py @@ -119,7 +119,8 @@ def _warn_partial_composite(completed: list[Path]) -> None: default=True, show_default=True, help="Use PyTorch's TorchDynamo ONNX exporter (default). " - "Pass --no-dynamo for the legacy TorchScript exporter (QNN-safe).", + "Pass --no-dynamo for the legacy TorchScript exporter, the validated " + "path for QNN/NPU compilation today.", ) @click.option( "--torch-module", diff --git a/tests/unit/core/test_dynamo_metadata_tagger.py b/tests/unit/core/test_dynamo_metadata_tagger.py index 45fcccbfe..2ecb7b1ba 100644 --- a/tests/unit/core/test_dynamo_metadata_tagger.py +++ b/tests/unit/core/test_dynamo_metadata_tagger.py @@ -19,8 +19,7 @@ from __future__ import annotations -import onnx -from onnx import helper +from onnx import ModelProto, NodeProto, StringStringEntryProto, helper from winml.modelkit.core.onnx_node_tagger import DynamoMetadataTagger @@ -37,7 +36,7 @@ def _make_node( class_hierarchy: list[str] | None = None, raw_name_scopes: str | None = None, raw_class_hierarchy: str | None = None, -) -> onnx.NodeProto: +) -> NodeProto: """Build an ONNX node with dynamo-style metadata_props. ``name_scopes``/``class_hierarchy`` are serialized with ``repr`` exactly as @@ -46,30 +45,30 @@ def _make_node( node = helper.make_node(op_type, inputs=["x"], outputs=["y"], name=name) if raw_name_scopes is not None: node.metadata_props.append( - onnx.StringStringEntryProto(key=NAME_SCOPES_KEY, value=raw_name_scopes) + StringStringEntryProto(key=NAME_SCOPES_KEY, value=raw_name_scopes) ) elif name_scopes is not None: node.metadata_props.append( - onnx.StringStringEntryProto(key=NAME_SCOPES_KEY, value=repr(name_scopes)) + StringStringEntryProto(key=NAME_SCOPES_KEY, value=repr(name_scopes)) ) if raw_class_hierarchy is not None: node.metadata_props.append( - onnx.StringStringEntryProto(key=CLASS_HIERARCHY_KEY, value=raw_class_hierarchy) + StringStringEntryProto(key=CLASS_HIERARCHY_KEY, value=raw_class_hierarchy) ) elif class_hierarchy is not None: node.metadata_props.append( - onnx.StringStringEntryProto(key=CLASS_HIERARCHY_KEY, value=repr(class_hierarchy)) + StringStringEntryProto(key=CLASS_HIERARCHY_KEY, value=repr(class_hierarchy)) ) return node -def _make_model(nodes: list[onnx.NodeProto]) -> onnx.ModelProto: +def _make_model(nodes: list[NodeProto]) -> ModelProto: """Wrap nodes in a minimal ModelProto (tagger only iterates graph.node).""" graph = helper.make_graph(nodes, "g", inputs=[], outputs=[]) return helper.make_model(graph) -def _resnet_like_nodes() -> list[onnx.NodeProto]: +def _resnet_like_nodes() -> list[NodeProto]: """Two nodes mirroring the empirically verified dynamo output shape.""" return [ _make_node( From 40d20992d9a932ea5bf18cffd6acdec6362d3376 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Wed, 22 Jul 2026 10:35:35 +0800 Subject: [PATCH 03/15] docs: correct dynamo opset wording (min opset 18 + down-convert) The dynamo exporter does not default to a higher final opset; winml passes the same opset (17) to both paths. torch's dynamo op library only implements ops at a minimum opset of 18, so it builds at 18 and then down-converts to the requested opset. Reword the export docs and concept doc to describe that lowering-and-conversion pass accurately instead of claiming a "newer opset" default. --- docs/commands/export.md | 11 +++++++---- docs/concepts/load-and-export.md | 2 +- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/docs/commands/export.md b/docs/commands/export.md index 218b7d2ef..80f354297 100644 --- a/docs/commands/export.md +++ b/docs/commands/export.md @@ -128,10 +128,13 @@ winml export -m microsoft/resnet-50 -o resnet50_clean.onnx --no-hierarchy hierarchy tags. Pass `--no-dynamo` to select the legacy TorchScript exporter. Shape staticness is independent of the exporter: both default to static shapes and only emit dynamic axes when you ask for them. The QNN-relevant - difference is op decomposition -- the dynamo exporter defaults to a newer opset - and lowers some ops differently, which can reduce QNN fusion coverage - (e.g. MatMul + Add fusion). Prefer `--no-dynamo` for hand exports targeting - QNN/NPU compilation until the dynamo graph is validated for your model. + difference is op decomposition: the dynamo exporter's op library targets a + minimum opset of 18, so it lowers ops at opset 18 and then down-converts to the + requested opset (17 by default), whereas TorchScript builds natively at opset + 17. That extra lowering-and-conversion pass can decompose some ops differently + and reduce QNN fusion coverage (e.g. MatMul + Add fusion). Prefer `--no-dynamo` + for hand exports targeting QNN/NPU compilation until the dynamo graph is + validated for your model. - **`--torch-module` is experimental.** The flag emits a warning and has no effect in the current release. Do not rely on it in automated pipelines yet. - **Output directory must be writable.** The command creates parent directories diff --git a/docs/concepts/load-and-export.md b/docs/concepts/load-and-export.md index 8a35c2e5c..17b86327c 100644 --- a/docs/concepts/load-and-export.md +++ b/docs/concepts/load-and-export.md @@ -16,7 +16,7 @@ Some community models host custom Python code in their repositories. The loader ## Exporting to ONNX -`winml export` converts the loaded model to ONNX. By default it uses PyTorch's TorchDynamo ONNX exporter (`torch.onnx.export(dynamo=True)`), which records rich per-node module metadata that is used to derive the `winml.hierarchy.*` node tags. Pass `--no-dynamo` to fall back to the legacy TorchScript exporter, which follows actual execution paths and tends to produce compact inference-oriented graphs. Both exporters default to static shapes; the QNN-relevant difference is that the dynamo exporter defaults to a newer opset and decomposes some ops differently, so `--no-dynamo` remains the validated choice for QNN/NPU hand exports today. +`winml export` converts the loaded model to ONNX. By default it uses PyTorch's TorchDynamo ONNX exporter (`torch.onnx.export(dynamo=True)`), which records rich per-node module metadata that is used to derive the `winml.hierarchy.*` node tags. Pass `--no-dynamo` to fall back to the legacy TorchScript exporter, which follows actual execution paths and tends to produce compact inference-oriented graphs. Both exporters default to static shapes; the QNN-relevant difference is op decomposition: the dynamo exporter's op library targets a minimum opset of 18, so it lowers ops at opset 18 and then down-converts to the requested opset (17 by default), whereas TorchScript builds natively at opset 17. That extra lowering-and-conversion pass can decompose some ops differently, so `--no-dynamo` remains the validated choice for QNN/NPU hand exports today. By default the exporter runs an eight-step process that includes hierarchy tracing and tag injection. The result is an ONNX file enriched with structural metadata that powers downstream features such as per-module benchmarking, inspector views, and optimizer scoping. From 3e846971742c22c9417e07b9a97e1b1f15e2b8a7 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Wed, 22 Jul 2026 11:05:15 +0800 Subject: [PATCH 04/15] Report actual exported opset and correct dynamo opset docs torch's dynamo exporter targets a minimum opset of 18 and does not always lower to a smaller requested version (e.g. ResNet's ReduceMean keeps an opset-18-only attribute), so the export metadata previously misreported the requested opset (17) while the real graph was opset 18. - exporter: read the produced model's default-domain opset from the exported file and record that in the export monitor/metadata; warn when it differs from the requested opset. No post-export conversion is performed -- the requested opset is passed straight to the exporter and we report what it actually produced. - docs: correct the dynamo vs TorchScript opset wording (dynamo exports at opset 18, TorchScript at the configured opset 17) with the ResNet head decomposition example. - tests: add unit coverage for _read_default_opset. --- docs/commands/export.md | 15 +++--- docs/concepts/load-and-export.md | 2 +- src/winml/modelkit/export/htp/exporter.py | 39 ++++++++++++++- tests/unit/export/test_htp_exporter_stats.py | 52 ++++++++++++++++++++ 4 files changed, 99 insertions(+), 9 deletions(-) diff --git a/docs/commands/export.md b/docs/commands/export.md index 80f354297..e9d9a0176 100644 --- a/docs/commands/export.md +++ b/docs/commands/export.md @@ -128,13 +128,14 @@ winml export -m microsoft/resnet-50 -o resnet50_clean.onnx --no-hierarchy hierarchy tags. Pass `--no-dynamo` to select the legacy TorchScript exporter. Shape staticness is independent of the exporter: both default to static shapes and only emit dynamic axes when you ask for them. The QNN-relevant - difference is op decomposition: the dynamo exporter's op library targets a - minimum opset of 18, so it lowers ops at opset 18 and then down-converts to the - requested opset (17 by default), whereas TorchScript builds natively at opset - 17. That extra lowering-and-conversion pass can decompose some ops differently - and reduce QNN fusion coverage (e.g. MatMul + Add fusion). Prefer `--no-dynamo` - for hand exports targeting QNN/NPU compilation until the dynamo graph is - validated for your model. + difference is opset and op decomposition: torch's dynamo op library targets a + minimum opset of 18, so the dynamo path exports at opset 18, whereas the + TorchScript path exports at the configured opset (17 by default). Dynamo also + lowers some ops differently -- for example ResNet's classification head becomes + `ReduceMean` + `Reshape` under dynamo but `GlobalAveragePool` + `Flatten` under + TorchScript -- which can change or reduce QNN fusion coverage. Prefer + `--no-dynamo` for hand exports targeting QNN/NPU compilation until the dynamo + graph is validated for your model. - **`--torch-module` is experimental.** The flag emits a warning and has no effect in the current release. Do not rely on it in automated pipelines yet. - **Output directory must be writable.** The command creates parent directories diff --git a/docs/concepts/load-and-export.md b/docs/concepts/load-and-export.md index 17b86327c..104707e9e 100644 --- a/docs/concepts/load-and-export.md +++ b/docs/concepts/load-and-export.md @@ -16,7 +16,7 @@ Some community models host custom Python code in their repositories. The loader ## Exporting to ONNX -`winml export` converts the loaded model to ONNX. By default it uses PyTorch's TorchDynamo ONNX exporter (`torch.onnx.export(dynamo=True)`), which records rich per-node module metadata that is used to derive the `winml.hierarchy.*` node tags. Pass `--no-dynamo` to fall back to the legacy TorchScript exporter, which follows actual execution paths and tends to produce compact inference-oriented graphs. Both exporters default to static shapes; the QNN-relevant difference is op decomposition: the dynamo exporter's op library targets a minimum opset of 18, so it lowers ops at opset 18 and then down-converts to the requested opset (17 by default), whereas TorchScript builds natively at opset 17. That extra lowering-and-conversion pass can decompose some ops differently, so `--no-dynamo` remains the validated choice for QNN/NPU hand exports today. +`winml export` converts the loaded model to ONNX. By default it uses PyTorch's TorchDynamo ONNX exporter (`torch.onnx.export(dynamo=True)`), which records rich per-node module metadata that is used to derive the `winml.hierarchy.*` node tags. Pass `--no-dynamo` to fall back to the legacy TorchScript exporter, which follows actual execution paths and tends to produce compact inference-oriented graphs. Both exporters default to static shapes; the QNN-relevant difference is opset and op decomposition: torch's dynamo op library targets a minimum opset of 18, so the dynamo path exports at opset 18, whereas the TorchScript path exports at the configured opset (17 by default) and lowers some ops differently (e.g. ResNet's head becomes `ReduceMean` + `Reshape` under dynamo but `GlobalAveragePool` + `Flatten` under TorchScript). Because the opset-17 TorchScript graph is what the QNN/NPU toolchain has been validated against, `--no-dynamo` remains the validated choice for QNN/NPU hand exports today. By default the exporter runs an eight-step process that includes hierarchy tracing and tag injection. The result is an ONNX file enriched with structural metadata that powers downstream features such as per-module benchmarking, inspector views, and optimizer scoping. diff --git a/src/winml/modelkit/export/htp/exporter.py b/src/winml/modelkit/export/htp/exporter.py index 211e394f2..c82d6dcca 100644 --- a/src/winml/modelkit/export/htp/exporter.py +++ b/src/winml/modelkit/export/htp/exporter.py @@ -295,9 +295,27 @@ def export( self._hierarchy_builder.get_outputs() if self._hierarchy_builder else None ) output_names = infer_output_names(traced_outputs) if traced_outputs is not None else [] + # Report the opset the exporter actually produced, not the requested + # one. torch's dynamo exporter targets a minimum opset of 18 and does + # not always down-convert to a lower requested value (e.g. ResNet stays + # at 18), so echoing export_config.opset_version would misreport the + # real graph. Fall back to the request only if the value can't be read. + actual_opset = self._read_default_opset(output_path) + if actual_opset is not None and actual_opset != export_config.opset_version: + logger.warning( + "Requested opset %d but the exporter produced opset %d. " + "torch's dynamo exporter targets a minimum opset of 18 and " + "cannot always lower to a smaller requested version; pass " + "--no-dynamo for a natively opset-%d graph.", + export_config.opset_version, + actual_opset, + export_config.opset_version, + ) monitor.update( ExportStep.ONNX_EXPORT, - opset_version=export_config.opset_version, + opset_version=( + actual_opset if actual_opset is not None else export_config.opset_version + ), do_constant_folding=export_config.do_constant_folding, onnx_size_mb=onnx_size_mb, output_names=output_names, @@ -385,6 +403,25 @@ def _trace_model_hierarchy(self, model: nn.Module, inputs: dict) -> None: self._hierarchy_data = summary["module_hierarchy"] self._export_stats["hierarchy_modules"] = len(self._hierarchy_data) + @staticmethod + def _read_default_opset(output_path: str) -> int | None: + """Return the default-domain (ai.onnx) opset of the exported model. + + Reads only the model proto (external weight files are not loaded), so it + stays cheap for large models. Returns ``None`` if the model cannot be + read or declares no default-domain opset import. + """ + import onnx + + try: + model = onnx.load(output_path, load_external_data=False) + except Exception: + return None + for opset in model.opset_import: + if opset.domain in ("", "ai.onnx"): + return opset.version + return None + def _verify_onnx_export( self, output_path: str, diff --git a/tests/unit/export/test_htp_exporter_stats.py b/tests/unit/export/test_htp_exporter_stats.py index 425a3713c..403d6cef7 100644 --- a/tests/unit/export/test_htp_exporter_stats.py +++ b/tests/unit/export/test_htp_exporter_stats.py @@ -6,11 +6,18 @@ from __future__ import annotations +from typing import TYPE_CHECKING from unittest.mock import MagicMock +from onnx import TensorProto, helper, save + from winml.modelkit.export.htp import HTPExporter +if TYPE_CHECKING: + from pathlib import Path + + class TestHTPExporterTaggedNodesStats: """tagged_nodes, empty_tags, and coverage must be 0 when embed_hierarchy_attributes=False.""" @@ -51,3 +58,48 @@ def test_stats_populated_when_hierarchy_enabled(self) -> None: assert exporter._export_stats["tagged_nodes"] == 2 assert exporter._export_stats["coverage_percentage"] == 50.0 assert exporter._export_stats["empty_tags"] == 0 + + +class TestHTPExporterReadDefaultOpset: + """_read_default_opset reports the produced model's ai.onnx opset. + + The dynamo exporter may not honor a lower requested opset (torch's dynamo op + set has a minimum of 18), so the exporter records the opset actually present + in the file instead of echoing the requested value. + """ + + @staticmethod + def _make_model(tmp_path: Path, opset_imports: list) -> str: + node = helper.make_node("Identity", ["x"], ["y"]) + graph = helper.make_graph( + [node], + "g", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [1])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [1])], + ) + model = helper.make_model(graph, opset_imports=opset_imports) + path = tmp_path / "m.onnx" + save(model, str(path)) + return str(path) + + def test_reads_default_domain_opset(self, tmp_path: Path) -> None: + path = self._make_model(tmp_path, [helper.make_opsetid("", 18)]) + assert HTPExporter._read_default_opset(path) == 18 + + def test_reads_ai_onnx_domain_opset(self, tmp_path: Path) -> None: + path = self._make_model(tmp_path, [helper.make_opsetid("ai.onnx", 17)]) + assert HTPExporter._read_default_opset(path) == 17 + + def test_prefers_default_domain_over_custom(self, tmp_path: Path) -> None: + path = self._make_model( + tmp_path, + [helper.make_opsetid("com.microsoft", 1), helper.make_opsetid("", 18)], + ) + assert HTPExporter._read_default_opset(path) == 18 + + def test_returns_none_without_default_domain(self, tmp_path: Path) -> None: + path = self._make_model(tmp_path, [helper.make_opsetid("com.microsoft", 1)]) + assert HTPExporter._read_default_opset(path) is None + + def test_returns_none_for_unreadable_path(self, tmp_path: Path) -> None: + assert HTPExporter._read_default_opset(str(tmp_path / "missing.onnx")) is None From 99e67aa10fd9e0d5bae77877711f9ba4148c0f1b Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Wed, 22 Jul 2026 11:47:56 +0800 Subject: [PATCH 05/15] Suppress onnxscript version-converter traceback unless verbose When the dynamo exporter cannot down-convert a model to the requested opset, onnxscript's version_converter logs a WARNING with a full call stack (e.g. ResNet's ReduceMean during an 18->17 attempt). winml already surfaces a concise opset warning for that case, so the raw traceback is redundant noise in normal output. Add onnxscript.version_converter to the noisy-library logger floor so it is pinned at ERROR by default and only follows the CLI level once the user passes -v/-vv, matching how optimum's chatter is handled. --- src/winml/modelkit/utils/logging.py | 12 ++++++---- tests/unit/utils/test_logging.py | 37 +++++++++++++++++++++++------ 2 files changed, 38 insertions(+), 11 deletions(-) diff --git a/src/winml/modelkit/utils/logging.py b/src/winml/modelkit/utils/logging.py index bacb54427..8696f4cd1 100644 --- a/src/winml/modelkit/utils/logging.py +++ b/src/winml/modelkit/utils/logging.py @@ -35,10 +35,14 @@ _DATE_FORMAT = "%H:%M:%S" # Third-party loggers whose INFO/WARNING chatter is noise for CLI users and can -# interleave with rich progress output (e.g. optimum's "No model type passed for the -# task ..." notice when a task maps to several loader classes). They are floored at -# ERROR in normal output and only follow the CLI level once the user passes -v/-vv. -_NOISY_LIBRARY_LOGGERS = ("optimum",) +# interleave with rich progress output. Examples: optimum's "No model type passed +# for the task ..." notice when a task maps to several loader classes, and +# onnxscript's version-converter fallback WARNING (with a full call stack) that +# fires when the dynamo exporter cannot down-convert a model to the requested opset +# -- winml already surfaces a concise opset warning for that case, so the raw +# traceback is redundant. They are floored at ERROR in normal output and only follow +# the CLI level once the user passes -v/-vv. +_NOISY_LIBRARY_LOGGERS = ("optimum", "onnxscript.version_converter") def configure_logging( diff --git a/tests/unit/utils/test_logging.py b/tests/unit/utils/test_logging.py index bbce52e27..f897e0dba 100644 --- a/tests/unit/utils/test_logging.py +++ b/tests/unit/utils/test_logging.py @@ -8,19 +8,20 @@ import pytest -from winml.modelkit.utils.logging import configure_logging +from winml.modelkit.utils.logging import _NOISY_LIBRARY_LOGGERS, configure_logging @pytest.fixture(autouse=True) def _restore_logger_levels(): """configure_logging mutates global logger state (root + the noisy library loggers); - restore both after each test so verbosity changes don't leak across tests.""" - root = logging.getLogger() - optimum = logging.getLogger("optimum") - root_before, optimum_before = root.level, optimum.level + restore all of them after each test so verbosity changes don't leak across tests.""" + saved = [(logging.getLogger(), logging.getLogger().level)] + for name in _NOISY_LIBRARY_LOGGERS: + logger = logging.getLogger(name) + saved.append((logger, logger.level)) yield - root.setLevel(root_before) - optimum.setLevel(optimum_before) + for logger, level in saved: + logger.setLevel(level) def test_library_loggers_floored_at_error_in_normal_mode(): @@ -54,3 +55,25 @@ def test_optimum_child_logger_gated_by_parent_floor(): configure_logging(verbosity=1) assert child.isEnabledFor(logging.WARNING) + + +def test_onnxscript_version_converter_floored_at_error_in_normal_mode(): + # The onnxscript version-converter fallback WARNING carries a full call stack when + # the dynamo exporter cannot down-convert to the requested opset. winml surfaces + # its own concise opset warning, so the raw traceback is floored out by default. + configure_logging(verbosity=0) + assert logging.getLogger("onnxscript.version_converter").level == logging.ERROR + + +@pytest.mark.parametrize("verbosity,expected", [(1, logging.INFO), (2, logging.DEBUG)]) +def test_onnxscript_version_converter_revealed_when_verbose(verbosity, expected): + # -v/-vv opts into the detail: the converter logger follows the CLI level so the + # call stack becomes visible on demand. + logger = logging.getLogger("onnxscript.version_converter") + + configure_logging(verbosity=0) + assert not logger.isEnabledFor(logging.WARNING) + + configure_logging(verbosity=verbosity) + assert logger.level == expected + assert logger.isEnabledFor(logging.WARNING) From 9607cfc6ce2db8577130be2ccd00c7115b436bec Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Wed, 22 Jul 2026 12:23:54 +0800 Subject: [PATCH 06/15] Fix onnx double-import and hide torch opset-downgrade notice Resolves the CodeQL double-import alert on exporter.py: _read_default_opset now reads the proto through winml's own load_onnx helper (load_weights=False, validate=False) instead of a bare `import onnx`, matching the runtime onnx-loading style already used elsewhere in the file. Also floors torch's one-line "Setting ONNX exporter to use operator set version 18 ..." notice (logger torch.onnx._internal.exporter._compat) at ERROR by default, alongside the onnxscript version-converter traceback. winml already surfaces a concise opset warning, so both are redundant in normal output and only follow the CLI level at -v/-vv. --- src/winml/modelkit/export/htp/exporter.py | 4 ++-- src/winml/modelkit/utils/logging.py | 19 ++++++++++++------- tests/unit/utils/test_logging.py | 23 +++++++++++++++++++++++ 3 files changed, 37 insertions(+), 9 deletions(-) diff --git a/src/winml/modelkit/export/htp/exporter.py b/src/winml/modelkit/export/htp/exporter.py index c82d6dcca..29820e446 100644 --- a/src/winml/modelkit/export/htp/exporter.py +++ b/src/winml/modelkit/export/htp/exporter.py @@ -411,10 +411,10 @@ def _read_default_opset(output_path: str) -> int | None: stays cheap for large models. Returns ``None`` if the model cannot be read or declares no default-domain opset import. """ - import onnx + from ...onnx import load_onnx try: - model = onnx.load(output_path, load_external_data=False) + model = load_onnx(output_path, load_weights=False, validate=False) except Exception: return None for opset in model.opset_import: diff --git a/src/winml/modelkit/utils/logging.py b/src/winml/modelkit/utils/logging.py index 8696f4cd1..15f1dfabd 100644 --- a/src/winml/modelkit/utils/logging.py +++ b/src/winml/modelkit/utils/logging.py @@ -36,13 +36,18 @@ # Third-party loggers whose INFO/WARNING chatter is noise for CLI users and can # interleave with rich progress output. Examples: optimum's "No model type passed -# for the task ..." notice when a task maps to several loader classes, and -# onnxscript's version-converter fallback WARNING (with a full call stack) that -# fires when the dynamo exporter cannot down-convert a model to the requested opset -# -- winml already surfaces a concise opset warning for that case, so the raw -# traceback is redundant. They are floored at ERROR in normal output and only follow -# the CLI level once the user passes -v/-vv. -_NOISY_LIBRARY_LOGGERS = ("optimum", "onnxscript.version_converter") +# for the task ..." notice when a task maps to several loader classes; onnxscript's +# version-converter fallback WARNING (with a full call stack) that fires when the +# dynamo exporter cannot down-convert a model to the requested opset; and torch's +# own one-line "Setting ONNX exporter to use operator set version 18 ..." notice for +# the same down-convert case -- winml already surfaces a concise opset warning, so +# both are redundant. They are floored at ERROR in normal output and only follow the +# CLI level once the user passes -v/-vv. +_NOISY_LIBRARY_LOGGERS = ( + "optimum", + "onnxscript.version_converter", + "torch.onnx._internal.exporter._compat", +) def configure_logging( diff --git a/tests/unit/utils/test_logging.py b/tests/unit/utils/test_logging.py index f897e0dba..7c49c9d63 100644 --- a/tests/unit/utils/test_logging.py +++ b/tests/unit/utils/test_logging.py @@ -77,3 +77,26 @@ def test_onnxscript_version_converter_revealed_when_verbose(verbosity, expected) configure_logging(verbosity=verbosity) assert logger.level == expected assert logger.isEnabledFor(logging.WARNING) + + +def test_torch_compat_opset_notice_floored_at_error_in_normal_mode(): + # torch's exporter emits a one-line "Setting ONNX exporter to use operator set + # version 18 ..." WARNING when it cannot honor a lower requested opset. winml + # surfaces its own concise opset warning, so torch's notice is floored by default. + configure_logging(verbosity=0) + logger = logging.getLogger("torch.onnx._internal.exporter._compat") + assert logger.level == logging.ERROR + assert not logger.isEnabledFor(logging.WARNING) + + +@pytest.mark.parametrize("verbosity,expected", [(1, logging.INFO), (2, logging.DEBUG)]) +def test_torch_compat_opset_notice_revealed_when_verbose(verbosity, expected): + # -v/-vv opts into the detail: the torch logger follows the CLI level. + logger = logging.getLogger("torch.onnx._internal.exporter._compat") + + configure_logging(verbosity=0) + assert not logger.isEnabledFor(logging.WARNING) + + configure_logging(verbosity=verbosity) + assert logger.level == expected + assert logger.isEnabledFor(logging.WARNING) From dc3c49c2e419e1b243669982dbe6e00fd8095b3d Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Wed, 22 Jul 2026 15:52:18 +0800 Subject: [PATCH 07/15] Recover dynamo module hierarchy and disambiguate same-class siblings Two dynamo-path fidelity fixes from review: - DynamoMetadataTagger now folds each module's local scope name into its tag segment (Class.local) for every level, not just numeric indices, so same-class named siblings (an attention block's query/key/value Linears) stay independently addressable instead of collapsing to a single /.../Linear tag. Indexed children keep the existing Class.N form. - Add DynamoMetadataTagger.build_module_hierarchy to reconstruct the flat module hierarchy from node metadata. The exporter now recovers it after loading the exported ONNX and issues the deferred HIERARCHY monitor update, so the report/console/metadata module tree and the hierarchy_modules stat are populated for dynamo exports (previously 0). Output names for the dynamo path are sourced from the ONNX graph outputs instead of the empty traced outputs. Adds same-class-sibling tag tests and build_module_hierarchy unit tests. --- src/winml/modelkit/core/onnx_node_tagger.py | 90 +++++++++++--- src/winml/modelkit/export/htp/exporter.py | 71 +++++++---- .../unit/core/test_dynamo_metadata_tagger.py | 111 +++++++++++++++++- 3 files changed, 226 insertions(+), 46 deletions(-) diff --git a/src/winml/modelkit/core/onnx_node_tagger.py b/src/winml/modelkit/core/onnx_node_tagger.py index faa47a53e..aa39b683e 100644 --- a/src/winml/modelkit/core/onnx_node_tagger.py +++ b/src/winml/modelkit/core/onnx_node_tagger.py @@ -20,7 +20,7 @@ import ast from collections import defaultdict -from typing import TYPE_CHECKING, ClassVar +from typing import TYPE_CHECKING, Any, ClassVar if TYPE_CHECKING: @@ -377,30 +377,49 @@ def _short_class(fq_class: str) -> str: """ return fq_class.rsplit(".", 1)[-1] if fq_class else "" - def _tag_from_metadata(self, class_hierarchy: list[str], name_scopes: list[str]) -> str | None: - """Build a ``/Root/Child.N/Leaf`` tag from aligned dynamo metadata lists. + @staticmethod + def _segment(short_class: str, scope: str) -> str: + """Build one hierarchy-tag segment from a class name and cumulative scope. + + The scope's final path component is the module's local attribute name + (e.g. ``query``) or its ``ModuleList``/``Sequential`` index (e.g. ``0``). + It is appended as ``Class.local`` so same-class siblings stay distinct: + an attention block's ``query``/``key``/``value`` (all ``Linear``) become + ``Linear.query``/``Linear.key``/``Linear.value`` and indexed children + become ``Class.N``. The root level has an empty scope and stays the bare + class name. + """ + local = scope.rsplit(".", 1)[-1] if scope else "" + return f"{short_class}.{local}" if local else short_class + + def _module_segments( + self, class_hierarchy: list[str], name_scopes: list[str] + ) -> list[tuple[str, str, str]]: + """Yield ``(scope, short_class, segment)`` for each module level. - The last element of each list is the node's own op target / node name - (not a module level) and is dropped. For every remaining level the - class supplies the segment name; when the cumulative scope's final - component is a pure index (a ``ModuleList``/``Sequential`` child) it is - folded in as ``Class.N`` to mirror the TorchScript tagger's indexed - module tags. + The last element of each aligned list is the node's own op target / node + name (not a module level) and is dropped. Levels whose class name is + empty are skipped so a malformed entry never emits a blank segment. """ classes = class_hierarchy[:-1] scopes = name_scopes[:-1] - segments: list[str] = [] + levels: list[tuple[str, str, str]] = [] for cls, scope in zip(classes, scopes, strict=False): short = self._short_class(cls) if not short: continue - local = scope.rsplit(".", 1)[-1] if scope else "" - if local.isdigit(): - segments.append(f"{short}.{local}") - else: - segments.append(short) + levels.append((scope, short, self._segment(short, scope))) + return levels + def _tag_from_metadata(self, class_hierarchy: list[str], name_scopes: list[str]) -> str | None: + """Build a ``/Root/Child.N/Leaf`` tag from aligned dynamo metadata lists. + + Every module level contributes a :meth:`_segment`, so both indexed + (``Blk.0``) and named (``Linear.query``) modules stay uniquely + addressable. Returns ``None`` when no module level survives. + """ + segments = [seg for _, _, seg in self._module_segments(class_hierarchy, name_scopes)] if not segments: return None return "/" + "/".join(segments) @@ -472,3 +491,44 @@ def get_tagging_statistics(self, onnx_model: onnx.ModelProto) -> dict[str, int]: "operation_matches": 0, "root_fallbacks": root_fallbacks, } + + def build_module_hierarchy(self, onnx_model: onnx.ModelProto) -> dict[str, dict[str, Any]]: + """Reconstruct the module hierarchy from dynamo node metadata. + + The dynamo exporter does not run the TorchScript trace that normally + feeds the export monitor's module tree, so this rebuilds the same flat + ``{scope_path -> module info}`` mapping directly from the per-node + ``name_scopes``/``class_hierarchy`` metadata. The result matches the + shape :class:`TracingHierarchyBuilder` produces, so the shared + report/console/metadata writers render the dynamo module tree + identically. + + The root module is keyed by ``""``; every other module is keyed by its + cumulative dotted scope path (e.g. ``"blocks.0"``, + ``"blocks.0.attention.query"``). Each value carries the module's short + class name, its full hierarchy tag, and first-seen execution order. + """ + self.model_root_tag = self._extract_model_root_tag(onnx_model) + + hierarchy: dict[str, dict[str, Any]] = {} + order = 0 + for node in onnx_model.graph.node: + metadata = self._node_metadata(node) + classes = self._parse_list(metadata.get(self.CLASS_HIERARCHY_KEY)) + scopes = self._parse_list(metadata.get(self.NAME_SCOPES_KEY)) + if not classes or not scopes: + continue + + segments: list[str] = [] + for scope, short, segment in self._module_segments(classes, scopes): + segments.append(segment) + if scope in hierarchy: + continue + hierarchy[scope] = { + "class_name": short, + "traced_tag": "/" + "/".join(segments), + "execution_order": order, + } + order += 1 + + return hierarchy diff --git a/src/winml/modelkit/export/htp/exporter.py b/src/winml/modelkit/export/htp/exporter.py index 29820e446..09a0c9cc1 100644 --- a/src/winml/modelkit/export/htp/exporter.py +++ b/src/winml/modelkit/export/htp/exporter.py @@ -244,7 +244,8 @@ def export( # Under dynamo the exporter records the module hierarchy natively on # each ONNX node's metadata_props, so the TorchScript trace is # redundant (and can fail for models that only export via dynamo). - # The hierarchy is recovered post-export from that metadata in Step 5. + # The hierarchy is recovered from that metadata after the ONNX is + # loaded (Step 5), so the monitor update is deferred for dynamo. if self._use_dynamo_hierarchy: self._hierarchy_data = {} else: @@ -256,16 +257,16 @@ def export( with self._get_optimum_patcher(model, task): self._trace_model_hierarchy(model, inputs) - execution_steps = ( - self._hierarchy_builder.get_execution_summary().get("execution_steps", 0) - if self._hierarchy_builder - else 0 - ) - monitor.update( - ExportStep.HIERARCHY, - hierarchy=self._hierarchy_data, - execution_steps=execution_steps, - ) + execution_steps = ( + self._hierarchy_builder.get_execution_summary().get("execution_steps", 0) + if self._hierarchy_builder + else 0 + ) + monitor.update( + ExportStep.HIERARCHY, + hierarchy=self._hierarchy_data, + execution_steps=execution_steps, + ) # Step 4: ONNX Export self._convert_model_to_onnx(model, output_path, inputs, export_config, task=task) @@ -285,16 +286,44 @@ def export( # Verify ONNX export self._verify_onnx_export(output_path, export_config) + # Load the exported ONNX once and create the node tagger now; the same + # model is reused for dynamo hierarchy recovery, node tagging, and + # graph-metadata embedding. + from ...onnx import load_onnx + + onnx_model = load_onnx(output_path, validate=False) + self._initialize_node_tagger(enable_operation_fallback) + + # Under dynamo the module hierarchy lives on the ONNX node metadata + # instead of a TorchScript trace. Recover it now that the graph exists + # and issue the deferred Step 3 update so the report/console/metadata + # module tree and the hierarchy_modules stat reflect the real modules. + if self._use_dynamo_hierarchy: + assert isinstance(self._node_tagger, DynamoMetadataTagger) + self._hierarchy_data = self._node_tagger.build_module_hierarchy(onnx_model) + monitor.update( + ExportStep.HIERARCHY, + hierarchy=self._hierarchy_data, + execution_steps=len(self._hierarchy_data), + ) + # Update monitor with ONNX export info onnx_size_mb = ( round(Path(output_path).stat().st_size / (1024 * 1024), 2) if Path(output_path).exists() else 0 ) - traced_outputs = ( - self._hierarchy_builder.get_outputs() if self._hierarchy_builder else None - ) - output_names = infer_output_names(traced_outputs) if traced_outputs is not None else [] + # Output names: dynamo exposes them authoritatively on the ONNX graph; + # the TorchScript path derives them from the traced outputs. + if self._use_dynamo_hierarchy: + output_names = [output.name for output in onnx_model.graph.output] + else: + traced_outputs = ( + self._hierarchy_builder.get_outputs() if self._hierarchy_builder else None + ) + output_names = ( + infer_output_names(traced_outputs) if traced_outputs is not None else [] + ) # Report the opset the exporter actually produced, not the requested # one. torch's dynamo exporter targets a minimum opset of 18 and does # not always down-convert to a lower requested value (e.g. ResNet stays @@ -321,17 +350,7 @@ def export( output_names=output_names, ) - # Step 5: Node Tagger Creation - from ...onnx import load_onnx - - onnx_model = load_onnx(output_path, validate=False) - - self._initialize_node_tagger(enable_operation_fallback) - - # Tagger creation is part of node tagging process - # No separate step needed - - # Step 6: Node Tagging + # Step 6: Node Tagging (tagger was created above) self._apply_hierarchy_tags(onnx_model) # Update monitor with tagging results diff --git a/tests/unit/core/test_dynamo_metadata_tagger.py b/tests/unit/core/test_dynamo_metadata_tagger.py index 2ecb7b1ba..b6d8e7f6f 100644 --- a/tests/unit/core/test_dynamo_metadata_tagger.py +++ b/tests/unit/core/test_dynamo_metadata_tagger.py @@ -103,10 +103,37 @@ def test_tags_indexed_and_named_modules(self) -> None: model = _make_model(_resnet_like_nodes()) tags = DynamoMetadataTagger().tag_all_nodes(model) - # Digit scope component ("blocks.0") folds into "Blk.0"; named module - # ("blocks.0.lin" -> Linear) uses the class short-name only. - assert tags["node_gemm"] == "/Net/Blk.0/Linear" - assert tags["node_relu"] == "/Net/Blk.0/ReLU" + # Indexed scope component ("blocks.0") folds into "Blk.0"; named modules + # ("blocks.0.lin" -> Linear, "blocks.0.act" -> ReLU) fold their local + # attribute name in the same way so same-class siblings stay distinct. + assert tags["node_gemm"] == "/Net/Blk.0/Linear.lin" + assert tags["node_relu"] == "/Net/Blk.0/ReLU.act" + + def test_same_class_named_siblings_stay_distinct(self) -> None: + # Attention query/key/value are all torch.nn.Linear; without folding the + # local attribute name they would collapse to a single "/.../Linear" tag + # and could no longer be benchmarked or scoped independently. + siblings = ("query", "key", "value") + nodes = [ + _make_node( + "MatMul", + f"node_{name}", + name_scopes=["", "attn", f"attn.{name}", "linear"], + class_hierarchy=[ + "pkg.Net", + "pkg.Attention", + "torch.nn.modules.linear.Linear", + "aten.linear.default", + ], + ) + for name in siblings + ] + tags = DynamoMetadataTagger().tag_all_nodes(_make_model(nodes)) + + assert tags["node_query"] == "/Net/Attention.attn/Linear.query" + assert tags["node_key"] == "/Net/Attention.attn/Linear.key" + assert tags["node_value"] == "/Net/Attention.attn/Linear.value" + assert len({tags["node_query"], tags["node_key"], tags["node_value"]}) == 3 def test_model_root_tag_from_first_class(self) -> None: tagger = DynamoMetadataTagger() @@ -182,7 +209,7 @@ def test_unnamed_node_key_uses_optype_fallback(self) -> None: # Key mirrors exporter convention: "_". (only_key,) = tags.keys() assert only_key.startswith("Gemm_") - assert tags[only_key] == "/Net/Blk.0/Linear" + assert tags[only_key] == "/Net/Blk.0/Linear.lin" def test_get_tagging_statistics(self) -> None: nodes = _resnet_like_nodes() @@ -198,3 +225,77 @@ def test_get_tagging_statistics(self) -> None: # Keys required by the shared console/report/metadata writers exist. for key in ("root_nodes", "parent_matches", "operation_matches"): assert key in stats + + +def _attention_nodes() -> list[NodeProto]: + """Query/key/value nodes for one attention block (all torch.nn.Linear).""" + return [ + _make_node( + "MatMul", + f"node_{name}", + name_scopes=["", "blocks.0", "blocks.0.attn", f"blocks.0.attn.{name}", "linear"], + class_hierarchy=[ + "pkg.Net", + "pkg.Blk", + "pkg.Attention", + "torch.nn.modules.linear.Linear", + "aten.linear.default", + ], + ) + for name in ("query", "key", "value") + ] + + +class TestBuildModuleHierarchy: + """Reconstruction of the flat module hierarchy from dynamo node metadata.""" + + def test_root_module_keyed_by_empty_string(self) -> None: + hierarchy = DynamoMetadataTagger().build_module_hierarchy(_make_model(_resnet_like_nodes())) + assert "" in hierarchy + assert hierarchy[""]["class_name"] == "Net" + assert hierarchy[""]["traced_tag"] == "/Net" + + def test_cumulative_scope_paths_and_tags(self) -> None: + hierarchy = DynamoMetadataTagger().build_module_hierarchy(_make_model(_resnet_like_nodes())) + # Each module level is keyed by its cumulative dotted scope path with the + # matching class name and full hierarchy tag. + assert hierarchy["blocks.0"]["class_name"] == "Blk" + assert hierarchy["blocks.0"]["traced_tag"] == "/Net/Blk.0" + assert hierarchy["blocks.0.lin"]["class_name"] == "Linear" + assert hierarchy["blocks.0.lin"]["traced_tag"] == "/Net/Blk.0/Linear.lin" + assert hierarchy["blocks.0.act"]["traced_tag"] == "/Net/Blk.0/ReLU.act" + + def test_same_class_siblings_get_distinct_entries(self) -> None: + hierarchy = DynamoMetadataTagger().build_module_hierarchy(_make_model(_attention_nodes())) + for name in ("query", "key", "value"): + key = f"blocks.0.attn.{name}" + assert hierarchy[key]["class_name"] == "Linear" + assert hierarchy[key]["traced_tag"] == f"/Net/Blk.0/Attention.attn/Linear.{name}" + # The shared attention parent is recorded once. + assert hierarchy["blocks.0.attn"]["class_name"] == "Attention" + + def test_execution_order_is_unique_per_module(self) -> None: + hierarchy = DynamoMetadataTagger().build_module_hierarchy(_make_model(_attention_nodes())) + orders = [info["execution_order"] for info in hierarchy.values()] + assert all(isinstance(o, int) for o in orders) + assert len(orders) == len(set(orders)) + + def test_nodes_without_metadata_are_skipped(self) -> None: + nodes = _resnet_like_nodes() + nodes.append(_make_node("Add", "node_bare")) # no metadata + nodes.append( + _make_node( + "Add", + "node_bad", + raw_name_scopes="not-a-list", + raw_class_hierarchy="[unclosed", + ) + ) + hierarchy = DynamoMetadataTagger().build_module_hierarchy(_make_model(nodes)) + # Only the real module scopes are present; malformed nodes contribute + # nothing and do not raise. + assert set(hierarchy) == {"", "blocks.0", "blocks.0.lin", "blocks.0.act"} + + def test_empty_graph_returns_empty_hierarchy(self) -> None: + hierarchy = DynamoMetadataTagger().build_module_hierarchy(_make_model([])) + assert hierarchy == {} From 9b360f881b8c7b8d96ab5619b90bdffd584425f9 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Wed, 22 Jul 2026 16:43:06 +0800 Subject: [PATCH 08/15] Fix dynamo hierarchy writer traversal, sibling keys, and opset docs Round-2 review follow-ups on the reconstructed dynamo module hierarchy. The tagger-level fixes from the previous commit were correct, but the shared tree writers could still drop or collapse modules: - metadata_writer: same-class named siblings (an attention block's query/key/value Linears) collapsed to one tree entry because children were keyed by bare class name. Key indexed children as Class.N (unchanged), disambiguate colliding same-class named siblings as Class.local, and keep the bare Class key for the unique case. Two-pass collision detection makes the keys order-independent, with a final full-remainder fallback for safety. - hierarchy_utils.find_immediate_children: rewrite traversal around a nearest-present-ancestor helper so a sparse hierarchy (an indexed child whose ModuleList container never emitted its own scope, e.g. blocks.0 with no blocks) attaches to its nearest real ancestor instead of being orphaned. metadata_writer now reuses this shared traversal, fixing the same root-scope limitation in the JSON report. - exporter: normalize the TorchScript output_names branch so output_names is always list[str] (infer_output_names returns list[str] | None), fixing the mypy union at the merge point. - docs: the dynamo path does not always export at opset 18. It targets a minimum opset of 18 and attempts to down-convert to a lower requested opset, which may succeed or fail; winml reports the opset actually produced. Update export.md and load-and-export.md accordingly. Add tests/unit/export/test_dynamo_hierarchy_report.py exercising the real MetadataWriter tree: all same-class siblings serialize and the sparse-root subtree is present; plus find_immediate_children sparse/dense nesting. --- docs/commands/export.md | 11 +- docs/concepts/load-and-export.md | 2 +- src/winml/modelkit/core/hierarchy_utils.py | 55 ++++---- src/winml/modelkit/export/htp/exporter.py | 7 +- .../modelkit/export/htp/metadata_writer.py | 92 ++++++------- .../export/test_dynamo_hierarchy_report.py | 123 ++++++++++++++++++ 6 files changed, 204 insertions(+), 86 deletions(-) create mode 100644 tests/unit/export/test_dynamo_hierarchy_report.py diff --git a/docs/commands/export.md b/docs/commands/export.md index e9d9a0176..7ccd18e58 100644 --- a/docs/commands/export.md +++ b/docs/commands/export.md @@ -128,9 +128,14 @@ winml export -m microsoft/resnet-50 -o resnet50_clean.onnx --no-hierarchy hierarchy tags. Pass `--no-dynamo` to select the legacy TorchScript exporter. Shape staticness is independent of the exporter: both default to static shapes and only emit dynamic axes when you ask for them. The QNN-relevant - difference is opset and op decomposition: torch's dynamo op library targets a - minimum opset of 18, so the dynamo path exports at opset 18, whereas the - TorchScript path exports at the configured opset (17 by default). Dynamo also + difference is opset and op decomposition: torch's dynamo op library is built + against a minimum opset of 18. When you request a lower opset (17 by default), + dynamo attempts to down-convert the graph to it; this can succeed (the graph is + saved at the requested opset) or fail (the graph stays at opset 18). `winml + export` reports the opset actually produced and warns when it differs from the + one you requested, so pass `--no-dynamo` if you need a graph natively at opset + 17. The TorchScript path always exports at the configured opset (17 by + default). Dynamo also lowers some ops differently -- for example ResNet's classification head becomes `ReduceMean` + `Reshape` under dynamo but `GlobalAveragePool` + `Flatten` under TorchScript -- which can change or reduce QNN fusion coverage. Prefer diff --git a/docs/concepts/load-and-export.md b/docs/concepts/load-and-export.md index 104707e9e..02cb68492 100644 --- a/docs/concepts/load-and-export.md +++ b/docs/concepts/load-and-export.md @@ -16,7 +16,7 @@ Some community models host custom Python code in their repositories. The loader ## Exporting to ONNX -`winml export` converts the loaded model to ONNX. By default it uses PyTorch's TorchDynamo ONNX exporter (`torch.onnx.export(dynamo=True)`), which records rich per-node module metadata that is used to derive the `winml.hierarchy.*` node tags. Pass `--no-dynamo` to fall back to the legacy TorchScript exporter, which follows actual execution paths and tends to produce compact inference-oriented graphs. Both exporters default to static shapes; the QNN-relevant difference is opset and op decomposition: torch's dynamo op library targets a minimum opset of 18, so the dynamo path exports at opset 18, whereas the TorchScript path exports at the configured opset (17 by default) and lowers some ops differently (e.g. ResNet's head becomes `ReduceMean` + `Reshape` under dynamo but `GlobalAveragePool` + `Flatten` under TorchScript). Because the opset-17 TorchScript graph is what the QNN/NPU toolchain has been validated against, `--no-dynamo` remains the validated choice for QNN/NPU hand exports today. +`winml export` converts the loaded model to ONNX. By default it uses PyTorch's TorchDynamo ONNX exporter (`torch.onnx.export(dynamo=True)`), which records rich per-node module metadata that is used to derive the `winml.hierarchy.*` node tags. Pass `--no-dynamo` to fall back to the legacy TorchScript exporter, which follows actual execution paths and tends to produce compact inference-oriented graphs. Both exporters default to static shapes; the QNN-relevant difference is opset and op decomposition: torch's dynamo op library is built against a minimum opset of 18. When a lower opset is requested (17 by default), dynamo attempts to down-convert to it, which may succeed (the graph is saved at the requested opset) or fail (it stays at opset 18); `winml export` reports the opset actually produced and warns when it differs from the requested value. The TorchScript path always exports at the configured opset (17 by default) and lowers some ops differently (e.g. ResNet's head becomes `ReduceMean` + `Reshape` under dynamo but `GlobalAveragePool` + `Flatten` under TorchScript). Because the opset-17 TorchScript graph is what the QNN/NPU toolchain has been validated against, `--no-dynamo` remains the validated choice for QNN/NPU hand exports today. By default the exporter runs an eight-step process that includes hierarchy tracing and tag injection. The result is an ONNX file enriched with structural metadata that powers downstream features such as per-module benchmarking, inspector views, and optimizer scoping. diff --git a/src/winml/modelkit/core/hierarchy_utils.py b/src/winml/modelkit/core/hierarchy_utils.py index 668ded9c2..f491d7740 100644 --- a/src/winml/modelkit/core/hierarchy_utils.py +++ b/src/winml/modelkit/core/hierarchy_utils.py @@ -14,11 +14,31 @@ from typing import Any +def _nearest_present_ancestor(path: str, hierarchy: dict[str, Any]) -> str: + """Return the longest proper scope-prefix of ``path`` present in ``hierarchy``. + + Scope paths are dotted (``blocks.0.attention.query``); this walks the prefixes + from longest to shortest and returns the first one that is itself a key. When + no intermediate ancestor exists it falls back to the root (``""``), so a sparse + hierarchy (e.g. ``blocks.0`` with no ``blocks`` container entry) still attaches + its subtree to the nearest real parent instead of being dropped from the tree. + """ + segments = path.split(".") + for cut in range(len(segments) - 1, 0, -1): + prefix = ".".join(segments[:cut]) + if prefix in hierarchy: + return prefix + return "" + + def find_immediate_children(parent_path: str, hierarchy: dict[str, Any]) -> list[str]: - """Find immediate children of a path using the WORKING logic from console writer. + """Find immediate children of a path using nearest-present-ancestor nesting. - This is the core hierarchy traversal logic that handles compound patterns - like 'layer.0', 'blocks.1', etc. correctly. + A path is an immediate child of ``parent_path`` when ``parent_path`` is the + longest ancestor scope actually present in ``hierarchy``. This handles both + compound patterns (``layer.0``/``blocks.1``) and sparse hierarchies where an + intermediate container scope (e.g. a ``ModuleList``) never emitted its own + entry, without special-casing any names or class types. Args: parent_path: Parent module path (empty string for root) @@ -40,28 +60,13 @@ def find_immediate_children(parent_path: str, hierarchy: dict[str, Any]) -> list ... }) ["encoder.layer.0", "encoder.layer.1"] # layer.0 is immediate despite having dots """ - if parent_path == "": - # Root case - return sorted([p for p in hierarchy if p and "." not in p]) - - # Non-root case - prefix = parent_path + "." - immediate = [] - - for path in hierarchy: - if not path.startswith(prefix) or path == parent_path: - continue - - suffix = path[len(prefix) :] - - # Check if immediate child - this is the KEY logic that was missing from reports! - if "." not in suffix: - # Simple immediate child (e.g., "attention" under "encoder") - immediate.append(path) - elif suffix.count(".") == 1 and suffix.split(".")[1].isdigit(): - # Compound pattern like layer.0 - treat as immediate child - # This handles ResNet patterns: encoder.layer.0, encoder.layer.1, etc. - immediate.append(path) + immediate = [ + path + for path in hierarchy + if path + and path != parent_path + and _nearest_present_ancestor(path, hierarchy) == parent_path + ] # Custom sort that handles numeric parts properly def sort_key(path: str) -> list[tuple[int, int | str]]: diff --git a/src/winml/modelkit/export/htp/exporter.py b/src/winml/modelkit/export/htp/exporter.py index 09a0c9cc1..77b07f764 100644 --- a/src/winml/modelkit/export/htp/exporter.py +++ b/src/winml/modelkit/export/htp/exporter.py @@ -315,15 +315,18 @@ def export( ) # Output names: dynamo exposes them authoritatively on the ONNX graph; # the TorchScript path derives them from the traced outputs. + output_names: list[str] if self._use_dynamo_hierarchy: output_names = [output.name for output in onnx_model.graph.output] else: traced_outputs = ( self._hierarchy_builder.get_outputs() if self._hierarchy_builder else None ) + # infer_output_names returns list[str] | None; normalize to a list + # so output_names is always list[str] (keeps mypy happy at merge). output_names = ( - infer_output_names(traced_outputs) if traced_outputs is not None else [] - ) + infer_output_names(traced_outputs) if traced_outputs is not None else None + ) or [] # Report the opset the exporter actually produced, not the requested # one. torch's dynamo exporter targets a minimum opset of 18 and does # not always down-convert to a lower requested value (e.g. ResNet stays diff --git a/src/winml/modelkit/export/htp/metadata_writer.py b/src/winml/modelkit/export/htp/metadata_writer.py index 88953322a..d65f27bce 100644 --- a/src/winml/modelkit/export/htp/metadata_writer.py +++ b/src/winml/modelkit/export/htp/metadata_writer.py @@ -18,6 +18,7 @@ from pathlib import Path from typing import TYPE_CHECKING, Any +from ...core.hierarchy_utils import find_immediate_children from ...core.time_utils import format_timestamp_iso from .base_writer import ExportData, ExportStep, StepAwareWriter, step from .metadata_builder import HTPMetadataBuilder @@ -382,6 +383,11 @@ def _build_children_for_parent( ) -> dict[str, Any] | None: """Build children dict for a parent module. + Traversal uses the shared :func:`find_immediate_children` so sparse and + compound scopes (``layer.0``, or an indexed child whose ``ModuleList`` + container never emitted its own entry) nest under the nearest present + ancestor instead of being dropped. + Args: parent_path: Parent module path (e.g., "", "encoder", "encoder.layer.0") flat_hierarchy: Flat dict of all modules @@ -389,63 +395,40 @@ def _build_children_for_parent( Returns: Dict of children or None if no children """ - children = {} - - # Find all direct children of this parent - for path, module_info in flat_hierarchy.items(): - if path == parent_path: - continue # Skip self - - # Check if this is a direct child - if parent_path == "": - # For root, direct children have no dots - if "." not in path: - child_name = path - else: - continue + child_paths = find_immediate_children(parent_path, flat_hierarchy) + if not child_paths: + return None + + # A child's key is normally its class name, but same-class named siblings + # (an attention block's query/key/value, all torch.nn.Linear) would then + # collide and overwrite each other. Detect collisions up front and fold + # the local scope identity into the key for the colliding classes so every + # sibling is serialized. + class_counts: dict[str, int] = {} + for path in child_paths: + class_name = flat_hierarchy[path].class_name + class_counts[class_name] = class_counts.get(class_name, 0) + 1 + + children: dict[str, Any] = {} + for path in child_paths: + module_info = flat_hierarchy[path] + remainder = path[len(parent_path) + 1 :] if parent_path else path + local = remainder.rsplit(".", 1)[-1] + + if local.isdigit(): + # Indexed container child (layer.0) -> Class.N (unchanged). + key = f"{module_info.class_name}.{local}" + elif class_counts[module_info.class_name] > 1: + # Same-class named siblings -> Class.local (query/key/value stay + # distinct instead of collapsing to a single Class key). + key = f"{module_info.class_name}.{local}" else: - # For non-root, check if path starts with parent and has exactly one more segment - if not path.startswith(parent_path + "."): - continue - - # Get the path after the parent - remainder = path[len(parent_path) + 1 :] - - # Check if this is a direct child by counting the depth - # For indexed modules like "layer.0", we need to check the full indexed name - # Split remainder into segments considering indexed modules - segments = [] - current_segment = "" - for part in remainder.split("."): - if current_segment and part.isdigit(): - # This is an index, append to previous segment - current_segment = f"{current_segment}.{part}" - else: - # Start new segment - if current_segment: - segments.append(current_segment) - current_segment = part - if current_segment: - segments.append(current_segment) - - # Direct child has exactly one segment - if len(segments) != 1: - continue - - child_name = segments[0] - - # Determine the key to use - # For indexed modules like layer.0, layer.1, use class_name.index - if "." in child_name and child_name.split(".")[-1].isdigit(): - # This is an indexed module like layer.0 - index = child_name.split(".")[-1] - key = f"{module_info.class_name}.{index}" - else: - # Use class_name as key for consistency with expected structure - # This ensures tests expecting BertEmbeddings, BertEncoder keys work correctly key = module_info.class_name - # Build child structure + # Guarantee uniqueness even for pathological sparse collisions. + if key in children: + key = f"{module_info.class_name}.{remainder}" + child: dict[str, Any] = { "class_name": module_info.class_name, "traced_tag": module_info.traced_tag, @@ -456,7 +439,6 @@ def _build_children_for_parent( if module_info.source: child["source"] = module_info.source - # Recursively build children for this child grandchildren = self._build_children_for_parent(path, flat_hierarchy) if grandchildren: child["children"] = grandchildren diff --git a/tests/unit/export/test_dynamo_hierarchy_report.py b/tests/unit/export/test_dynamo_hierarchy_report.py new file mode 100644 index 000000000..ea09067fc --- /dev/null +++ b/tests/unit/export/test_dynamo_hierarchy_report.py @@ -0,0 +1,123 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Writer-level tests for the dynamo module-hierarchy report. + +These exercise the actual ``MetadataWriter`` tree output (not just the flat map +``build_module_hierarchy`` returns), covering the two failure modes the shared +writers previously had: + + * same-class named siblings (an attention block's query/key/value Linears) + collapsing to a single tree entry because children were keyed by class name, + and + * a sparse hierarchy whose intermediate container scope (a ``ModuleList``) + never emitted its own entry, orphaning the whole subtree from the tree. +""" + +from __future__ import annotations + +from onnx import ModelProto, NodeProto, StringStringEntryProto, helper + +from winml.modelkit.core.hierarchy_utils import find_immediate_children +from winml.modelkit.core.onnx_node_tagger import DynamoMetadataTagger +from winml.modelkit.export.htp.metadata_writer import MetadataWriter +from winml.modelkit.export.htp.step_data import ModuleInfo + + +NAME_SCOPES_KEY = "pkg.torch.onnx.name_scopes" +CLASS_HIERARCHY_KEY = "pkg.torch.onnx.class_hierarchy" + + +def _node(op_type: str, name: str, name_scopes: list[str], class_hierarchy: list[str]) -> NodeProto: + node = helper.make_node(op_type, inputs=["x"], outputs=["y"], name=name) + node.metadata_props.append(StringStringEntryProto(key=NAME_SCOPES_KEY, value=repr(name_scopes))) + node.metadata_props.append( + StringStringEntryProto(key=CLASS_HIERARCHY_KEY, value=repr(class_hierarchy)) + ) + return node + + +def _model(nodes: list[NodeProto]) -> ModelProto: + return helper.make_model(helper.make_graph(nodes, "g", inputs=[], outputs=[])) + + +def _attention_model() -> ModelProto: + """One block with an attention module exposing query/key/value Linears. + + The ``blocks`` ModuleList never appears as its own scope (torch skips + container modules), so the root's only child is the compound ``blocks.0``. + """ + return _model( + [ + _node( + "MatMul", + f"n_{name}", + ["", "blocks.0", "blocks.0.attn", f"blocks.0.attn.{name}", "linear"], + [ + "pkg.Net", + "pkg.Blk", + "pkg.Attention", + "torch.nn.modules.linear.Linear", + "aten.linear.default", + ], + ) + for name in ("query", "key", "value") + ] + ) + + +def _to_module_info(flat: dict[str, dict]) -> dict[str, ModuleInfo]: + return { + scope: ModuleInfo( + class_name=info["class_name"], + traced_tag=info["traced_tag"], + execution_order=info["execution_order"], + ) + for scope, info in flat.items() + } + + +class TestFindImmediateChildrenSparse: + """Nearest-present-ancestor nesting for sparse/compound scopes.""" + + def test_compound_root_child_attaches_to_root(self) -> None: + # "blocks.0" has no "blocks" ancestor entry, so it is a root child and + # its subtree is not dropped. + hierarchy = {"": {}, "blocks.0": {}, "blocks.0.attn": {}} + assert find_immediate_children("", hierarchy) == ["blocks.0"] + assert find_immediate_children("blocks.0", hierarchy) == ["blocks.0.attn"] + + def test_present_container_reparents_index(self) -> None: + # When the "blocks" container IS present, "blocks.0" nests under it. + hierarchy = {"": {}, "blocks": {}, "blocks.0": {}} + assert find_immediate_children("", hierarchy) == ["blocks"] + assert find_immediate_children("blocks", hierarchy) == ["blocks.0"] + + +class TestMetadataWriterTree: + """The persisted MetadataWriter tree preserves every reconstructed module.""" + + def _tree(self) -> dict: + flat = DynamoMetadataTagger().build_module_hierarchy(_attention_model()) + writer = MetadataWriter("unused.json") + return writer._build_hierarchical_modules(_to_module_info(flat)) + + def test_same_class_siblings_all_serialized(self) -> None: + tree = self._tree() + # root -> Blk.0 -> Attention -> {Linear.query, Linear.key, Linear.value} + attn = tree["children"]["Blk.0"]["children"]["Attention"] + linears = attn["children"] + assert set(linears) == {"Linear.query", "Linear.key", "Linear.value"} + scopes = {child["scope"] for child in linears.values()} + assert scopes == { + "blocks.0.attn.query", + "blocks.0.attn.key", + "blocks.0.attn.value", + } + + def test_sparse_root_subtree_present(self) -> None: + tree = self._tree() + # The compound root child "blocks.0" is present, not orphaned. + assert "Blk.0" in tree["children"] + assert tree["children"]["Blk.0"]["scope"] == "blocks.0" From 3c07a4387c4e820359167edb1e96d69ce48090b2 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Wed, 22 Jul 2026 17:36:14 +0800 Subject: [PATCH 09/15] Label dynamo hierarchy by source instead of faking a trace The dynamo export path reconstructs the module hierarchy from ONNX node metadata after export, but the reporting pipeline still described it as a TorchScript forward trace: the console/report claimed "Tracing module execution", the export summary said "Traced modules", and the metadata's execution_steps was set to the reconstructed-module count (a fabricated step total). Docs and the CLI docstring likewise presented forward-hook tracing as the default behaviour. Thread an explicit hierarchy source through the reporting pipeline so writers use source-appropriate wording and never report a trace that did not run: - step_data: add HIERARCHY_SOURCE_TRACE / HIERARCHY_SOURCE_ONNX_METADATA; HierarchyData.execution_steps is now int | None (None = no trace) and carries a source field. - exporter: tag the TorchScript HIERARCHY update as trace and the dynamo one as onnx_metadata, dropping the fabricated execution_steps count. - console/markdown writers: dynamo shows "Reconstructing/Recovered" with no execution-step line; the export summary label switches between "Recovered modules" and "Traced modules" by source. - metadata: record source + the real producer (DynamoMetadataTagger vs TracingHierarchyBuilder) and emit execution_steps=0 for dynamo instead of the module count; document the new source field in the schema. - docs + CLI docstring: distinguish the default dynamo metadata recovery from the --no-dynamo forward-hook trace. Add source-wording tests (console, markdown, and metadata) for both paths. --- docs/commands/export.md | 7 +- docs/concepts/load-and-export.md | 8 +- src/winml/modelkit/commands/export.py | 6 +- .../modelkit/export/htp/console_writer.py | 29 +++-- src/winml/modelkit/export/htp/exporter.py | 6 +- .../export/htp/htp_metadata_schema.json | 13 +- .../export/htp/markdown_report_writer.py | 25 +++- .../modelkit/export/htp/metadata_builder.py | 19 ++- .../modelkit/export/htp/metadata_writer.py | 18 ++- src/winml/modelkit/export/htp/monitor.py | 14 ++- src/winml/modelkit/export/htp/step_data.py | 14 ++- .../export/test_dynamo_hierarchy_report.py | 118 +++++++++++++++++- 12 files changed, 241 insertions(+), 36 deletions(-) diff --git a/docs/commands/export.md b/docs/commands/export.md index 7ccd18e58..f3e86cac3 100644 --- a/docs/commands/export.md +++ b/docs/commands/export.md @@ -37,9 +37,12 @@ $ winml export [options] `winml export` loads the model via Hugging Face `transformers`, then runs the eight-step Hierarchy-preserving Tags Protocol (HTP): model preparation, input -generation, module-hierarchy tracing, TorchScript ONNX export, node-tagger +generation, module-hierarchy recovery, ONNX export, node-tagger creation, per-node tagging, tag injection into ONNX `metadata_props`, and -optional report generation. The hierarchy metadata allows downstream tools to +optional report generation. By default the hierarchy is reconstructed from the +TorchDynamo exporter's per-node module metadata; with `--no-dynamo` it is +captured by tracing the model's forward pass under the legacy TorchScript +exporter instead. The hierarchy metadata allows downstream tools to reason about operators grouped by their originating module rather than flat graph position. When `--no-hierarchy` is specified, hierarchy steps are bypassed and a bare ONNX file is written, useful for third-party tools that do not diff --git a/docs/concepts/load-and-export.md b/docs/concepts/load-and-export.md index 02cb68492..54d891ed4 100644 --- a/docs/concepts/load-and-export.md +++ b/docs/concepts/load-and-export.md @@ -18,7 +18,7 @@ Some community models host custom Python code in their repositories. The loader `winml export` converts the loaded model to ONNX. By default it uses PyTorch's TorchDynamo ONNX exporter (`torch.onnx.export(dynamo=True)`), which records rich per-node module metadata that is used to derive the `winml.hierarchy.*` node tags. Pass `--no-dynamo` to fall back to the legacy TorchScript exporter, which follows actual execution paths and tends to produce compact inference-oriented graphs. Both exporters default to static shapes; the QNN-relevant difference is opset and op decomposition: torch's dynamo op library is built against a minimum opset of 18. When a lower opset is requested (17 by default), dynamo attempts to down-convert to it, which may succeed (the graph is saved at the requested opset) or fail (it stays at opset 18); `winml export` reports the opset actually produced and warns when it differs from the requested value. The TorchScript path always exports at the configured opset (17 by default) and lowers some ops differently (e.g. ResNet's head becomes `ReduceMean` + `Reshape` under dynamo but `GlobalAveragePool` + `Flatten` under TorchScript). Because the opset-17 TorchScript graph is what the QNN/NPU toolchain has been validated against, `--no-dynamo` remains the validated choice for QNN/NPU hand exports today. -By default the exporter runs an eight-step process that includes hierarchy tracing and tag injection. The result is an ONNX file enriched with structural metadata that powers downstream features such as per-module benchmarking, inspector views, and optimizer scoping. +By default the exporter runs an eight-step process that includes hierarchy recovery and tag injection. The result is an ONNX file enriched with structural metadata that powers downstream features such as per-module benchmarking, inspector views, and optimizer scoping. ### Hierarchy tagging in detail @@ -31,7 +31,7 @@ During export the HTP (Hierarchy-preserving Tags Protocol) exporter attaches two #### How tags are built -The exporter registers PyTorch forward hooks on each module. When a module executes, a pre-hook pushes its class name onto a tag stack; the post-hook pops it. This produces hierarchical paths that mirror the PyTorch module tree: +The module path attached to each node is obtained differently depending on the exporter. Under the default TorchDynamo exporter, dynamo records the originating module for every node as ONNX metadata, and the exporter reconstructs the hierarchy directly from that metadata after export — no forward hooks are involved. Under `--no-dynamo`, the TorchScript path instead registers PyTorch forward hooks on each module: when a module executes, a pre-hook pushes its class name onto a tag stack and the post-hook pops it. Either way the result is hierarchical paths that mirror the PyTorch module tree: ```mermaid flowchart LR @@ -42,11 +42,11 @@ flowchart LR E --> F[Tag stack → path] ``` -Only modules that are actually executed during tracing receive tags — unused modules are excluded. For example, `prajjwal1/bert-tiny` has 48 registered modules but only 18 are reached during a forward pass. +Under `--no-dynamo`, only modules that are actually executed during tracing receive tags — unused modules are excluded. For example, `prajjwal1/bert-tiny` has 48 registered modules but only 18 are reached during a forward pass. (Under the default dynamo path the hierarchy instead reflects the modules present in the exported graph's metadata.) #### Concrete example: BERT-tiny -Running `winml export -m prajjwal1/bert-tiny -o model.onnx -v` produces the following hierarchy tree (18 traced modules, 132 ONNX nodes, 100 % coverage): +Running `winml export -m prajjwal1/bert-tiny -o model.onnx -v --no-dynamo` produces the following hierarchy tree (18 traced modules, 132 ONNX nodes, 100 % coverage) — the numbers below are from the TorchScript trace path: ``` BertModel (132 nodes) diff --git a/src/winml/modelkit/commands/export.py b/src/winml/modelkit/commands/export.py index 3449ec858..b8aae23c4 100644 --- a/src/winml/modelkit/commands/export.py +++ b/src/winml/modelkit/commands/export.py @@ -187,8 +187,10 @@ def export( The export process (8 steps): 1. Model Preparation - Load and configure model 2. Input Generation - Generate example inputs - 3. Hierarchy Building - Trace module execution - 4. ONNX Export - Convert to ONNX format (TorchScript by default) + 3. Hierarchy Building - Recover the module hierarchy (from the ONNX metadata + under dynamo, or a forward-execution trace with --no-dynamo) + 4. ONNX Export - Convert to ONNX format (TorchDynamo by default; TorchScript + with --no-dynamo) 5. Node Tagger Creation - Create tagger from hierarchy 6. Node Tagging - Apply hierarchy tags to nodes 7. Tag Injection - Embed tags in ONNX node metadata_props diff --git a/src/winml/modelkit/export/htp/console_writer.py b/src/winml/modelkit/export/htp/console_writer.py index c132feb4b..e0296aecc 100644 --- a/src/winml/modelkit/export/htp/console_writer.py +++ b/src/winml/modelkit/export/htp/console_writer.py @@ -22,6 +22,7 @@ from ...core.hierarchy_utils import build_rich_tree from .base_writer import ExportData, ExportStep, StepAwareWriter, step +from .step_data import HIERARCHY_SOURCE_ONNX_METADATA if TYPE_CHECKING: @@ -171,13 +172,23 @@ def write_hierarchy(self, export_step: ExportStep, data: ExportData) -> int: self.console.print(f"\n🏗️ {self._bold(step_header)}") self.console.print("=" * self.SEPARATOR_LENGTH) - self.console.print("🔍 Tracing module execution with dummy inputs...") - self.console.print( - f"✅ Traced {self._bright_cyan(len(data.hierarchy.hierarchy))} modules in hierarchy" - ) - self.console.print( - f"📊 Total execution steps: {self._bright_cyan(data.hierarchy.execution_steps)}" - ) + if data.hierarchy.source == HIERARCHY_SOURCE_ONNX_METADATA: + # Dynamo path: no forward trace ran; the hierarchy is reconstructed + # from ONNX node metadata, so avoid trace/execution-step wording. + self.console.print("🔍 Reconstructing module hierarchy from ONNX node metadata...") + self.console.print( + f"✅ Recovered {self._bright_cyan(len(data.hierarchy.hierarchy))} modules " + "in hierarchy" + ) + else: + self.console.print("🔍 Tracing module execution with dummy inputs...") + self.console.print( + f"✅ Traced {self._bright_cyan(len(data.hierarchy.hierarchy))} modules in hierarchy" + ) + if data.hierarchy.execution_steps is not None: + self.console.print( + f"📊 Total execution steps: {self._bright_cyan(data.hierarchy.execution_steps)}" + ) # Build and display hierarchy tree if data.hierarchy.hierarchy: @@ -335,9 +346,7 @@ def _build_truncated_tree(self, source_tree: Tree, target_tree: Tree, max_lines: line_count = 1 # Start with root # Helper to add nodes up to limit - def add_nodes_to_limit( - source_children: Any, target_parent: Any, current_count: int - ) -> int: + def add_nodes_to_limit(source_children: Any, target_parent: Any, current_count: int) -> int: count = current_count for child in source_children: if count >= max_lines: diff --git a/src/winml/modelkit/export/htp/exporter.py b/src/winml/modelkit/export/htp/exporter.py index 77b07f764..01719341f 100644 --- a/src/winml/modelkit/export/htp/exporter.py +++ b/src/winml/modelkit/export/htp/exporter.py @@ -39,6 +39,7 @@ from .base_writer import ExportStep from .hierarchy import TracingHierarchyBuilder from .monitor import HTPExportMonitor +from .step_data import HIERARCHY_SOURCE_ONNX_METADATA, HIERARCHY_SOURCE_TRACE if TYPE_CHECKING: @@ -266,6 +267,7 @@ def export( ExportStep.HIERARCHY, hierarchy=self._hierarchy_data, execution_steps=execution_steps, + source=HIERARCHY_SOURCE_TRACE, ) # Step 4: ONNX Export @@ -298,13 +300,15 @@ def export( # instead of a TorchScript trace. Recover it now that the graph exists # and issue the deferred Step 3 update so the report/console/metadata # module tree and the hierarchy_modules stat reflect the real modules. + # No forward trace runs on this path, so execution_steps stays unset + # (None) rather than reporting the module count as a fake step total. if self._use_dynamo_hierarchy: assert isinstance(self._node_tagger, DynamoMetadataTagger) self._hierarchy_data = self._node_tagger.build_module_hierarchy(onnx_model) monitor.update( ExportStep.HIERARCHY, hierarchy=self._hierarchy_data, - execution_steps=len(self._hierarchy_data), + source=HIERARCHY_SOURCE_ONNX_METADATA, ) # Update monitor with ONNX export info diff --git a/src/winml/modelkit/export/htp/htp_metadata_schema.json b/src/winml/modelkit/export/htp/htp_metadata_schema.json index 624705392..dcf027e79 100644 --- a/src/winml/modelkit/export/htp/htp_metadata_schema.json +++ b/src/winml/modelkit/export/htp/htp_metadata_schema.json @@ -141,20 +141,29 @@ "type": "object" }, "TracingInfo": { - "description": "Tracing execution information", + "description": "Hierarchy build information (execution trace or ONNX-metadata recovery)", "properties": { "builder": { "description": "Hierarchy builder used", "type": "string", "default": "TracingHierarchyBuilder" }, + "source": { + "description": "How the hierarchy was obtained: a forward-execution trace (TorchScript) or reconstruction from ONNX node metadata (dynamo)", + "type": "string", + "enum": [ + "trace", + "onnx_metadata" + ], + "default": "trace" + }, "modules_traced": { "description": "Number of modules traced during execution", "type": "integer", "minimum": 0 }, "execution_steps": { - "description": "Number of execution steps", + "description": "Number of execution steps (0 when the hierarchy was recovered from ONNX metadata rather than an execution trace)", "type": "integer", "minimum": 0 }, diff --git a/src/winml/modelkit/export/htp/markdown_report_writer.py b/src/winml/modelkit/export/htp/markdown_report_writer.py index 35a42b687..a156e356f 100644 --- a/src/winml/modelkit/export/htp/markdown_report_writer.py +++ b/src/winml/modelkit/export/htp/markdown_report_writer.py @@ -23,6 +23,7 @@ ) from ...core.time_utils import format_timestamp_iso from .base_writer import ExportData, ExportStep, StepAwareWriter, step +from .step_data import HIERARCHY_SOURCE_ONNX_METADATA class MarkdownReportWriter(StepAwareWriter): @@ -254,11 +255,18 @@ def _write_hierarchy_section(self, data: ExportData) -> None: self.doc.add_heading("✅ Step 3/6: Hierarchy Building", level=3) if data.hierarchy: - items = [ - f"**Modules Traced**: {len(data.hierarchy.hierarchy)}", - f"**Execution Steps**: {data.hierarchy.execution_steps}", - "**Status**: Module hierarchy successfully traced", - ] + if data.hierarchy.source == HIERARCHY_SOURCE_ONNX_METADATA: + items = [ + f"**Modules Reconstructed**: {len(data.hierarchy.hierarchy)}", + "**Source**: ONNX node metadata (dynamo, no execution trace)", + "**Status**: Module hierarchy reconstructed from ONNX metadata", + ] + else: + items = [ + f"**Modules Traced**: {len(data.hierarchy.hierarchy)}", + f"**Execution Steps**: {data.hierarchy.execution_steps}", + "**Status**: Module hierarchy successfully traced", + ] self.doc.add_unordered_list(items) # Add module hierarchy tree preview @@ -527,9 +535,14 @@ def _write_summary_section(self, data: ExportData) -> None: data.node_tagging.tagging_stats.get("empty_tags", 0) if data.node_tagging else 0 ) + modules_label = ( + "Recovered Modules" + if data.hierarchy and data.hierarchy.source == HIERARCHY_SOURCE_ONNX_METADATA + else "Traced Modules" + ) stats = [ f"**Hierarchy Modules**: {total_modules}", - f"**Traced Modules**: {traced_modules}/{total_modules}", + f"**{modules_label}**: {traced_modules}/{total_modules}", f"**ONNX Nodes**: {onnx_nodes}", f"**Tagged Nodes**: {tagged_nodes} ({coverage:.1f}%)", f"**Empty Tags**: {empty_tags}", diff --git a/src/winml/modelkit/export/htp/metadata_builder.py b/src/winml/modelkit/export/htp/metadata_builder.py index fd0698faf..b12fe72fd 100644 --- a/src/winml/modelkit/export/htp/metadata_builder.py +++ b/src/winml/modelkit/export/htp/metadata_builder.py @@ -15,6 +15,7 @@ from typing import Any from . import __spec_version__ as htp_version +from .step_data import HIERARCHY_SOURCE_TRACE @dataclass @@ -52,9 +53,10 @@ def to_dict(self) -> dict[str, Any]: @dataclass class TracingInfo: - """Tracing execution information.""" + """Hierarchy build information (execution trace or ONNX-metadata recovery).""" builder: str = "TracingHierarchyBuilder" + source: str = HIERARCHY_SOURCE_TRACE modules_traced: int = 0 execution_steps: int = 0 model_type: str | None = None @@ -205,9 +207,17 @@ def with_tracing_info( task: str | None = None, inputs: dict[str, dict[str, Any]] | None = None, outputs: list[str] | None = None, + source: str = HIERARCHY_SOURCE_TRACE, + builder: str | None = None, ) -> HTPMetadataBuilder: - """Set tracing information.""" - self._tracing_info = TracingInfo( + """Set hierarchy build information. + + ``source`` records how the hierarchy was obtained; ``builder`` names the + component that produced it (defaults to the TorchScript tracer when not + given, so the dynamo path passes its reconstructor explicitly). + """ + info = TracingInfo( + source=source, modules_traced=modules_traced, execution_steps=execution_steps, model_type=model_type, @@ -215,6 +225,9 @@ def with_tracing_info( inputs=inputs, outputs=outputs, ) + if builder is not None: + info.builder = builder + self._tracing_info = info return self def with_modules(self, modules: dict[str, dict[str, Any]]) -> HTPMetadataBuilder: diff --git a/src/winml/modelkit/export/htp/metadata_writer.py b/src/winml/modelkit/export/htp/metadata_writer.py index d65f27bce..4fd7520c1 100644 --- a/src/winml/modelkit/export/htp/metadata_writer.py +++ b/src/winml/modelkit/export/htp/metadata_writer.py @@ -22,6 +22,7 @@ from ...core.time_utils import format_timestamp_iso from .base_writer import ExportData, ExportStep, StepAwareWriter, step from .metadata_builder import HTPMetadataBuilder +from .step_data import HIERARCHY_SOURCE_ONNX_METADATA if TYPE_CHECKING: @@ -133,20 +134,33 @@ def write_hierarchy(self, export_step: ExportStep, data: ExportData) -> int: # Get input info from previous step data input_data = self._steps_data.get("input_generation", {}) + # The dynamo path reconstructs the hierarchy from ONNX metadata instead + # of a forward trace, so there is no execution-step count and a different + # producer. Record the source explicitly and avoid a fabricated count. + source = data.hierarchy.source + if source == HIERARCHY_SOURCE_ONNX_METADATA: + builder = "DynamoMetadataTagger" + else: + builder = "TracingHierarchyBuilder" + execution_steps = data.hierarchy.execution_steps or 0 + self.builder.with_tracing_info( modules_traced=len(data.hierarchy.hierarchy), - execution_steps=data.hierarchy.execution_steps, + execution_steps=execution_steps, model_type=input_data.get("model_type"), task=input_data.get("task"), inputs=input_data.get("inputs"), + source=source, + builder=builder, ) # Record step completion self._steps_data["hierarchy_building"] = { "completed": True, "timestamp": format_timestamp_iso(data.get_step_timestamp(export_step)) or "", + "source": source, "modules_traced": len(data.hierarchy.hierarchy), - "execution_steps": data.hierarchy.execution_steps, + "execution_steps": execution_steps, } return 1 diff --git a/src/winml/modelkit/export/htp/monitor.py b/src/winml/modelkit/export/htp/monitor.py index 2cf94f549..218c55d5b 100644 --- a/src/winml/modelkit/export/htp/monitor.py +++ b/src/winml/modelkit/export/htp/monitor.py @@ -23,6 +23,8 @@ from .markdown_report_writer import MarkdownReportWriter from .metadata_writer import MetadataWriter from .step_data import ( + HIERARCHY_SOURCE_ONNX_METADATA, + HIERARCHY_SOURCE_TRACE, HierarchyData, InputGenData, ModelPrepData, @@ -162,7 +164,8 @@ def _update_step_data(self, step: ExportStep, kwargs: dict) -> None: self.data.hierarchy = HierarchyData( hierarchy=hierarchy, - execution_steps=kwargs.get("execution_steps", 0), + execution_steps=kwargs.get("execution_steps"), + source=kwargs.get("source", HIERARCHY_SOURCE_TRACE), module_list=kwargs.get("module_list", []), ) @@ -210,8 +213,15 @@ def _print_summary(self) -> None: console.print("📊 Export Summary:") console.print(f" • Total time: [bold cyan]{total_time:.2f}s[/bold cyan]") console.print(f" • Hierarchy modules: [bold cyan]{total_modules}[/bold cyan]") + # "Traced" only fits the TorchScript path; the dynamo path recovers modules + # from ONNX metadata, so label the count by how it was actually obtained. + modules_label = ( + "Recovered modules" + if self.data.hierarchy and self.data.hierarchy.source == HIERARCHY_SOURCE_ONNX_METADATA + else "Traced modules" + ) console.print( - f" • Traced modules: [bold cyan]{traced_modules}/{total_modules}[/bold cyan]" + f" • {modules_label}: [bold cyan]{traced_modules}/{total_modules}[/bold cyan]" ) console.print(f" • ONNX nodes: [bold cyan]{nodes}[/bold cyan]") console.print( diff --git a/src/winml/modelkit/export/htp/step_data.py b/src/winml/modelkit/export/htp/step_data.py index de4f79461..b596cf1dc 100644 --- a/src/winml/modelkit/export/htp/step_data.py +++ b/src/winml/modelkit/export/htp/step_data.py @@ -15,6 +15,15 @@ from typing import Any +# How a module hierarchy was obtained. The TorchScript path records a real +# forward-execution trace (hook-based), so execution steps are meaningful. The +# Dynamo path has no such trace: the hierarchy is reconstructed from the ONNX +# node metadata after export, so there is no execution-step count. Writers use +# this to pick source-appropriate wording instead of always claiming a trace. +HIERARCHY_SOURCE_TRACE = "trace" +HIERARCHY_SOURCE_ONNX_METADATA = "onnx_metadata" + + def _timestamp_field() -> Any: """Create a timestamp field that captures time at instance creation. @@ -68,7 +77,10 @@ class HierarchyData: """Data for hierarchy building step.""" hierarchy: dict[str, ModuleInfo] - execution_steps: int + # None when the hierarchy source has no execution trace (e.g. dynamo, which + # reconstructs modules from ONNX metadata). Do not fabricate a count. + execution_steps: int | None = None + source: str = HIERARCHY_SOURCE_TRACE module_list: list[tuple[str, str]] = field(default_factory=list) timestamp: float = _timestamp_field() diff --git a/tests/unit/export/test_dynamo_hierarchy_report.py b/tests/unit/export/test_dynamo_hierarchy_report.py index ea09067fc..d2b51a4a3 100644 --- a/tests/unit/export/test_dynamo_hierarchy_report.py +++ b/tests/unit/export/test_dynamo_hierarchy_report.py @@ -17,12 +17,25 @@ from __future__ import annotations +import io + from onnx import ModelProto, NodeProto, StringStringEntryProto, helper +from rich.console import Console from winml.modelkit.core.hierarchy_utils import find_immediate_children from winml.modelkit.core.onnx_node_tagger import DynamoMetadataTagger +from winml.modelkit.export.htp.base_writer import ExportData, ExportStep +from winml.modelkit.export.htp.console_writer import ConsoleWriter +from winml.modelkit.export.htp.markdown_report_writer import MarkdownReportWriter from winml.modelkit.export.htp.metadata_writer import MetadataWriter -from winml.modelkit.export.htp.step_data import ModuleInfo +from winml.modelkit.export.htp.monitor import HTPExportMonitor +from winml.modelkit.export.htp.step_data import ( + HIERARCHY_SOURCE_ONNX_METADATA, + HIERARCHY_SOURCE_TRACE, + HierarchyData, + ModelPrepData, + ModuleInfo, +) NAME_SCOPES_KEY = "pkg.torch.onnx.name_scopes" @@ -121,3 +134,106 @@ def test_sparse_root_subtree_present(self) -> None: # The compound root child "blocks.0" is present, not orphaned. assert "Blk.0" in tree["children"] assert tree["children"]["Blk.0"]["scope"] == "blocks.0" + + +def _hierarchy_data(source: str, execution_steps: int | None) -> HierarchyData: + return HierarchyData( + hierarchy={"": ModuleInfo(class_name="Net", traced_tag="/Net")}, + execution_steps=execution_steps, + source=source, + ) + + +def _console_output(hierarchy: HierarchyData) -> str: + buf = io.StringIO() + console = Console(file=buf, width=120, no_color=True) + writer = ConsoleWriter(console=console) + writer.write(ExportStep.HIERARCHY, ExportData(hierarchy=hierarchy)) + return buf.getvalue() + + +class TestHierarchySourceWording: + """Console/metadata wording must match how the hierarchy was obtained.""" + + def test_dynamo_source_uses_reconstruction_wording(self) -> None: + out = _console_output(_hierarchy_data(HIERARCHY_SOURCE_ONNX_METADATA, None)) + assert "Reconstructing module hierarchy from ONNX" in out + # No forward trace ran, so no trace/execution-step claims. + assert "Tracing module execution" not in out + assert "execution steps" not in out.lower() + + def test_trace_source_uses_trace_wording(self) -> None: + out = _console_output(_hierarchy_data(HIERARCHY_SOURCE_TRACE, 42)) + assert "Tracing module execution" in out + assert "Total execution steps" in out + assert "42" in out + + def test_metadata_records_dynamo_source_without_fake_step_count(self) -> None: + writer = MetadataWriter("unused.json") + writer.write( + ExportStep.HIERARCHY, + ExportData(hierarchy=_hierarchy_data(HIERARCHY_SOURCE_ONNX_METADATA, None)), + ) + info = writer.builder._tracing_info + assert info.source == HIERARCHY_SOURCE_ONNX_METADATA + assert info.builder == "DynamoMetadataTagger" + # Module count must not be reported as an execution-step total. + assert info.execution_steps == 0 + + def test_metadata_records_trace_source_and_steps(self) -> None: + writer = MetadataWriter("unused.json") + writer.write( + ExportStep.HIERARCHY, + ExportData(hierarchy=_hierarchy_data(HIERARCHY_SOURCE_TRACE, 7)), + ) + info = writer.builder._tracing_info + assert info.source == HIERARCHY_SOURCE_TRACE + assert info.builder == "TracingHierarchyBuilder" + assert info.execution_steps == 7 + + +def _summary_data(source: str) -> ExportData: + return ExportData( + hierarchy=HierarchyData( + hierarchy={"": ModuleInfo(class_name="Net", traced_tag="/Net")}, + source=source, + ), + model_prep=ModelPrepData(model_class="Net", total_modules=3, total_parameters=0), + ) + + +class TestExportSummaryModulesLabel: + """The export summary must not claim a trace when modules were recovered.""" + + def _monitor_summary(self, source: str) -> str: + monitor = HTPExportMonitor("unused.onnx", verbose=True) + buf = io.StringIO() + monitor.console = Console(file=buf, width=120, no_color=True) + monitor.data = _summary_data(source) + monitor._print_summary() + return buf.getvalue() + + def test_console_summary_dynamo_says_recovered(self) -> None: + out = self._monitor_summary(HIERARCHY_SOURCE_ONNX_METADATA) + assert "Recovered modules:" in out + assert "Traced modules:" not in out + + def test_console_summary_trace_says_traced(self) -> None: + out = self._monitor_summary(HIERARCHY_SOURCE_TRACE) + assert "Traced modules:" in out + assert "Recovered modules:" not in out + + def _markdown_summary(self, source: str) -> str: + writer = MarkdownReportWriter("unused.md") + writer._write_summary_section(_summary_data(source)) + return str(writer.doc) + + def test_markdown_summary_dynamo_says_recovered(self) -> None: + out = self._markdown_summary(HIERARCHY_SOURCE_ONNX_METADATA) + assert "Recovered Modules" in out + assert "Traced Modules" not in out + + def test_markdown_summary_trace_says_traced(self) -> None: + out = self._markdown_summary(HIERARCHY_SOURCE_TRACE) + assert "Traced Modules" in out + assert "Recovered Modules" not in out From 5fcc345ecdcb24b8af9eecd2912161ef0dd40304 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Wed, 22 Jul 2026 18:07:33 +0800 Subject: [PATCH 10/15] Preserve exporter choice in submodule configs --- src/winml/modelkit/config/build.py | 5 +++++ tests/unit/config/test_build.py | 23 +++++++++++++++++++++++ 2 files changed, 28 insertions(+) diff --git a/src/winml/modelkit/config/build.py b/src/winml/modelkit/config/build.py index 07f3b2091..4722a2c14 100644 --- a/src/winml/modelkit/config/build.py +++ b/src/winml/modelkit/config/build.py @@ -1061,6 +1061,11 @@ def _input_name(i: int) -> str: input_tensors=input_tensors or None, output_tensors=output_tensors or None, dynamic_axes={}, # Static shapes for submodules + dynamo=( + parent_config.export.dynamo + if parent_config.export is not None + else WinMLExportConfig().dynamo + ), # opset_version and batch_size use dataclass defaults from WinMLExportConfig ), optim=copy.deepcopy(parent_config.optim), diff --git a/tests/unit/config/test_build.py b/tests/unit/config/test_build.py index 3cd3cd186..ca986cc32 100644 --- a/tests/unit/config/test_build.py +++ b/tests/unit/config/test_build.py @@ -872,6 +872,29 @@ def test_single_input_single_output(self, parent_config: WinMLBuildConfig) -> No assert len(result.export.output_tensors) == 1 assert result.export.output_tensors[0].name == "output_0" + @pytest.mark.parametrize("dynamo", [False, True]) + def test_preserves_parent_exporter_choice( + self, + parent_config: WinMLBuildConfig, + dynamo: bool, + ) -> None: + """An explicit parent exporter choice survives submodule specialization.""" + assert parent_config.export is not None + parent_config.export.dynamo = dynamo + sub_info = SubmoduleInfo( + class_name="Linear", + module_path="encoder.proj", + input_shapes=[[1, 8]], + output_shapes=[[1, 8]], + input_dtypes=["float32"], + output_dtypes=["float32"], + ) + + result = _build_submodule_config(sub_info, parent_config) + + assert result.export is not None + assert result.export.dynamo is dynamo + def test_multi_input(self, parent_config: WinMLBuildConfig) -> None: """SubmoduleInfo with 2 input_shapes creates 2 InputTensorSpec.""" sub_info = SubmoduleInfo( From e6c2a201ddf0eebe6d3b3b25a61fe31c4c1fc52e Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Wed, 22 Jul 2026 18:09:29 +0800 Subject: [PATCH 11/15] Warn when dynamo hierarchy metadata is unavailable --- src/winml/modelkit/export/htp/exporter.py | 15 ++++++++-- .../export/test_dynamo_hierarchy_report.py | 29 ++++++++++++++++++- 2 files changed, 41 insertions(+), 3 deletions(-) diff --git a/src/winml/modelkit/export/htp/exporter.py b/src/winml/modelkit/export/htp/exporter.py index 01719341f..249cc7964 100644 --- a/src/winml/modelkit/export/htp/exporter.py +++ b/src/winml/modelkit/export/htp/exporter.py @@ -303,8 +303,7 @@ def export( # No forward trace runs on this path, so execution_steps stays unset # (None) rather than reporting the module count as a fake step total. if self._use_dynamo_hierarchy: - assert isinstance(self._node_tagger, DynamoMetadataTagger) - self._hierarchy_data = self._node_tagger.build_module_hierarchy(onnx_model) + self._hierarchy_data = self._recover_dynamo_hierarchy(onnx_model) monitor.update( ExportStep.HIERARCHY, hierarchy=self._hierarchy_data, @@ -675,6 +674,18 @@ def _initialize_node_tagger(self, enable_operation_fallback: bool) -> None: self._hierarchy_data, enable_operation_fallback=enable_operation_fallback ) + def _recover_dynamo_hierarchy(self, onnx_model: onnx.ModelProto) -> dict[str, dict[str, Any]]: + """Recover module hierarchy from dynamo metadata and surface total degradation.""" + assert isinstance(self._node_tagger, DynamoMetadataTagger) + hierarchy = self._node_tagger.build_module_hierarchy(onnx_model) + if onnx_model.graph.node and not hierarchy: + logger.warning( + "Dynamo export produced no usable module hierarchy metadata; " + "hierarchy tags will use the model-root fallback. " + "Pass --no-dynamo to build hierarchy from a TorchScript execution trace." + ) + return hierarchy + def _apply_hierarchy_tags(self, onnx_model: onnx.ModelProto) -> None: """Tag nodes internally.""" assert self._node_tagger is not None, ( diff --git a/tests/unit/export/test_dynamo_hierarchy_report.py b/tests/unit/export/test_dynamo_hierarchy_report.py index d2b51a4a3..882405b90 100644 --- a/tests/unit/export/test_dynamo_hierarchy_report.py +++ b/tests/unit/export/test_dynamo_hierarchy_report.py @@ -18,17 +18,18 @@ from __future__ import annotations import io +import logging from onnx import ModelProto, NodeProto, StringStringEntryProto, helper from rich.console import Console from winml.modelkit.core.hierarchy_utils import find_immediate_children from winml.modelkit.core.onnx_node_tagger import DynamoMetadataTagger +from winml.modelkit.export.htp import HTPExporter, HTPExportMonitor from winml.modelkit.export.htp.base_writer import ExportData, ExportStep from winml.modelkit.export.htp.console_writer import ConsoleWriter from winml.modelkit.export.htp.markdown_report_writer import MarkdownReportWriter from winml.modelkit.export.htp.metadata_writer import MetadataWriter -from winml.modelkit.export.htp.monitor import HTPExportMonitor from winml.modelkit.export.htp.step_data import ( HIERARCHY_SOURCE_ONNX_METADATA, HIERARCHY_SOURCE_TRACE, @@ -192,6 +193,32 @@ def test_metadata_records_trace_source_and_steps(self) -> None: assert info.execution_steps == 7 +class TestDynamoHierarchyRecoveryWarning: + """A completely missing dynamo hierarchy is visible but non-fatal.""" + + def test_warns_when_nonempty_graph_has_no_usable_metadata(self, caplog) -> None: + exporter = HTPExporter() + exporter._node_tagger = DynamoMetadataTagger() + model = _model([helper.make_node("Identity", ["x"], ["y"], name="bare")]) + + with caplog.at_level(logging.WARNING): + hierarchy = exporter._recover_dynamo_hierarchy(model) + + assert hierarchy == {} + assert "no usable module hierarchy metadata" in caplog.text + assert "--no-dynamo" in caplog.text + + def test_does_not_warn_when_hierarchy_is_recovered(self, caplog) -> None: + exporter = HTPExporter() + exporter._node_tagger = DynamoMetadataTagger() + + with caplog.at_level(logging.WARNING): + hierarchy = exporter._recover_dynamo_hierarchy(_attention_model()) + + assert hierarchy + assert "no usable module hierarchy metadata" not in caplog.text + + def _summary_data(source: str) -> ExportData: return ExportData( hierarchy=HierarchyData( From dc167b42fde7a21a4717117c7fb57ba5579d5007 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Wed, 22 Jul 2026 18:11:22 +0800 Subject: [PATCH 12/15] Make opset mismatch warnings source-aware --- src/winml/modelkit/export/htp/exporter.py | 34 ++++++++++++++++---- tests/unit/export/test_htp_exporter_stats.py | 32 ++++++++++++++++++ 2 files changed, 59 insertions(+), 7 deletions(-) diff --git a/src/winml/modelkit/export/htp/exporter.py b/src/winml/modelkit/export/htp/exporter.py index 249cc7964..f95325262 100644 --- a/src/winml/modelkit/export/htp/exporter.py +++ b/src/winml/modelkit/export/htp/exporter.py @@ -336,15 +336,11 @@ def export( # at 18), so echoing export_config.opset_version would misreport the # real graph. Fall back to the request only if the value can't be read. actual_opset = self._read_default_opset(output_path) - if actual_opset is not None and actual_opset != export_config.opset_version: - logger.warning( - "Requested opset %d but the exporter produced opset %d. " - "torch's dynamo exporter targets a minimum opset of 18 and " - "cannot always lower to a smaller requested version; pass " - "--no-dynamo for a natively opset-%d graph.", + if actual_opset is not None: + self._warn_on_opset_mismatch( export_config.opset_version, actual_opset, - export_config.opset_version, + dynamo=self._use_dynamo_hierarchy, ) monitor.update( ExportStep.ONNX_EXPORT, @@ -447,6 +443,30 @@ def _read_default_opset(output_path: str) -> int | None: return opset.version return None + @staticmethod + def _warn_on_opset_mismatch( + requested_opset: int, + actual_opset: int, + *, + dynamo: bool, + ) -> None: + """Warn accurately when an exporter does not honor the requested opset.""" + if requested_opset == actual_opset: + return + if dynamo and actual_opset > requested_opset: + logger.warning( + "Requested opset %d, but dynamo could not lower the exported model " + "and kept opset %d. Pass --no-dynamo if the exact lower opset is required.", + requested_opset, + actual_opset, + ) + return + logger.warning( + "Requested opset %d but the exporter produced opset %d.", + requested_opset, + actual_opset, + ) + def _verify_onnx_export( self, output_path: str, diff --git a/tests/unit/export/test_htp_exporter_stats.py b/tests/unit/export/test_htp_exporter_stats.py index 403d6cef7..5e719d127 100644 --- a/tests/unit/export/test_htp_exporter_stats.py +++ b/tests/unit/export/test_htp_exporter_stats.py @@ -6,9 +6,11 @@ from __future__ import annotations +import logging from typing import TYPE_CHECKING from unittest.mock import MagicMock +import pytest from onnx import TensorProto, helper, save from winml.modelkit.export.htp import HTPExporter @@ -103,3 +105,33 @@ def test_returns_none_without_default_domain(self, tmp_path: Path) -> None: def test_returns_none_for_unreadable_path(self, tmp_path: Path) -> None: assert HTPExporter._read_default_opset(str(tmp_path / "missing.onnx")) is None + + +class TestHTPExporterOpsetMismatchWarning: + """Mismatch guidance must match the exporter and conversion direction.""" + + def test_failed_dynamo_downconversion_recommends_no_dynamo(self, caplog) -> None: + with caplog.at_level(logging.WARNING): + HTPExporter._warn_on_opset_mismatch(17, 18, dynamo=True) + + assert "could not lower" in caplog.text + assert "--no-dynamo" in caplog.text + + @pytest.mark.parametrize( + ("requested", "actual", "dynamo"), + [(18, 17, True), (17, 18, False)], + ) + def test_other_mismatches_use_generic_wording( + self, + caplog, + requested: int, + actual: int, + dynamo: bool, + ) -> None: + with caplog.at_level(logging.WARNING): + HTPExporter._warn_on_opset_mismatch(requested, actual, dynamo=dynamo) + + assert f"Requested opset {requested}" in caplog.text + assert f"produced opset {actual}" in caplog.text + assert "--no-dynamo" not in caplog.text + assert "could not lower" not in caplog.text From 18b7d72f4fc58249168cb21621008c94c2300d8b Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Wed, 22 Jul 2026 18:12:29 +0800 Subject: [PATCH 13/15] Align HTP terminology with dynamo hierarchy recovery --- docs/concepts/load-and-export.md | 10 +++++----- src/winml/modelkit/export/htp/__init__.py | 20 ++++++++----------- src/winml/modelkit/export/htp/exporter.py | 4 +--- .../export/htp/htp_metadata_schema.json | 7 ++++++- 4 files changed, 20 insertions(+), 21 deletions(-) diff --git a/docs/concepts/load-and-export.md b/docs/concepts/load-and-export.md index 54d891ed4..e5c96ca58 100644 --- a/docs/concepts/load-and-export.md +++ b/docs/concepts/load-and-export.md @@ -2,13 +2,13 @@ The first stage of the winml-cli pipeline is the most deterministic: bring a model into memory and convert it to ONNX. Everything that follows — optimization, quantization, compilation — operates on that ONNX artifact. A well-exported graph with accurate metadata travels cleanly through the rest of the pipeline without requiring patching or re-export. -Loading is an internal operation: the loader module resolves model provenance, selects the right HuggingFace model class, and prepares the weights for tracing. The `winml export` command is the surface users interact with directly. +Loading is an internal operation: the loader module resolves model provenance, selects the right HuggingFace model class, and prepares the weights for export. The `winml export` command is the surface users interact with directly. ## Loading a model When you point winml-cli at a model identifier, the internal loader resolves it in one of two ways. If the identifier looks like a HuggingFace Hub path (e.g., `prajjwal1/bert-tiny`), the loader downloads the model weights and configuration to the standard HuggingFace cache at `~/.cache/huggingface`. Subsequent runs are served from that cache without re-downloading. If the identifier is a path to a local PyTorch checkpoint directory, the loader reads it directly without network access. -In both cases the loader auto-detects the task — image classification, text feature extraction, and so on — and selects a corresponding HuggingFace model class. The result is a PyTorch model object ready for tracing. +In both cases the loader auto-detects the task — image classification, text feature extraction, and so on — and selects a corresponding HuggingFace model class. The result is a PyTorch model object ready for export. Before committing to a full export you can verify that the loader resolved everything correctly with `winml inspect`. It prints the detected task, the HuggingFace model class, the export configuration, and the WinML inference class — all without downloading weights. Add `--hierarchy` to reconstruct the PyTorch module tree from random-weight tracing. @@ -31,7 +31,7 @@ During export the HTP (Hierarchy-preserving Tags Protocol) exporter attaches two #### How tags are built -The module path attached to each node is obtained differently depending on the exporter. Under the default TorchDynamo exporter, dynamo records the originating module for every node as ONNX metadata, and the exporter reconstructs the hierarchy directly from that metadata after export — no forward hooks are involved. Under `--no-dynamo`, the TorchScript path instead registers PyTorch forward hooks on each module: when a module executes, a pre-hook pushes its class name onto a tag stack and the post-hook pops it. Either way the result is hierarchical paths that mirror the PyTorch module tree: +The module path attached to each node is obtained differently depending on the exporter. Under the default TorchDynamo exporter, dynamo records originating-module information as ONNX metadata, and the exporter reconstructs the hierarchy directly from that metadata after export — no forward hooks are involved. Under `--no-dynamo`, the TorchScript path instead registers PyTorch forward hooks on each module: when a module executes, a pre-hook pushes its class name onto a tag stack and the post-hook pops it. Either way the result is hierarchical paths that mirror the PyTorch module tree: ```mermaid flowchart LR @@ -75,7 +75,7 @@ Each ONNX node gets its tag from the module it belongs to. Here are a few exampl #### Node-to-module mapping -After the ONNX graph is produced by `torch.onnx.export`, a 4-priority system assigns each ONNX node to the closest matching module: +Under the default dynamo path, each node is assigned directly from its `pkg.torch.onnx.name_scopes` and `pkg.torch.onnx.class_hierarchy` metadata; nodes without usable metadata receive the model-root fallback. Under `--no-dynamo`, after the ONNX graph is produced by `torch.onnx.export`, the legacy tagger instead uses a 4-priority system to assign each ONNX node to the closest traced module: 1. **Direct match** (61 %) — the node's scope name maps exactly to a traced module. 2. **Parent match** (24 %) — walk up the scope hierarchy until a traced module is found. @@ -100,7 +100,7 @@ These I/O specs enable tools like `winml perf` to generate correct dummy inputs Alongside the `.onnx` file, the exporter writes a `*_htp_metadata.json` sidecar containing: - **`nodes`** — complete mapping of every ONNX node name → hierarchy tag -- **`modules`** — traced module information (class name, tag, execution order) +- **`modules`** — hierarchy module information (class name, tag, discovery order) - **`statistics`** — export time, node counts, coverage percentage - **`outputs`** — I/O tensor specifications diff --git a/src/winml/modelkit/export/htp/__init__.py b/src/winml/modelkit/export/htp/__init__.py index 2ea63c8cb..97f0b2e10 100644 --- a/src/winml/modelkit/export/htp/__init__.py +++ b/src/winml/modelkit/export/htp/__init__.py @@ -2,21 +2,17 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. # -------------------------------------------------------------------------- -"""HTP (Hierarchical Trace-and-Project) Strategy with IO/ABC Architecture. +"""HTP (Hierarchical Trace-and-Project) strategy with IO/ABC architecture. -This strategy uses execution tracing with PyTorch hooks to capture module context -during forward pass, then projects this onto ONNX operations. +HTP preserves PyTorch module context on ONNX nodes. The hierarchy source follows +the selected exporter: -Key Features: -- Works with complex models and control flow -- Built-in module tracking for better accuracy -- Conservative tag propagation -- Optimized for HuggingFace transformers -- New IO/ABC-based monitoring architecture +- TorchDynamo (default): reconstruct hierarchy from ONNX node metadata after export. +- TorchScript (``--no-dynamo``): trace module execution with PyTorch forward hooks, + then project the traced hierarchy onto ONNX nodes. -Variations: -- Standard HTP: Hook-based execution tracing -- Built-in HTP: Uses PyTorch's internal module tracking +Both paths emit the same ``winml.hierarchy.*`` metadata contract for downstream +inspection, benchmarking, and optimization. TODO: Future folder structure refactoring Currently keeping the folder structure flat for simplicity. diff --git a/src/winml/modelkit/export/htp/exporter.py b/src/winml/modelkit/export/htp/exporter.py index f95325262..e7764e6fb 100644 --- a/src/winml/modelkit/export/htp/exporter.py +++ b/src/winml/modelkit/export/htp/exporter.py @@ -76,8 +76,6 @@ class HTPConfig: "opset_version": 17, "do_constant_folding": True, "verbose": False, # ONNX internal verbose - # PyTorch dynamo export disabled by default. Use --dynamo flag to enable - # for rich node metadata (namespace, class_hierarchy, etc.) # dynamic_axes: Not set (defaults to None = static dimensions) # This prevents dynamic batch which causes MatMulAddFusion failure } @@ -93,7 +91,7 @@ class HTPConfig: DEFAULT_EXPORT_STATS: ClassVar[dict[str, Any]] = { # Seconds elapsed from export() entry to final stat collection. "export_time": 0.0, - # Number of named hierarchy modules discovered during tracing. + # Number of named modules in the recovered or traced hierarchy. "hierarchy_modules": 0, # Total ONNX graph nodes in the exported model. "onnx_nodes": 0, diff --git a/src/winml/modelkit/export/htp/htp_metadata_schema.json b/src/winml/modelkit/export/htp/htp_metadata_schema.json index dcf027e79..0d896b1e0 100644 --- a/src/winml/modelkit/export/htp/htp_metadata_schema.json +++ b/src/winml/modelkit/export/htp/htp_metadata_schema.json @@ -158,7 +158,7 @@ "default": "trace" }, "modules_traced": { - "description": "Number of modules traced during execution", + "description": "Number of modules present in the recovered or traced hierarchy", "type": "integer", "minimum": 0 }, @@ -378,6 +378,11 @@ "timestamp": { "$ref": "#/$defs/ISOTimestamp" }, + "source": { + "description": "How the hierarchy was obtained", + "type": "string", + "enum": ["trace", "onnx_metadata"] + }, "modules_traced": {"type": "integer", "minimum": 0}, "execution_steps": {"type": "integer", "minimum": 0} } From 7627d992dc5d78bbc5ec5c0103bb18b68e697443 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Wed, 22 Jul 2026 18:23:24 +0800 Subject: [PATCH 14/15] Clarify source-specific hierarchy documentation --- docs/concepts/load-and-export.md | 4 +++- src/winml/modelkit/export/htp/exporter.py | 7 ++++--- src/winml/modelkit/export/htp/htp_metadata_schema.json | 2 +- 3 files changed, 8 insertions(+), 5 deletions(-) diff --git a/docs/concepts/load-and-export.md b/docs/concepts/load-and-export.md index e5c96ca58..8fbc378d7 100644 --- a/docs/concepts/load-and-export.md +++ b/docs/concepts/load-and-export.md @@ -31,7 +31,9 @@ During export the HTP (Hierarchy-preserving Tags Protocol) exporter attaches two #### How tags are built -The module path attached to each node is obtained differently depending on the exporter. Under the default TorchDynamo exporter, dynamo records originating-module information as ONNX metadata, and the exporter reconstructs the hierarchy directly from that metadata after export — no forward hooks are involved. Under `--no-dynamo`, the TorchScript path instead registers PyTorch forward hooks on each module: when a module executes, a pre-hook pushes its class name onto a tag stack and the post-hook pops it. Either way the result is hierarchical paths that mirror the PyTorch module tree: +The module path attached to each node is obtained differently depending on the exporter. Under the default TorchDynamo exporter, dynamo records originating-module information as ONNX metadata, and the exporter reconstructs the hierarchy directly from that metadata after export — no forward hooks are involved. Under `--no-dynamo`, the TorchScript path instead registers PyTorch forward hooks on each module: when a module executes, a pre-hook pushes its class name onto a tag stack and the post-hook pops it. Either way the result is hierarchical paths that mirror the PyTorch module tree. + +The hook flow below applies only to the `--no-dynamo` TorchScript path: ```mermaid flowchart LR diff --git a/src/winml/modelkit/export/htp/exporter.py b/src/winml/modelkit/export/htp/exporter.py index e7764e6fb..d456e6b64 100644 --- a/src/winml/modelkit/export/htp/exporter.py +++ b/src/winml/modelkit/export/htp/exporter.py @@ -7,11 +7,12 @@ """HTP (Hierarchy-preserving Tags Protocol) Exporter. This exporter preserves the hierarchical structure of HuggingFace models -when converting to ONNX format by tracing module execution and tagging -ONNX nodes with their source module information. +when converting to ONNX format. TorchDynamo recovers module context from +ONNX node metadata; the legacy TorchScript path traces module execution. +Both paths tag ONNX nodes with their source module information. Key Features: -- Direct module context capture during execution +- Source-aware hierarchy recovery for both ONNX exporters - Precise hierarchy tag generation - Comprehensive metadata export - Optional detailed reporting diff --git a/src/winml/modelkit/export/htp/htp_metadata_schema.json b/src/winml/modelkit/export/htp/htp_metadata_schema.json index 0d896b1e0..3591020f7 100644 --- a/src/winml/modelkit/export/htp/htp_metadata_schema.json +++ b/src/winml/modelkit/export/htp/htp_metadata_schema.json @@ -484,7 +484,7 @@ "minimum": 0 }, "traced_modules": { - "description": "Number of modules traced during execution", + "description": "Number of modules present in the recovered or traced hierarchy", "type": "integer", "minimum": 0 }, From 144cba4eae295159e2a9e1fe410f3f7a56afa4b4 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Wed, 22 Jul 2026 18:41:49 +0800 Subject: [PATCH 15/15] Clarify hierarchy encounter order Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../modelkit/export/htp/htp_metadata_schema.json | 2 +- .../modelkit/export/htp/markdown_report_writer.py | 11 ++++++++--- tests/unit/export/test_dynamo_hierarchy_report.py | 12 ++++++++++++ 3 files changed, 21 insertions(+), 4 deletions(-) diff --git a/src/winml/modelkit/export/htp/htp_metadata_schema.json b/src/winml/modelkit/export/htp/htp_metadata_schema.json index 3591020f7..5873afb79 100644 --- a/src/winml/modelkit/export/htp/htp_metadata_schema.json +++ b/src/winml/modelkit/export/htp/htp_metadata_schema.json @@ -113,7 +113,7 @@ "type": "string" }, "execution_order": { - "description": "Order of execution during tracing", + "description": "Order the module was first encountered (execution order for trace, discovery order in the exported graph for onnx_metadata)", "title": "Execution Order", "type": "integer", "minimum": 0 diff --git a/src/winml/modelkit/export/htp/markdown_report_writer.py b/src/winml/modelkit/export/htp/markdown_report_writer.py index a156e356f..d714027e4 100644 --- a/src/winml/modelkit/export/htp/markdown_report_writer.py +++ b/src/winml/modelkit/export/htp/markdown_report_writer.py @@ -427,7 +427,12 @@ def _write_module_hierarchy_section(self, data: ExportData) -> None: hierarchy = data.hierarchy.hierarchy # Module table with reordered columns - self.doc.add_heading("Module List (Sorted by Execution Order)", level=3) + order_label = ( + "Discovery Order" + if data.hierarchy.source == HIERARCHY_SOURCE_ONNX_METADATA + else "Execution Order" + ) + self.doc.add_heading(f"Module List (Sorted by {order_label})", level=3) # Count direct and total nodes for each module if available direct_counts: dict[str, int] = {} @@ -437,10 +442,10 @@ def _write_module_hierarchy_section(self, data: ExportData) -> None: data.node_tagging.tagged_nodes ) - headers = ["Execution Order", "Class Name", "Nodes", "Tag", "Scope"] + headers = [order_label, "Class Name", "Nodes", "Tag", "Scope"] rows = [] - # Sort by execution order + # Sort by the source-specific encounter order sorted_items = sorted( hierarchy.items(), key=lambda x: ( diff --git a/tests/unit/export/test_dynamo_hierarchy_report.py b/tests/unit/export/test_dynamo_hierarchy_report.py index 882405b90..919d0f520 100644 --- a/tests/unit/export/test_dynamo_hierarchy_report.py +++ b/tests/unit/export/test_dynamo_hierarchy_report.py @@ -169,6 +169,18 @@ def test_trace_source_uses_trace_wording(self) -> None: assert "Total execution steps" in out assert "42" in out + def test_dynamo_module_list_uses_discovery_order(self) -> None: + writer = MarkdownReportWriter("unused.md") + writer._write_module_hierarchy_section( + ExportData( + hierarchy=_hierarchy_data(HIERARCHY_SOURCE_ONNX_METADATA, None), + ) + ) + + out = str(writer.doc) + assert "Discovery Order" in out + assert "Execution Order" not in out + def test_metadata_records_dynamo_source_without_fake_step_count(self) -> None: writer = MetadataWriter("unused.json") writer.write(