feat(layoutlm): add QNN NPU document QA recipes

examples/recipes/impira_layoutlm-document-qa/qnn/npu/question-answering_fp16_config.json

@@ -0,0 +1,124 @@
+{
+  "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,
+          50265
+        ]
+      },
+      {
+        "name": "bbox",
+        "dtype": "int32",
+        "shape": [
+          1,
+          512,
+          4
+        ],
+        "value_range": [
+          0,
+          1000
+        ]
+      },
+      {
+        "name": "attention_mask",
+        "dtype": "int32",
+        "shape": [
+          1,
+          512
+        ],
+        "value_range": [
+          0,
+          2
+        ]
+      },
+      {
+        "name": "token_type_ids",
+        "dtype": "int32",
+        "shape": [
+          1,
+          512
+        ],
+        "value_range": [
+          0,
+          1
+        ]
+      }
+    ],
+    "output_tensors": [
+      {
+        "name": "start_logits"
+      },
+      {
+        "name": "end_logits"
+      }
+    ],
+    "compatibility": {
+      "transformers_attention": "eager"
+    }
+  },
+  "optim": {
+    "clamp_constant_values": true,
+    "gelu_fusion": false,
+    "matmul_add_fusion": false
+  },
+  "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": "question-answering",
+    "model_id": "impira/layoutlm-document-qa",
+    "model_type": "layoutlm",
+    "fp16_keep_io_types": true,
+    "fp16_op_block_list": null
+  },
+  "compile": {
+    "execution_provider": "qnn",
+    "provider_options": {
+      "device_type": "NPU",
+      "htp_performance_mode": "burst",
+      "htp_graph_finalization_optimization_mode": "3"
+    },
+    "provider_option_file_keys": [],
+    "enable_ep_context": true,
+    "embed_context": false,
+    "compiler": "ort",
+    "qnn_sdk_root": null,
+    "device": "npu",
+    "ep_device": null,
+    "validate": false
+  },
+  "loader": {
+    "task": "question-answering",
+    "model_class": "LayoutLMForQuestionAnswering",
+    "model_type": "layoutlm"
+  }
+}

examples/recipes/impira_layoutlm-document-qa/qnn/npu/question-answering_w8a16_config.json

@@ -0,0 +1,122 @@
+{
+  "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,
+          50265
+        ]
+      },
+      {
+        "name": "bbox",
+        "dtype": "int32",
+        "shape": [
+          1,
+          512,
+          4
+        ],
+        "value_range": [
+          0,
+          1000
+        ]
+      },
+      {
+        "name": "attention_mask",
+        "dtype": "int32",
+        "shape": [
+          1,
+          512
+        ],
+        "value_range": [
+          0,
+          2
+        ]
+      },
+      {
+        "name": "token_type_ids",
+        "dtype": "int32",
+        "shape": [
+          1,
+          512
+        ],
+        "value_range": [
+          0,
+          1
+        ]
+      }
+    ],
+    "output_tensors": [
+      {
+        "name": "start_logits"
+      },
+      {
+        "name": "end_logits"
+      }
+    ],
+    "compatibility": {
+      "transformers_attention": "eager"
+    }
+  },
+  "optim": {
+    "clamp_constant_values": true,
+    "gelu_fusion": true,
+    "matmul_add_fusion": false
+  },
+  "quant": {
+    "mode": "static",
+    "samples": 10,
+    "calibration_method": "minmax",
+    "weight_type": "uint8",
+    "activation_type": "uint16",
+    "per_channel": false,
+    "symmetric": false,
+    "weight_symmetric": null,
+    "activation_symmetric": null,
+    "save_calibration": false,
+    "distribution": "uniform",
+    "seed": 42,
+    "calibration_load_path": null,
+    "calibration_save_path": null,
+    "op_types_to_quantize": null,
+    "nodes_to_exclude": null,
+    "task": "question-answering",
+    "model_id": "impira/layoutlm-document-qa",
+    "model_type": "layoutlm"
+  },
+  "compile": {
+    "execution_provider": "qnn",
+    "provider_options": {
+      "device_type": "NPU",
+      "htp_performance_mode": "burst",
+      "htp_graph_finalization_optimization_mode": "3"
+    },
+    "provider_option_file_keys": [],
+    "enable_ep_context": true,
+    "embed_context": false,
+    "compiler": "ort",
+    "qnn_sdk_root": null,
+    "device": "npu",
+    "ep_device": null,
+    "validate": false
+  },
+  "loader": {
+    "task": "question-answering",
+    "model_class": "LayoutLMForQuestionAnswering",
+    "model_type": "layoutlm"
+  }
+}

src/winml/modelkit/build/common.py

