Add dynamo export support and make it the default exporter

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, 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`). |
@@ -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
@@ -79,6 +82,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
@@ -114,9 +126,26 @@ 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.
+  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 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
+  `--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
   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

docs/concepts/load-and-export.md

@@ -2,23 +2,23 @@
 
 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.
 
 Some community models host custom Python code in their repositories. The loader refuses to execute it by default. Pass `--trust-remote-code` to `winml config` when generating a build configuration for such a model.
 
 ## 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 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,9 @@ 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 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
@@ -42,11 +44,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)
@@ -75,7 +77,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 +102,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
 

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. |

src/winml/modelkit/commands/export.py

@@ -116,9 +116,11 @@ 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, the validated "
+    "path for QNN/NPU compilation today.",
 )
 @click.option(
     "--torch-module",
@@ -185,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
@@ -206,8 +210,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 +327,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``."""

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),

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]]: