Fix large-model FP16 conversion and add XLM-RoBERTa NER recipes

examples/recipes/Davlan_xlm-roberta-large-ner-hrl/cpu/cpu/token-classification_fp16_config.json

@@ -0,0 +1,76 @@
+{
+  "export": {
+    "opset_version": 17,
+    "batch_size": 1,
+    "export_params": true,
+    "do_constant_folding": true,
+    "verbose": false,
+    "dynamo": false,
+    "enable_hierarchy_tags": true,
+    "clean_onnx": false,
+    "hierarchy_tag_format": "full",
+    "input_tensors": [
+      {
+        "name": "input_ids",
+        "dtype": "int32",
+        "shape": [
+          1,
+          512
+        ],
+        "value_range": [
+          0,
+          250002
+        ]
+      },
+      {
+        "name": "attention_mask",
+        "dtype": "int32",
+        "shape": [
+          1,
+          512
+        ],
+        "value_range": [
+          0,
+          2
+        ]
+      }
+    ],
+    "output_tensors": [
+      {
+        "name": "logits"
+      }
+    ]
+  },
+  "optim": {
+    "clamp_constant_values": true
+  },
+  "quant": {
+    "mode": "fp16",
+    "samples": 10,
+    "calibration_method": "minmax",
+    "weight_type": "uint8",
+    "activation_type": "uint8",
+    "per_channel": false,
+    "symmetric": false,
+    "weight_symmetric": null,
+    "activation_symmetric": null,
+    "save_calibration": false,
+    "distribution": "uniform",
+    "seed": null,
+    "calibration_load_path": null,
+    "calibration_save_path": null,
+    "op_types_to_quantize": null,
+    "nodes_to_exclude": null,
+    "task": "token-classification",
+    "model_id": "Davlan/xlm-roberta-large-ner-hrl",
+    "model_type": "xlm-roberta",
+    "fp16_keep_io_types": true,
+    "fp16_op_block_list": null
+  },
+  "compile": null,
+  "loader": {
+    "task": "token-classification",
+    "model_class": "XLMRobertaForTokenClassification",
+    "model_type": "xlm-roberta"
+  }
+}

examples/recipes/Davlan_xlm-roberta-large-ner-hrl/cpu/cpu/token-classification_fp32_config.json

@@ -0,0 +1,54 @@
+{
+  "export": {
+    "opset_version": 17,
+    "batch_size": 1,
+    "export_params": true,
+    "do_constant_folding": true,
+    "verbose": false,
+    "dynamo": false,
+    "enable_hierarchy_tags": true,
+    "clean_onnx": false,
+    "hierarchy_tag_format": "full",
+    "input_tensors": [
+      {
+        "name": "input_ids",
+        "dtype": "int32",
+        "shape": [
+          1,
+          512
+        ],
+        "value_range": [
+          0,
+          250002
+        ]
+      },
+      {
+        "name": "attention_mask",
+        "dtype": "int32",
+        "shape": [
+          1,
+          512
+        ],
+        "value_range": [
+          0,
+          2
+        ]
+      }
+    ],
+    "output_tensors": [
+      {
+        "name": "logits"
+      }
+    ]
+  },
+  "optim": {
+    "clamp_constant_values": true
+  },
+  "quant": null,
+  "compile": null,
+  "loader": {
+    "task": "token-classification",
+    "model_class": "XLMRobertaForTokenClassification",
+    "model_type": "xlm-roberta"
+  }
+}

src/winml/modelkit/commands/analyze.py

@@ -979,10 +979,11 @@ def analyze(
                 )
             execution_pairs = _sort_ep_device_pairs(execution_pairs)
 
-        # Local pairs are still needed to gate --run-unknown-op probing
-        # (_resolve_run_unknown_op). Single-target `auto` selection is already
-        # local by construction, so no extra intersection/warning is required.
-        local_pairs = set(_get_local_ep_device_pairs())
+        # Local pairs are only needed to gate runtime unknown-op probing.
+        # Static rule analysis (the default) must not load optional provider
+        # plugins merely to print an informational list: on Windows that can
+        # leave native plugin/catalog state alive until interpreter teardown.
+        local_pairs = set(_get_local_ep_device_pairs()) if run_unknown_op else set()
 
         if not execution_pairs:
             raise click.UsageError("No EP/device combination matched the current selection.")

