From ec130eaaf71bce478ca0792afc0c19e9c1ffd614 Mon Sep 17 00:00:00 2001 From: Shiyi Zheng Date: Mon, 24 Aug 2026 21:40:06 +0800 Subject: [PATCH] Add mxbai reranking FP16 boundary support --- .../cpu/cpu/reranking_fp16_config.json | 153 ++++++++++++++ ...config.json => reranking_fp32_config.json} | 15 +- .../cpu/text-classification_fp32_config.json | 54 ----- src/winml/modelkit/quant/config.py | 16 ++ src/winml/modelkit/quant/fp16.py | 13 +- src/winml/modelkit/quant/passes/fp16.py | 2 + tests/unit/config/test_build.py | 50 +++++ tests/unit/optim/test_fp16.py | 190 +++++++++++++++++- tests/unit/recipes/test_cpu_recipes.py | 41 ++++ tests/unit/test_quant_passes.py | 55 ++++- 10 files changed, 515 insertions(+), 74 deletions(-) create mode 100644 examples/recipes/mixedbread-ai_mxbai-rerank-base-v1/cpu/cpu/reranking_fp16_config.json rename examples/recipes/mixedbread-ai_mxbai-rerank-base-v1/cpu/cpu/{text-classification_fp16_config.json => reranking_fp32_config.json} (82%) delete mode 100644 examples/recipes/mixedbread-ai_mxbai-rerank-base-v1/cpu/cpu/text-classification_fp32_config.json diff --git a/examples/recipes/mixedbread-ai_mxbai-rerank-base-v1/cpu/cpu/reranking_fp16_config.json b/examples/recipes/mixedbread-ai_mxbai-rerank-base-v1/cpu/cpu/reranking_fp16_config.json new file mode 100644 index 000000000..45b9ee79f --- /dev/null +++ b/examples/recipes/mixedbread-ai_mxbai-rerank-base-v1/cpu/cpu/reranking_fp16_config.json @@ -0,0 +1,153 @@ +{ + "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, + 128100 + ] + }, + { + "name": "attention_mask", + "dtype": "int32", + "shape": [ + 1, + 512 + ], + "value_range": [ + 0, + 2 + ] + } + ], + "output_tensors": [ + { + "name": "logits" + } + ], + "compatibility": { + "transformers_attention": "eager" + } + }, + "optim": {}, + "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": "reranking", + "model_id": "mixedbread-ai/mxbai-rerank-base-v1", + "model_type": "deberta-v2", + "fp16_keep_io_types": true, + "fp16_op_block_list": null, + "fp16_nodes_to_exclude": [ + "InsertedPrecisionFreeCast_/deberta/embeddings/LayerNorm/LayerNormalization_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.0/attention/self/Transpose_3_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.0/attention/self/Reshape_1_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.0/attention/self/Reshape_3_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.0/attention/self/Transpose_8_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.0/attention/self/Reshape_5_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.0/attention/self/Reshape_13_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.1/attention/self/Transpose_3_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.1/attention/self/Reshape_1_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.1/attention/self/Reshape_3_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.1/attention/self/Transpose_8_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.1/attention/self/Reshape_5_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.1/attention/self/Reshape_13_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.2/attention/self/Transpose_3_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.2/attention/self/Reshape_1_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.2/attention/self/Reshape_3_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.2/attention/self/Transpose_8_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.2/attention/self/Reshape_5_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.2/attention/self/Reshape_13_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.3/attention/self/Transpose_3_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.3/attention/self/Reshape_1_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.3/attention/self/Reshape_3_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.3/attention/self/Transpose_8_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.3/attention/self/Reshape_5_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.3/attention/self/Reshape_13_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.4/attention/self/Transpose_3_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.4/attention/self/Reshape_1_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.4/attention/self/Reshape_3_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.4/attention/self/Transpose_8_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.4/attention/self/Reshape_5_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.4/attention/self/Reshape_13_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.5/attention/self/Transpose_3_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.5/attention/self/Reshape_1_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.5/attention/self/Reshape_3_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.5/attention/self/Transpose_8_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.5/attention/self/Reshape_5_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.5/attention/self/Reshape_13_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.6/attention/self/Transpose_3_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.6/attention/self/Reshape_1_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.6/attention/self/Reshape_3_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.6/attention/self/Transpose_8_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.6/attention/self/Reshape_5_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.6/attention/self/Reshape_13_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.7/attention/self/Transpose_3_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.7/attention/self/Reshape_1_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.7/attention/self/Reshape_3_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.7/attention/self/Transpose_8_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.7/attention/self/Reshape_5_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.7/attention/self/Reshape_13_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.8/attention/self/Transpose_3_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.8/attention/self/Reshape_1_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.8/attention/self/Reshape_3_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.8/attention/self/Transpose_8_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.8/attention/self/Reshape_5_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.8/attention/self/Reshape_13_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.9/attention/self/Transpose_3_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.9/attention/self/Reshape_1_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.9/attention/self/Reshape_3_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.9/attention/self/Transpose_8_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.9/attention/self/Reshape_5_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.9/attention/self/Reshape_13_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.10/attention/self/Transpose_3_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.10/attention/self/Reshape_1_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.10/attention/self/Reshape_3_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.10/attention/self/Transpose_8_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.10/attention/self/Reshape_5_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.10/attention/self/Reshape_13_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.11/attention/self/Transpose_3_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.11/attention/self/Reshape_1_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.11/attention/self/Reshape_3_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.11/attention/self/Transpose_8_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.11/attention/self/Reshape_5_output_0", + "InsertedPrecisionFreeCast_/deberta/encoder/layer.11/attention/self/Reshape_13_output_0", + "InsertedPrecisionFreeCast_/pooler/Gather_output_0" + ] + }, + "compile": null, + "loader": { + "task": "reranking", + "model_class": "DebertaV2ForSequenceClassification", + "model_type": "deberta-v2" + } +} \ No newline at end of file diff --git a/examples/recipes/mixedbread-ai_mxbai-rerank-base-v1/cpu/cpu/text-classification_fp16_config.json b/examples/recipes/mixedbread-ai_mxbai-rerank-base-v1/cpu/cpu/reranking_fp32_config.json similarity index 82% rename from examples/recipes/mixedbread-ai_mxbai-rerank-base-v1/cpu/cpu/text-classification_fp16_config.json rename to examples/recipes/mixedbread-ai_mxbai-rerank-base-v1/cpu/cpu/reranking_fp32_config.json index 7b36e6952..f32bd8c6c 100644 --- a/examples/recipes/mixedbread-ai_mxbai-rerank-base-v1/cpu/cpu/text-classification_fp16_config.json +++ b/examples/recipes/mixedbread-ai_mxbai-rerank-base-v1/cpu/cpu/reranking_fp32_config.json @@ -39,16 +39,17 @@ { "name": "logits" } - ] - }, - "optim": { - "clamp_constant_values": true + ], + "compatibility": { + "transformers_attention": "eager" + } }, + "optim": {}, "quant": null, "compile": null, "loader": { - "task": "text-classification", - "model_class": "AutoModelForSequenceClassification", + "task": "reranking", + "model_class": "DebertaV2ForSequenceClassification", "model_type": "deberta-v2" } -} +} \ No newline at end of file diff --git a/examples/recipes/mixedbread-ai_mxbai-rerank-base-v1/cpu/cpu/text-classification_fp32_config.json b/examples/recipes/mixedbread-ai_mxbai-rerank-base-v1/cpu/cpu/text-classification_fp32_config.json deleted file mode 100644 index 7b36e6952..000000000 --- a/examples/recipes/mixedbread-ai_mxbai-rerank-base-v1/cpu/cpu/text-classification_fp32_config.json +++ /dev/null @@ -1,54 +0,0 @@ -{ - "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, - 128100 - ] - }, - { - "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": "text-classification", - "model_class": "AutoModelForSequenceClassification", - "model_type": "deberta-v2" - } -} diff --git a/src/winml/modelkit/quant/config.py b/src/winml/modelkit/quant/config.py index d1fa15cfc..ca17d9d08 100644 --- a/src/winml/modelkit/quant/config.py +++ b/src/winml/modelkit/quant/config.py @@ -119,6 +119,20 @@ class WinMLQuantizationConfig: # FP16 conversion settings (only used when mode="fp16") fp16_keep_io_types: bool = True fp16_op_block_list: list[str] | None = None + fp16_nodes_to_exclude: list[str] | None = None + + def __post_init__(self) -> None: + """Normalize FP16 node exclusions while preserving declaration order.""" + node_exclusions: object = self.fp16_nodes_to_exclude + if node_exclusions is None: + return + if not isinstance(node_exclusions, list): + msg = "fp16_nodes_to_exclude must be a list of non-empty strings" + raise TypeError(msg) + if not all(isinstance(name, str) and name for name in node_exclusions): + msg = "fp16_nodes_to_exclude entries must be non-empty strings" + raise ValueError(msg) + self.fp16_nodes_to_exclude = list(dict.fromkeys(node_exclusions)) def to_dict(self) -> dict: """Convert to dictionary for serialization. @@ -168,6 +182,7 @@ def to_dict(self) -> dict: if self.mode == "fp16": result["fp16_keep_io_types"] = self.fp16_keep_io_types result["fp16_op_block_list"] = self.fp16_op_block_list + result["fp16_nodes_to_exclude"] = self.fp16_nodes_to_exclude return result @classmethod @@ -217,6 +232,7 @@ def from_dict(cls, data: dict) -> WinMLQuantizationConfig: reduce_range=data.get("reduce_range", False), fp16_keep_io_types=data.get("fp16_keep_io_types", True), fp16_op_block_list=data.get("fp16_op_block_list"), + fp16_nodes_to_exclude=data.get("fp16_nodes_to_exclude"), ) diff --git a/src/winml/modelkit/quant/fp16.py b/src/winml/modelkit/quant/fp16.py index ff207b414..6bc6a5b79 100644 --- a/src/winml/modelkit/quant/fp16.py +++ b/src/winml/modelkit/quant/fp16.py @@ -2822,7 +2822,7 @@ def _validate_local_function_conversion(model: ModelProto) -> None: raise RuntimeError(msg) from error -def _validate_converted_types(model: ModelProto) -> None: +def _validate_converted_types(model: ModelProto, *, required: bool = True) -> bool: """Reject converted graphs whose concrete types no longer agree.""" from onnx import checker, shape_inference @@ -2831,11 +2831,16 @@ def _validate_converted_types(model: ModelProto) -> None: inferred = shape_inference.infer_shapes(model, check_type=True, strict_mode=True) checker.check_model(inferred) except ( + AttributeError, + EncodeError, checker.ValidationError, shape_inference.InferenceError, ) as error: + if not required: + return False msg = "FP16 conversion produced incompatible FP16 types." raise RuntimeError(msg) from error + return True def _node_dependencies(node: NodeProto) -> set[str]: @@ -2889,6 +2894,7 @@ def convert_to_fp16( *, keep_io_types: bool = True, op_block_list: list[str] | None = None, + node_block_list: list[str] | None = None, ) -> ModelProto: """Convert an ONNX model from FP32 to FP16 precision. @@ -3024,11 +3030,13 @@ def convert_to_fp16( if op_block_list: logger.info(" Keeping ops in FP32: %s", op_block_list) + supports_strict_type_validation = _validate_converted_types(conversion_model, required=False) try: converted: ModelProto = convert_float_to_float16( conversion_model, keep_io_types=keep_io_types, op_block_list=op_block_list, + node_block_list=node_block_list, ) except EncodeError: logger.warning( @@ -3041,6 +3049,7 @@ def convert_to_fp16( keep_io_types=keep_io_types, disable_shape_infer=True, op_block_list=op_block_list, + node_block_list=node_block_list, ) converted_graphs = _ort_traversed_graphs(converted, op_block_list) @@ -3060,7 +3069,7 @@ def convert_to_fp16( _repair_sparse_float_initializers(converted, op_block_list, keep_io_types=keep_io_types) _graph_topological_sort(converted.graph) _validate_initializer_output_types(converted) - if requires_attribute_validation or external_attribute_tensors or requires_container_validation: + if supports_strict_type_validation: _validate_converted_types(converted) _validate_local_function_conversion(converted) diff --git a/src/winml/modelkit/quant/passes/fp16.py b/src/winml/modelkit/quant/passes/fp16.py index 6669fda02..02bb13979 100644 --- a/src/winml/modelkit/quant/passes/fp16.py +++ b/src/winml/modelkit/quant/passes/fp16.py @@ -29,6 +29,7 @@ class FP16Pass(BaseQuantPass): - ``fp16_keep_io_types`` — keep model inputs/outputs in their original dtype - ``fp16_op_block_list`` — op types that must not be cast to FP16 + - ``fp16_nodes_to_exclude`` — exact node names that must not be cast to FP16 Example:: @@ -64,6 +65,7 @@ def run( model, keep_io_types=self._config.fp16_keep_io_types, op_block_list=self._config.fp16_op_block_list, + node_block_list=self._config.fp16_nodes_to_exclude, ) output_path.parent.mkdir(parents=True, exist_ok=True) save_onnx(model, output_path, use_external_data=use_external_data) diff --git a/tests/unit/config/test_build.py b/tests/unit/config/test_build.py index fe9c494e5..8872efeeb 100644 --- a/tests/unit/config/test_build.py +++ b/tests/unit/config/test_build.py @@ -2912,6 +2912,56 @@ def test_from_dict_reads_model_id(self) -> None: assert config.model_id == "microsoft/resnet-50" +class TestFP16NodeExclusions: + def test_absent_null_and_non_empty_round_trip(self) -> None: + assert WinMLQuantizationConfig.from_dict({"mode": "fp16"}).fp16_nodes_to_exclude is None + assert ( + WinMLQuantizationConfig.from_dict( + {"mode": "fp16", "fp16_nodes_to_exclude": None} + ).fp16_nodes_to_exclude + is None + ) + config = WinMLQuantizationConfig.from_dict( + {"mode": "fp16", "fp16_nodes_to_exclude": ["cast_b", "cast_a", "cast_b"]} + ) + assert config.fp16_nodes_to_exclude == ["cast_b", "cast_a"] + assert WinMLQuantizationConfig.from_dict(config.to_dict()) == config + + def test_field_is_serialized_only_for_fp16_mode(self) -> None: + assert "fp16_nodes_to_exclude" not in WinMLQuantizationConfig().to_dict() + assert WinMLQuantizationConfig(mode="fp16").to_dict()["fp16_nodes_to_exclude"] is None + + def test_non_string_or_empty_entries_are_rejected(self) -> None: + with pytest.raises(TypeError, match="must be a list"): + WinMLQuantizationConfig.from_dict({"mode": "fp16", "fp16_nodes_to_exclude": "cast"}) + with pytest.raises(TypeError, match="must be a list"): + WinMLQuantizationConfig( + mode="fp16", + fp16_nodes_to_exclude=("cast",), # type: ignore[arg-type] + ) + with pytest.raises(ValueError, match="non-empty strings"): + WinMLQuantizationConfig(mode="fp16", fp16_nodes_to_exclude=[""]) + with pytest.raises(ValueError, match="non-empty strings"): + WinMLQuantizationConfig.from_dict( + {"mode": "fp16", "fp16_nodes_to_exclude": ["cast", 1]} + ) + + def test_node_exclusions_affect_cache_key_but_null_matches_omitted(self) -> None: + omitted = WinMLBuildConfig(quant=WinMLQuantizationConfig(mode="fp16")) + explicit_null = WinMLBuildConfig( + quant=WinMLQuantizationConfig.from_dict({"mode": "fp16", "fp16_nodes_to_exclude": None}) + ) + cast_a = WinMLBuildConfig( + quant=WinMLQuantizationConfig(mode="fp16", fp16_nodes_to_exclude=["cast_a"]) + ) + cast_b = WinMLBuildConfig( + quant=WinMLQuantizationConfig(mode="fp16", fp16_nodes_to_exclude=["cast_b"]) + ) + + assert omitted.generate_cache_key() == explicit_null.generate_cache_key() + assert cast_a.generate_cache_key() != cast_b.generate_cache_key() + + # ============================================================================= # TestInt16QuantTypes - Tests for int16/uint16 quantization type support # ============================================================================= diff --git a/tests/unit/optim/test_fp16.py b/tests/unit/optim/test_fp16.py index caa2f4486..b52e80cd1 100644 --- a/tests/unit/optim/test_fp16.py +++ b/tests/unit/optim/test_fp16.py @@ -15,10 +15,9 @@ from __future__ import annotations -from typing import TYPE_CHECKING - import numpy as np import onnxruntime as ort +import pytest from google.protobuf.message import EncodeError from onnx import ( AttributeProto, @@ -32,13 +31,10 @@ shape_inference, ) +import winml.modelkit.quant.fp16 as fp16_module from winml.modelkit.quant.fp16 import convert_to_fp16 -if TYPE_CHECKING: - import pytest - - # ============================================================================= # HELPERS # ============================================================================= @@ -3330,6 +3326,188 @@ def _branch(name: str, initializer_name: str, value: float) -> GraphProto: class TestConvertToFP16: """Test convert_to_fp16 utility function.""" + def test_exact_cast_node_exclusion_preserves_only_selected_float_boundary(self) -> None: + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1]) + token_ids = helper.make_tensor_value_info("token_ids", TensorProto.INT32, [1]) + selected_output = helper.make_tensor_value_info("selected_out", TensorProto.FLOAT, [1]) + neighbor_output = helper.make_tensor_value_info("neighbor_out", TensorProto.FLOAT, [1]) + embeddings_output = helper.make_tensor_value_info("embeddings_out", TensorProto.FLOAT, [1]) + selected = helper.make_node( + "Cast", + ["x"], + ["selected_out"], + name="InsertedPrecisionFreeCast_selected", + to=TensorProto.FLOAT, + ) + neighbor = helper.make_node( + "Cast", + ["x"], + ["neighbor_out"], + name="InsertedPrecisionFreeCast_neighbor", + to=TensorProto.FLOAT, + ) + embeddings_cast = helper.make_node( + "Cast", + ["token_ids"], + ["embeddings_out"], + name="InsertedPrecisionFreeCast_/deberta/embeddings/Cast_output_0", + to=TensorProto.FLOAT, + ) + model = helper.make_model( + helper.make_graph( + [selected, neighbor, embeddings_cast], + "named_cast_boundaries", + [x, token_ids], + [selected_output, neighbor_output, embeddings_output], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + + result = convert_to_fp16( + model, + keep_io_types=False, + node_block_list=["InsertedPrecisionFreeCast_selected"], + ) + cast_targets = { + node.name: next(attribute.i for attribute in node.attribute if attribute.name == "to") + for node in result.graph.node + if node.op_type == "Cast" and node.name.startswith("InsertedPrecisionFreeCast_") + } + + assert cast_targets["InsertedPrecisionFreeCast_selected"] == TensorProto.FLOAT + assert cast_targets["InsertedPrecisionFreeCast_neighbor"] == TensorProto.FLOAT16 + assert ( + cast_targets["InsertedPrecisionFreeCast_/deberta/embeddings/Cast_output_0"] + == TensorProto.FLOAT16 + ) + checker.check_model(result) + shape_inference.infer_shapes(result, check_type=True, strict_mode=True) + ort.InferenceSession(result.SerializeToString(), providers=["CPUExecutionProvider"]) + + def test_unknown_node_exclusion_has_no_heuristic_effect(self) -> None: + original = _build_simple_fp32_model() + expected = convert_to_fp16(ModelProto.FromString(original.SerializeToString())) + actual = convert_to_fp16(original, node_block_list=["unknown_model_specific_prefix"]) + + assert actual.SerializeToString() == expected.SerializeToString() + + def test_strict_validation_rejects_ordinary_top_level_type_mismatch( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + model = _build_simple_fp32_model() + + def incompatible_conversion(model_arg, **kwargs): + model_arg.graph.output[0].type.tensor_type.elem_type = TensorProto.INT64 + return model_arg + + monkeypatch.setattr( + "onnxruntime.transformers.float16.convert_float_to_float16", + incompatible_conversion, + ) + + with np.testing.assert_raises_regex(RuntimeError, "incompatible FP16 types"): + convert_to_fp16(model) + + def test_strict_capable_conversion_runs_converted_strict_sequence( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + model = _build_simple_fp32_model() + calls: list[bool] = [] + validate = fp16_module._validate_converted_types + + def recording_validate(model_arg: ModelProto, *, required: bool = True) -> bool: + calls.append(required) + return validate(model_arg, required=required) + + monkeypatch.setattr(fp16_module, "_validate_converted_types", recording_validate) + + convert_to_fp16(model) + + assert calls == [False, True] + + @pytest.mark.parametrize( + ("failure_source", "error"), + [ + ("checker", AttributeError("missing checker capability")), + ("inference", EncodeError("model cannot be serialized")), + ("checker", checker.ValidationError("checker capability unavailable")), + ("inference", shape_inference.InferenceError("inference capability unavailable")), + ], + ids=["attribute-error", "encode-error", "validation-error", "inference-error"], + ) + def test_input_capability_exceptions_skip_only_converted_generic_strict_sequence( + self, + monkeypatch: pytest.MonkeyPatch, + failure_source: str, + error: Exception, + ) -> None: + model = _build_simple_fp32_model() + original = model.SerializeToString() + converter_calls = 0 + infer_shapes = shape_inference.infer_shapes + + def fail_checker(model_arg: ModelProto) -> None: + del model_arg + raise error + + def fail_inference(model_arg: ModelProto, **kwargs: object) -> ModelProto: + if kwargs.get("check_type") is True and kwargs.get("strict_mode") is True: + raise error + return infer_shapes(model_arg, **kwargs) + + def unchanged_conversion(model_arg: ModelProto, **kwargs: object) -> ModelProto: + nonlocal converter_calls + del kwargs + converter_calls += 1 + return model_arg + + if failure_source == "checker": + monkeypatch.setattr(checker, "check_model", fail_checker) + else: + monkeypatch.setattr(shape_inference, "infer_shapes", fail_inference) + monkeypatch.setattr( + "onnxruntime.transformers.float16.convert_float_to_float16", + unchanged_conversion, + ) + + result = convert_to_fp16(model) + + assert result is model + assert converter_calls == 1 + assert model.SerializeToString() == original + + def test_input_capability_preflight_does_not_swallow_arbitrary_exceptions( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + model = _build_simple_fp32_model() + + def fail_checker(model_arg: ModelProto) -> None: + del model_arg + raise ValueError("unexpected checker failure") + + monkeypatch.setattr(checker, "check_model", fail_checker) + + with pytest.raises(ValueError, match="unexpected checker failure"): + convert_to_fp16(model) + + def test_specialized_external_data_safeguard_precedes_generic_strict_gate( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + model = _build_external_float_attribute_model(clear_data=True) + + def unexpected_generic_validation(*args: object, **kwargs: object) -> bool: + del args, kwargs + pytest.fail("generic strict validation must not replace specialized safeguards") + + monkeypatch.setattr( + fp16_module, + "_validate_converted_types", + unexpected_generic_validation, + ) + + with pytest.raises(RuntimeError, match="unloaded external data"): + convert_to_fp16(model, keep_io_types=False, op_block_list=[]) + def test_converts_weights_to_fp16(self) -> None: """FP16 conversion converts float32 initializers to float16.""" model = _build_simple_fp32_model() diff --git a/tests/unit/recipes/test_cpu_recipes.py b/tests/unit/recipes/test_cpu_recipes.py index bc3106c57..4df5730bf 100644 --- a/tests/unit/recipes/test_cpu_recipes.py +++ b/tests/unit/recipes/test_cpu_recipes.py @@ -13,6 +13,24 @@ REPO_ROOT = Path(__file__).resolve().parents[3] +MXBAI_FP16_BOUNDARY_CASTS = [ + "InsertedPrecisionFreeCast_/deberta/embeddings/LayerNorm/LayerNormalization_output_0", + *[ + f"InsertedPrecisionFreeCast_/deberta/encoder/layer.{layer}/attention/self/{boundary}_output_0" + for layer in range(12) + for boundary in ( + "Transpose_3", + "Reshape_1", + "Reshape_3", + "Transpose_8", + "Reshape_5", + "Reshape_13", + ) + ], + "InsertedPrecisionFreeCast_/pooler/Gather_output_0", +] +MXBAI_EMBEDDINGS_INT32_CAST = "InsertedPrecisionFreeCast_/deberta/embeddings/Cast_output_0" + recipes = [ { "path": REPO_ROOT @@ -77,3 +95,26 @@ def test_cpu_recipes(rec): else: assert config.quant is not None assert config.quant.mode == rec["quant_mode"] + + +def test_mxbai_fp16_recipe_has_exact_frozen_boundary_casts() -> None: + path = ( + REPO_ROOT + / "examples" + / "recipes" + / "mixedbread-ai_mxbai-rerank-base-v1" + / "cpu" + / "cpu" + / "reranking_fp16_config.json" + ) + data = json.loads(path.read_text(encoding="utf-8")) + exclusions = data["quant"]["fp16_nodes_to_exclude"] + + assert exclusions == MXBAI_FP16_BOUNDARY_CASTS + assert len(exclusions) == 74 + assert len(set(exclusions)) == 74 + assert MXBAI_EMBEDDINGS_INT32_CAST not in exclusions + + config = WinMLBuildConfig.from_dict(data) + assert config.quant is not None + assert config.quant.fp16_nodes_to_exclude == MXBAI_FP16_BOUNDARY_CASTS diff --git a/tests/unit/test_quant_passes.py b/tests/unit/test_quant_passes.py index fb269116b..e6a7263f6 100644 --- a/tests/unit/test_quant_passes.py +++ b/tests/unit/test_quant_passes.py @@ -306,11 +306,13 @@ class TestFP16PassConfig: def test_reads_fp16_fields_from_config( self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: - """FP16Pass should pass fp16_keep_io_types and fp16_op_block_list to convert_to_fp16.""" + """FP16Pass should pass all FP16 controls to convert_to_fp16.""" + selected_nodes = ["selected_cast_2", "selected_cast_1", "selected_cast_3"] config = WinMLQuantizationConfig( mode="fp16", fp16_keep_io_types=False, fp16_op_block_list=["Gather"], + fp16_nodes_to_exclude=selected_nodes, ) model_path = tmp_path / "model.onnx" model_path.write_text("x") @@ -319,8 +321,14 @@ def test_reads_fp16_fields_from_config( calls: list[dict] = [] fake_model = SimpleNamespace() - def fake_convert(model, *, keep_io_types, op_block_list): - calls.append({"keep_io_types": keep_io_types, "op_block_list": op_block_list}) + def fake_convert(model, *, keep_io_types, op_block_list, node_block_list): + calls.append( + { + "keep_io_types": keep_io_types, + "op_block_list": op_block_list, + "node_block_list": node_block_list, + } + ) return model # Patch the source modules that are lazily imported inside run() @@ -336,7 +344,37 @@ def fake_convert(model, *, keep_io_types, op_block_list): result = FP16Pass(config).run(model_path, output_path) assert result.success - assert calls == [{"keep_io_types": False, "op_block_list": ["Gather"]}] + assert calls == [ + { + "keep_io_types": False, + "op_block_list": ["Gather"], + "node_block_list": selected_nodes, + } + ] + + def test_none_node_exclusions_remain_none( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + config = WinMLQuantizationConfig(mode="fp16") + model_path = tmp_path / "model.onnx" + model_path.write_text("x") + calls: list[list[str] | None] = [] + fake_model = SimpleNamespace() + + fake_onnx_mod = ModuleType("winml.modelkit.onnx") + fake_onnx_mod.load_onnx = lambda *a, **k: fake_model # type: ignore[attr-defined] + fake_onnx_mod.save_onnx = lambda *a, **k: None # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "winml.modelkit.onnx", fake_onnx_mod) + + fake_fp16_mod = ModuleType("winml.modelkit.quant.fp16") + fake_fp16_mod.convert_to_fp16 = ( # type: ignore[attr-defined] + lambda model, **kwargs: calls.append(kwargs["node_block_list"]) or model + ) + monkeypatch.setitem(sys.modules, "winml.modelkit.quant.fp16", fake_fp16_mod) + + FP16Pass(config).run(model_path, tmp_path / "out.onnx") + + assert calls == [None] class TestFP16Conversion: @@ -346,6 +384,7 @@ def test_retries_without_shape_inference_when_proto_serialization_fails( """Large external-data models can exceed protobuf's in-memory serialize limit.""" calls: list[dict] = [] model = SimpleNamespace(graph=SimpleNamespace(initializer=[], node=[])) + selected_nodes = ["selected_cast_2", "selected_cast_1", "selected_cast_3"] def fake_convert(model_arg, **kwargs): calls.append(kwargs) @@ -367,15 +406,21 @@ def fake_convert(model_arg, **kwargs): model, keep_io_types=True, op_block_list=["Softmax"], + node_block_list=selected_nodes, ) is model ) assert calls == [ - {"keep_io_types": True, "op_block_list": ["Softmax"]}, + { + "keep_io_types": True, + "op_block_list": ["Softmax"], + "node_block_list": selected_nodes, + }, { "keep_io_types": True, "disable_shape_infer": True, "op_block_list": ["Softmax"], + "node_block_list": selected_nodes, }, ]