@@ -22,7 +22,7 @@
 from ..analyze import analyze_onnx
 from ..compiler import compile_onnx
 from ..onnx import copy_onnx_model, is_quantized_onnx
-from ..optim import optimize_onnx
+from ..optim import WinMLOptimizationConfig, optimize_onnx
 from ..quant import quantize_onnx
 
 
@@ -385,6 +385,7 @@ def _run_analyze_loop(
     analyze_iterations = 0
     analyze_black_nodes = 0
     discovered_optim: dict[str, bool] = {}
+    explicitly_disabled_optim = {name for name, enabled in config.optim.items() if enabled is False}
     analysis = None
     _not_converged = False
 
@@ -426,6 +427,16 @@ def _run_analyze_loop(
             if on_patterns_discovered is not None:
                 on_patterns_discovered(optim_config)
 
+            optim_config = WinMLOptimizationConfig.from_dict(
+                {
+                    name: enabled
+                    for name, enabled in optim_config.items()
+                    if name not in explicitly_disabled_optim
+                }
+            )
+            if not optim_config:
+                break
+
             # Notify: re-optimizing with discovered flags
             if on_reoptimize is not None:
                 on_reoptimize(optim_config)

src/winml/modelkit/commands/build.py

@@ -1155,7 +1155,13 @@ def _patch_device(cfg: WinMLBuildConfig) -> None:
                     provider = cfg.compile.ep_config.provider
                     patched = WinMLCompileConfig.for_provider(provider, device=runtime_device)
                     if patched is not None:
-                        cfg.compile = patched
+                        cfg.compile.ep_config.provider = patched.ep_config.provider
+                        cfg.compile.ep_config.device = patched.ep_config.device
+                        cfg.compile.ep_config.provider_options = {
+                            **cfg.compile.ep_config.provider_options,
+                            **patched.ep_config.provider_options,
+                        }
+                        cfg.compile.ep_device = patched.ep_device
 
             if isinstance(config_or_configs, list):
                 for _cfg in config_or_configs:

src/winml/modelkit/models/hf/layoutlm.py

@@ -13,14 +13,15 @@
 from optimum.utils.input_generators import DummyBboxInputGenerator, DummyVisionInputGenerator
 
 from ...export import MaxLengthTextInputGenerator, register_onnx_overwrite
+from .roberta import _adjust_position_embeddings
 
 
 if TYPE_CHECKING:
     import torch
 
 
-class ZeroTokenTypeLayoutLMTextInputGenerator(MaxLengthTextInputGenerator):
-    """LayoutLM text dummy generator that keeps token_type_ids within type_vocab_size=1."""
+class LayoutLMTextInputGenerator(MaxLengthTextInputGenerator):
+    """LayoutLM text dummy generator constrained by the configured token-type vocabulary."""
 
     def generate(
         self,
@@ -29,8 +30,18 @@ def generate(
         int_dtype: str = "int64",
         float_dtype: str = "fp32",
     ) -> torch.Tensor:
-        """Generate LayoutLM text inputs, replacing token_type_ids with zeros."""
-        tensor = cast(
+        """Generate LayoutLM text inputs within the configured token-type range."""
+        if input_name == "token_type_ids":
+            return cast(
+                "torch.Tensor",
+                self.random_int_tensor(
+                    [self.batch_size, self.sequence_length],
+                    max_value=self.normalized_config.type_vocab_size,
+                    framework=framework,
+                    dtype=int_dtype,
+                ),
+            )
+        return cast(
             "torch.Tensor",
             super().generate(
                 input_name,
@@ -39,9 +50,6 @@ def generate(
                 float_dtype=float_dtype,
             ),
         )
-        if input_name == "token_type_ids":
-            return tensor.new_zeros(tensor.shape)
-        return tensor
 
 
 @register_onnx_overwrite("layoutlm", "question-answering", library_name="transformers")
@@ -59,10 +67,15 @@ class LayoutLMQAIOConfig(LayoutLMOnnxConfig):  # type: ignore[misc]  # optimum b
     # recipe's `value_range` instead.
     NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args(
         sequence_length="max_position_embeddings",
+        type_vocab_size="type_vocab_size",
         allow_new=True,
     )
     DUMMY_INPUT_GENERATOR_CLASSES: tuple[type[Any], ...] = (
-        ZeroTokenTypeLayoutLMTextInputGenerator,
+        LayoutLMTextInputGenerator,
         DummyVisionInputGenerator,
         DummyBboxInputGenerator,
     )
+
+    def __init__(self, config: Any, task: str, **kwargs: Any) -> None:
+        _adjust_position_embeddings(config)
+        super().__init__(config, task, **kwargs)