src/winml/modelkit/quant/fp16.py

@@ -11,21 +11,23 @@
 from __future__ import annotations
 
 import logging
-from typing import TYPE_CHECKING
+import tempfile
+from pathlib import Path
+from typing import TYPE_CHECKING, cast
 
 
 if TYPE_CHECKING:
-    from onnx import ModelProto
+    import onnx
 
 logger = logging.getLogger(__name__)
 
 
 def convert_to_fp16(
-    model: ModelProto,
+    model: onnx.ModelProto | str | Path,
     *,
     keep_io_types: bool = True,
     op_block_list: list[str] | None = None,
-) -> ModelProto:
+) -> onnx.ModelProto:
     """Convert an ONNX model from FP32 to FP16 precision.
 
     Uses onnxruntime.transformers.float16.convert_float_to_float16 internally.
@@ -34,7 +36,10 @@
     Note: ORT's converter mutates the model in-place and returns the same object.
 
     Args:
-        model: Input ONNX ModelProto (will be mutated in-place by ORT).
+        model: Input ONNX ModelProto or path. Paths use ORT's file-based shape
+            inference, which supports models above protobuf's 2 GB in-memory
+            serialization limit and resolves external data relative to the
+            model file.
         keep_io_types: If True, preserve FP32 model inputs/outputs by inserting
             Cast nodes at boundaries. Recommended for CPU-safe inference.
         op_block_list: Op types to keep in FP32 (e.g., ["LayerNorm", "Softmax"]).
@@ -44,31 +49,71 @@
     Returns:
         The converted model (same object as input due to ORT in-place mutation).
     """
-    from onnx import TensorProto
+    import onnx
     from onnxruntime.transformers.float16 import convert_float_to_float16
 
     # Skip if model is already FP16 (check floating-point initializer dtypes)
-    fp32_types = {TensorProto.FLOAT, TensorProto.DOUBLE, TensorProto.BFLOAT16}
-    initializers = model.graph.initializer
+    fp32_types = {onnx.TensorProto.FLOAT, onnx.TensorProto.DOUBLE, onnx.TensorProto.BFLOAT16}
+    model_path = Path(model) if isinstance(model, str | Path) else None
+    if model_path is not None:
+        inspection_model = onnx.load(str(model_path), load_external_data=False)
+    else:
+        inspection_model = cast("onnx.ModelProto", model)
+
+    initializers = inspection_model.graph.initializer
     if initializers:
-        float_inits = [t for t in initializers if t.data_type in fp32_types | {TensorProto.FLOAT16}]
-        if float_inits and all(t.data_type == TensorProto.FLOAT16 for t in float_inits):
+        float_inits = [
+            t for t in initializers if t.data_type in fp32_types | {onnx.TensorProto.FLOAT16}
+        ]
+        if float_inits and all(t.data_type == onnx.TensorProto.FLOAT16 for t in float_inits):
             logger.info("Model is already FP16 — skipping conversion.")
-            return model
+            if model_path is not None:
+                # A graph-only load retains external-data locations relative to
+                # the source model. Materialize those tensors before returning
+                # so callers can safely persist the result in another directory.
+                onnx.load_external_data_for_model(inspection_model, str(model_path.parent))
+            return inspection_model
 
-    original_nodes = len(model.graph.node)
+    original_nodes = len(inspection_model.graph.node)
 
     logger.info("Converting model to FP16...")
     if keep_io_types:
         logger.info("  Keeping I/O types as FP32")
     if op_block_list:
         logger.info("  Keeping ops in FP32: %s", op_block_list)
 
-    converted: ModelProto = convert_float_to_float16(
-        model,
-        keep_io_types=keep_io_types,
-        op_block_list=op_block_list,
-    )
+    if model_path is not None:
+        # ORT's converter uses NamedTemporaryFile while it is still open,
+        # which cannot be reopened by ONNX on Windows. Own the temporary path
+        # here, close it before inference, and retain the file-based API that
+        # supports protobufs above 2 GB.
+        with tempfile.NamedTemporaryFile(
+            dir=model_path.parent, suffix=".shape_inferred.onnx", delete=False
+        ) as temporary:
+            inferred_path = Path(temporary.name)
+        try:
+            onnx.shape_inference.infer_shapes_path(str(model_path), str(inferred_path))
+            inferred_model = onnx.load(str(inferred_path))
+        finally:
+            inferred_path.unlink(missing_ok=True)
+        converted = cast(
+            "onnx.ModelProto",
+            convert_float_to_float16(
+                inferred_model,
+                keep_io_types=keep_io_types,
+                disable_shape_infer=True,
+                op_block_list=op_block_list,
+            ),
+        )
+    else:
+        converted = cast(
+            "onnx.ModelProto",
+            convert_float_to_float16(
+                model,
+                keep_io_types=keep_io_types,
+                op_block_list=op_block_list,
+            ),
+        )
 
     # ORT's converter appends Cast nodes at the end of the node list (for
     # keep_io_types), which breaks topological ordering. Re-sort the graph

src/winml/modelkit/quant/passes/fp16.py

@@ -44,7 +44,7 @@ def run(
         use_external_data: bool = True,
     ) -> QuantizeResult:
         """Convert *model_path* to FP16 and write the result to *output_path*."""
-        from ...onnx import load_onnx, save_onnx
+        from ...onnx import save_onnx
         from ..config import QuantizeResult
         from ..fp16 import convert_to_fp16
 
@@ -59,9 +59,8 @@ def run(
         warnings: list[str] = []
 
         logger.info("Running FP16-only conversion (no quantization)...")
-        model = load_onnx(model_path, validate=False)
         model = convert_to_fp16(
-            model,
+            model_path,
             keep_io_types=self._config.fp16_keep_io_types,
             op_block_list=self._config.fp16_op_block_list,
         )

src/winml/modelkit/session/session.py

@@ -165,7 +165,13 @@ def __init__(
                 (``add_provider_for_devices`` mutates the options and cannot
                 be replayed). Defaults to ``ort.SessionOptions``.
         """
-        WinMLSession._init_winml_eps_once()
+        # The built-in CPU EP needs no optional WindowsML plugin discovery.
+        # Avoid loading the WindowsML catalog for an explicitly requested CPU
+        # session: besides unnecessary startup work, mixing catalog/plugin
+        # lifetime with a CPU-only ORT session can crash native teardown on
+        # Windows after otherwise successful inference.
+        if ep not in ("cpu", "CPUExecutionProvider"):
+            WinMLSession._init_winml_eps_once()
 
         self._onnx_path = Path(onnx_path)
         if not self._onnx_path.exists():
@@ -355,9 +361,7 @@ def run(
             # Run inference (with optional perf tracking)
             output_names = [o.name for o in session.get_outputs()]
             if self._perf_stats:
-                outputs = self._perf_stats.record(
-                    lambda: session.run(output_names, ort_inputs)
-                )
+                outputs = self._perf_stats.record(lambda: session.run(output_names, ort_inputs))
             else:
                 outputs = session.run(output_names, ort_inputs)
 

src/winml/modelkit/sysinfo/device.py

@@ -213,10 +213,21 @@ def resolve_device(
     if device != "auto" and device not in _VALID_DEVICES:
         raise ValueError(f"Unknown device '{device}'. Expected 'auto', 'npu', 'gpu', or 'cpu'.")
 
+    # CPUExecutionProvider is built into ORT and does not require WindowsML's
+    # optional provider catalog. Keep explicit --ep cpu resolution on this
+    # direct path so CPU-only commands do not load unrelated plugin handles
+    # whose native teardown may outlive an ORT session on Windows.
+    ep_full = normalize_ep_name(ep)
+    if device == "cpu" and ep_full == "CPUExecutionProvider":
+        import onnxruntime as ort
+
+        if "CPUExecutionProvider" not in ort.get_available_providers():
+            raise ValueError("Requested EP 'cpu' is not available on this system.")
+        return "cpu", ["cpu"]
+
     device_ep_map = dict(_get_device_ep_map_from_ort())
 
     if ep is not None:
-        ep_full = normalize_ep_name(ep)
         if ep_full not in EP_SUPPORTED_DEVICES:
             raise ValueError(f"Unknown EP '{ep}'. Expected one of: {sorted(EP_SUPPORTED_DEVICES)}")
         # Static policy gate (see issue #860): reject EP/device combos the EP