From 35d04dac963235c23e141a22402c1a03eae3241d Mon Sep 17 00:00:00 2001 From: Zhenchao Ni Date: Mon, 14 Sep 2026 14:25:36 +0800 Subject: [PATCH 1/2] Fix fp16 quantization --- src/winml/modelkit/quant/fp16.py | 109 ++++++++++++++ tests/unit/optim/test_fp16.py | 237 +++++++++++++++++++++++++------ 2 files changed, 300 insertions(+), 46 deletions(-) diff --git a/src/winml/modelkit/quant/fp16.py b/src/winml/modelkit/quant/fp16.py index a86521286..bfce9c90d 100644 --- a/src/winml/modelkit/quant/fp16.py +++ b/src/winml/modelkit/quant/fp16.py @@ -1775,6 +1775,110 @@ def reserve_late_tensor( reserve_late_tensor(graph, node, "output", output_index) +def _rename_generated_io_cast_name_collisions( + model: ModelProto, + *, + keep_io_types: bool, + op_block_list: list[str] | None, +) -> None: + """Free ORT's I/O Cast names without changing public I/O or tensor bindings. + + Only top-level internal tensors are renamed. Captures are updated in all + descendant graphs, including blocked graphs, but local definitions shadow + the rename. Public I/O collisions are left to the existing rejection guard. + Inferred types select reserved names and conversion-relevant nodes without + copying inferred metadata back into the model. + """ + if not keep_io_types: + return + + inferred_model = _ort_inference_preflight_model(model) + name_mapping, reserved_nodes = _ort_keep_io_name_mapping( + inferred_model, keep_io_types=keep_io_types + ) + reserved_tensors = set(name_mapping.values()) + public_names = { + value.name + for values in (getattr(model.graph, "input", []), getattr(model.graph, "output", [])) + for value in values + } + if reserved_tensors & public_names: + return + + tensor_collisions = reserved_tensors & _graph_tensor_names(model.graph) + nodes_to_rename = [ + node + for graph, inferred_graph in zip( + _ort_traversed_graphs(model, op_block_list), + _ort_traversed_graphs(inferred_model, op_block_list), + strict=True, + ) + for node, inferred_node in zip(graph.node, inferred_graph.node, strict=True) + if node.name in reserved_nodes + and not _node_is_conversion_neutral(inferred_model, inferred_graph, inferred_node) + ] + if not tensor_collisions and not nodes_to_rename: + return + + graphs = _all_graphs(model) + occupied_names = reserved_nodes | reserved_tensors + for graph in graphs: + occupied_names.update(_graph_tensor_names(graph)) + occupied_names.update(node.name for node in getattr(graph, "node", [])) + for annotation in getattr(graph, "quantization_annotation", []): + occupied_names.add(annotation.tensor_name) + occupied_names.update(entry.value for entry in annotation.quant_parameter_tensor_names) + + def allocate_name(name: str) -> str: + suffix = 1 + candidate = f"{name}__existing_{suffix}" + while candidate in occupied_names: + suffix += 1 + candidate = f"{name}__existing_{suffix}" + occupied_names.add(candidate) + return candidate + + tensor_renames = {name: allocate_name(name) for name in sorted(tensor_collisions)} + node_renames = [(node, allocate_name(node.name)) for node in nodes_to_rename] + graph_renames = [(model.graph, tensor_renames)] + for graph, renames in graph_renames: + for child in _iter_all_child_graphs(graph): + child_renames = { + name: replacement + for name, replacement in renames.items() + if not _graph_defines_name(child, name) + } + if child_renames: + graph_renames.append((child, child_renames)) + + for graph, renames in graph_renames: + for node in getattr(graph, "node", []): + for values in (node.input, node.output): + for index, name in enumerate(values): + if name in renames: + values[index] = renames[name] + for values in ( + getattr(graph, "input", []), + getattr(graph, "output", []), + getattr(graph, "value_info", []), + getattr(graph, "initializer", []), + ): + for value in values: + if value.name in renames: + value.name = renames[value.name] + for sparse in getattr(graph, "sparse_initializer", []): + if sparse.values.name in renames: + sparse.values.name = renames[sparse.values.name] + for annotation in getattr(graph, "quantization_annotation", []): + if annotation.tensor_name in renames: + annotation.tensor_name = renames[annotation.tensor_name] + for entry in annotation.quant_parameter_tensor_names: + if entry.value in renames: + entry.value = renames[entry.value] + for node, replacement in node_renames: + node.name = replacement + + def _reject_generated_io_cast_name_collisions( model: ModelProto, *, @@ -2891,6 +2995,11 @@ def convert_to_fp16( _reject_sparse_initializer_tensor_metadata(model, op_block_list) _reject_duplicate_float_initializer_names(model, op_block_list) + _rename_generated_io_cast_name_collisions( + model, + keep_io_types=keep_io_types, + op_block_list=op_block_list, + ) io_preflight_model = _ort_inference_preflight_model(model) blocked_ops = _effective_blocked_ops(op_block_list) _reject_unpreserved_float_container_io( diff --git a/tests/unit/optim/test_fp16.py b/tests/unit/optim/test_fp16.py index acb505a29..43260f416 100644 --- a/tests/unit/optim/test_fp16.py +++ b/tests/unit/optim/test_fp16.py @@ -19,6 +19,7 @@ import numpy as np import onnxruntime as ort +import pytest from google.protobuf.message import EncodeError from onnx import ( AttributeProto, @@ -36,7 +37,7 @@ if TYPE_CHECKING: - import pytest + from collections.abc import Callable # ============================================================================= @@ -556,6 +557,36 @@ def _branch(name: str) -> GraphProto: return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) +def _build_io_cast_name_collision_model() -> ModelProto: + """Build existing I/O Casts, fan-out, and an occupied rename candidate.""" + condition = helper.make_tensor_value_info("condition", TensorProto.BOOL, []) + source = helper.make_tensor_value_info("source", TensorProto.FLOAT, [1, 4]) + flag = helper.make_tensor_value_info("flag", TensorProto.BOOL, []) + result = helper.make_tensor_value_info("result", TensorProto.FLOAT, [1, 4]) + nodes = [ + helper.make_node("Identity", [condition.name], [flag.name], name="flag_identity"), + helper.make_node( + "Cast", [source.name], ["graph_input_cast_1"], + name="graph_input_cast1", to=TensorProto.FLOAT, + ), + helper.make_node( + "Identity", ["graph_input_cast_1"], ["graph_input_cast_1__existing_1"], + name="graph_input_cast1__existing_1", + ), + helper.make_node( + "Add", ["graph_input_cast_1", "graph_input_cast_1__existing_1"], + ["graph_output_cast_1"], name="sum", + ), + helper.make_node( + "Cast", ["graph_output_cast_1"], [result.name], + name="graph_output_cast1", to=TensorProto.FLOAT, + ), + ] + graph = helper.make_graph(nodes, "io_cast_collision", [condition, source], [flag, result]) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + return shape_inference.infer_shapes(model, check_type=True, strict_mode=True) + + def _build_initializer_output_name_collision_model() -> ModelProto: """Build a legal graph with a name that collides with ORT's generated alias.""" existing = helper.make_tensor_value_info("graph_output_cast_0", TensorProto.FLOAT, [1]) @@ -3322,6 +3353,165 @@ def _branch(name: str, initializer_name: str, value: float) -> GraphProto: return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) +# ============================================================================= +# I/O CAST NAME COLLISION TESTS +# ============================================================================= + + +class TestRenameGeneratedIOCastNameCollisions: + """Test name-only preprocessing and its integration with FP16 conversion.""" + + def test_rename_preserves_bindings_and_computation(self) -> None: + """Free reserved names without changing types, attributes, or computation.""" + from winml.modelkit.quant.fp16 import ( + _reject_generated_io_cast_name_collisions, + _rename_generated_io_cast_name_collisions, + ) + + model = _build_io_cast_name_collision_model() + original = ModelProto() + original.CopyFrom(model) + feeds = { + "condition": np.array(True), + "source": np.random.default_rng(0).standard_normal((1, 4)).astype(np.float32), + } + expected = ort.InferenceSession( + original.SerializeToString(), providers=["CPUExecutionProvider"] + ).run(None, feeds) + + _rename_generated_io_cast_name_collisions(model, keep_io_types=True, op_block_list=[]) + + _reject_generated_io_cast_name_collisions(model, keep_io_types=True, op_block_list=[]) + checker.check_model(model) + shape_inference.infer_shapes(model, check_type=True, strict_mode=True) + assert model.graph.input == original.graph.input + assert model.graph.output == original.graph.output + assert len(model.graph.node) == len(original.graph.node) + node_names = [node.name for node in model.graph.node] + tensor_names = [name for node in model.graph.node for name in node.output] + assert len(node_names) == len(set(node_names)) + assert len(tensor_names) == len(set(tensor_names)) + assert model.graph.node[2].name == original.graph.node[2].name + assert model.graph.node[2].output == original.graph.node[2].output + assert model.graph.node[2].input[0] == model.graph.node[1].output[0] + assert model.graph.node[3].input[0] == model.graph.node[1].output[0] + assert model.graph.node[3].input[1] == model.graph.node[2].output[0] + assert model.graph.node[4].input[0] == model.graph.node[3].output[0] + for before, after in zip(original.graph.node, model.graph.node, strict=True): + assert after.op_type == before.op_type + assert after.domain == before.domain + assert after.attribute == before.attribute + for before, after in zip(original.graph.value_info, model.graph.value_info, strict=True): + assert after.type == before.type + assert after.name in tensor_names + actual = ort.InferenceSession( + model.SerializeToString(), providers=["CPUExecutionProvider"] + ).run(None, feeds) + for expected_output, actual_output in zip(expected, actual, strict=True): + assert actual_output.dtype == expected_output.dtype + np.testing.assert_array_equal(actual_output, expected_output) + + def test_rename_is_deterministic_and_idempotent(self) -> None: + """Occupied suffixes must not cause duplicate or unstable fresh names.""" + from winml.modelkit.quant.fp16 import _rename_generated_io_cast_name_collisions + + model = _build_io_cast_name_collision_model() + duplicate = ModelProto() + duplicate.CopyFrom(model) + _rename_generated_io_cast_name_collisions(model, keep_io_types=True, op_block_list=[]) + renamed = model.SerializeToString() + _rename_generated_io_cast_name_collisions(duplicate, keep_io_types=True, op_block_list=[]) + assert duplicate.SerializeToString() == renamed + _rename_generated_io_cast_name_collisions(model, keep_io_types=True, op_block_list=[]) + assert model.SerializeToString() == renamed + + @pytest.mark.parametrize( + ("build_model", "keep_io_types"), + [(_build_io_cast_name_collision_model, False), (_build_simple_fp32_model, True)], + ids=["io-types-not-kept", "no-collisions"], + ) + def test_rename_is_noop_when_not_needed( + self, build_model: Callable[[], ModelProto], keep_io_types: bool + ) -> None: + """Leave unrelated graphs and conversion without kept I/O unchanged.""" + from winml.modelkit.quant.fp16 import _rename_generated_io_cast_name_collisions + + model = build_model() + original = model.SerializeToString() + _rename_generated_io_cast_name_collisions( + model, keep_io_types=keep_io_types, op_block_list=[] + ) + assert model.SerializeToString() == original + + def test_public_io_collision_remains_rejected(self) -> None: + """Renaming internal bindings must not alter the public I/O contract.""" + from winml.modelkit.quant.fp16 import ( + _reject_generated_io_cast_name_collisions, + _rename_generated_io_cast_name_collisions, + ) + + model = _build_initializer_output_name_collision_model() + original = model.SerializeToString() + _rename_generated_io_cast_name_collisions(model, keep_io_types=True, op_block_list=[]) + assert model.SerializeToString() == original + with np.testing.assert_raises_regex(RuntimeError, "existing names collide"): + _reject_generated_io_cast_name_collisions(model, keep_io_types=True, op_block_list=[]) + + @pytest.mark.parametrize( + "build_model", + [ + _build_io_cast_name_collision_model, + _build_initializer_output_node_name_collision_model, + _build_nested_node_name_collision_model, + _build_regular_output_nested_node_name_collision_model, + _build_inferred_output_nested_node_name_collision_model, + ], + ids=["existing-io-casts", "initializer-output", "nested-node", "nested-capture", "inferred-io"], + ) + def test_conversion_resolves_collisions_and_preserves_fp32_io( + self, build_model: Callable[[], ModelProto] + ) -> None: + """The conversion entry point resolves collisions before calling ORT.""" + model = build_model() + reference_model = shape_inference.infer_shapes(model, check_type=True, strict_mode=True) + checker.check_model(reference_model) + reference = ort.InferenceSession( + reference_model.SerializeToString(), providers=["CPUExecutionProvider"] + ) + io_before = [ + value.SerializeToString() + for values in (reference_model.graph.input, reference_model.graph.output) + for value in values + ] + feeds = { + value.name: ( + np.array(True) + if value.type.tensor_type.elem_type == TensorProto.BOOL + else np.random.default_rng(index).standard_normal( + tuple(dimension.dim_value for dimension in value.type.tensor_type.shape.dim) + ).astype(np.float32) + ) + for index, value in enumerate(reference_model.graph.input) + } + expected = reference.run(None, feeds) + + result = convert_to_fp16(model, keep_io_types=True, op_block_list=[]) + + assert result is model + assert [ + value.SerializeToString() + for values in (result.graph.input, result.graph.output) + for value in values + ] == io_before + checker.check_model(result) + shape_inference.infer_shapes(result, check_type=True, strict_mode=True) + session = ort.InferenceSession(result.SerializeToString(), providers=["CPUExecutionProvider"]) + actual = session.run(None, feeds) + for expected_output, actual_output in zip(expected, actual, strict=True): + assert actual_output.dtype == expected_output.dtype + np.testing.assert_allclose(actual_output, expected_output, rtol=1e-3, atol=1e-3) + + # ============================================================================= # CONVERT_TO_FP16 TESTS # ============================================================================= @@ -4102,37 +4292,6 @@ def test_generated_tensor_name_collision_is_rejected_before_mutation(self) -> No convert_to_fp16(model, keep_io_types=True, op_block_list=[]) assert model.SerializeToString() == original - def test_generated_cast_node_name_collision_is_rejected_before_mutation(self) -> None: - """A user node occupying ORT's deterministic Cast name fails safely.""" - model = _build_initializer_output_node_name_collision_model() - original = model.SerializeToString() - - with np.testing.assert_raises_regex(RuntimeError, "existing names collide"): - convert_to_fp16(model, keep_io_types=True, op_block_list=[]) - - assert model.SerializeToString() == original - - def test_nested_generated_node_name_collision_is_rejected_before_mutation(self) -> None: - """Nested nodes also participate in ORT's global generated-name set.""" - model = _build_nested_node_name_collision_model() - original = model.SerializeToString() - - with np.testing.assert_raises_regex(RuntimeError, "existing names collide"): - convert_to_fp16(model, keep_io_types=True, op_block_list=[]) - - assert model.SerializeToString() == original - - def test_regular_output_nested_node_name_collision_is_rejected_before_mutation( - self, - ) -> None: - """Generated I/O Cast names are reserved even without initializer outputs.""" - model = _build_regular_output_nested_node_name_collision_model() - original = model.SerializeToString() - - with np.testing.assert_raises_regex(RuntimeError, "graph_output_cast0"): - convert_to_fp16(model, keep_io_types=True, op_block_list=[]) - assert model.SerializeToString() == original - def test_neutral_nested_node_name_collision_is_allowed(self) -> None: """Skipping an already-concrete INT node cannot alter precision.""" model = _build_neutral_nested_node_name_collision_model() @@ -4184,20 +4343,6 @@ def test_neutral_collision_can_use_shadowed_int_formal( np.testing.assert_array_equal(output, np.array([2.0], dtype=np.float32)) np.testing.assert_array_equal(loop_output, np.array([7], dtype=np.int64)) - def test_inferred_output_nested_node_name_collision_is_rejected_before_mutation( - self, - ) -> None: - """Generated Cast reservations use the same inferred I/O types as ORT.""" - model = _build_inferred_output_nested_node_name_collision_model() - inferred = shape_inference.infer_shapes(model, strict_mode=True) - assert inferred.graph.output[0].type.tensor_type.elem_type == TensorProto.FLOAT - checker.check_model(inferred) - original = model.SerializeToString() - - with np.testing.assert_raises_regex(RuntimeError, "graph_output_cast0"): - convert_to_fp16(model, keep_io_types=True, op_block_list=[]) - assert model.SerializeToString() == original - def test_nested_local_generated_tensor_alias_does_not_collide(self) -> None: """A nested local binding may legally shadow a generated top-level alias.""" model = _build_nested_local_generated_tensor_alias_model() From 5e4add07e5c4981d791ff10f46c2cbadecba8aa6 Mon Sep 17 00:00:00 2001 From: Zhenchao Ni Date: Mon, 14 Sep 2026 15:22:04 +0800 Subject: [PATCH 2/2] Preserve caller model when FP16 collision conversion fails --- src/winml/modelkit/quant/fp16.py | 13 +++++++++---- tests/unit/optim/test_fp16.py | 29 +++++++++++++++++++++++++++-- 2 files changed, 36 insertions(+), 6 deletions(-) diff --git a/src/winml/modelkit/quant/fp16.py b/src/winml/modelkit/quant/fp16.py index bfce9c90d..44ac04534 100644 --- a/src/winml/modelkit/quant/fp16.py +++ b/src/winml/modelkit/quant/fp16.py @@ -2995,6 +2995,9 @@ def convert_to_fp16( _reject_sparse_initializer_tensor_metadata(model, op_block_list) _reject_duplicate_float_initializer_names(model, op_block_list) + original_model = model + if keep_io_types: + model = deepcopy(model) _rename_generated_io_cast_name_collisions( model, keep_io_types=keep_io_types, @@ -3102,7 +3105,9 @@ def convert_to_fp16( if needs_safe_conversion: _reject_unloaded_external_initializer_outputs(model, op_block_list) original_nodes = len(model.graph.node) - conversion_model = deepcopy(model) if needs_safe_conversion else model + conversion_model = ( + deepcopy(model) if needs_safe_conversion and model is original_model else model + ) if needs_safe_conversion: _internalize_external_initializer_outputs(conversion_model, op_block_list) _internalize_selected_external_initializers( @@ -3158,9 +3163,9 @@ def convert_to_fp16( _validate_converted_types(converted) _validate_local_function_conversion(converted) - if converted is not model: - model.CopyFrom(converted) - converted = model + if converted is not original_model: + original_model.CopyFrom(converted) + converted = original_model converted_nodes = len(converted.graph.node) if converted_nodes != original_nodes: diff --git a/tests/unit/optim/test_fp16.py b/tests/unit/optim/test_fp16.py index 43260f416..3027bac40 100644 --- a/tests/unit/optim/test_fp16.py +++ b/tests/unit/optim/test_fp16.py @@ -3443,6 +3443,23 @@ def test_rename_is_noop_when_not_needed( ) assert model.SerializeToString() == original + def test_collision_with_unloaded_weight_is_rejected_before_mutation(self) -> None: + """A later preflight failure must not publish internal collision renames.""" + model = _build_io_cast_name_collision_model() + weight = numpy_helper.from_array( + np.random.default_rng(0).standard_normal((1, 4)).astype(np.float32), + "weight", + ) + model.graph.initializer.append(weight) + model.graph.node[3].input[1] = weight.name + _mark_initializers_as_external(model.graph, clear_data=True) + original = model.SerializeToString() + + with pytest.raises(RuntimeError, match="unloaded external data"): + convert_to_fp16(model, keep_io_types=True, op_block_list=[]) + + assert model.SerializeToString() == original + def test_public_io_collision_remains_rejected(self) -> None: """Renaming internal bindings must not alter the public I/O contract.""" from winml.modelkit.quant.fp16 import ( @@ -3466,7 +3483,13 @@ def test_public_io_collision_remains_rejected(self) -> None: _build_regular_output_nested_node_name_collision_model, _build_inferred_output_nested_node_name_collision_model, ], - ids=["existing-io-casts", "initializer-output", "nested-node", "nested-capture", "inferred-io"], + ids=[ + "existing-io-casts", + "initializer-output", + "nested-node", + "nested-capture", + "inferred-io", + ], ) def test_conversion_resolves_collisions_and_preserves_fp32_io( self, build_model: Callable[[], ModelProto] @@ -3505,7 +3528,9 @@ def test_conversion_resolves_collisions_and_preserves_fp32_io( ] == io_before checker.check_model(result) shape_inference.infer_shapes(result, check_type=True, strict_mode=True) - session = ort.InferenceSession(result.SerializeToString(), providers=["CPUExecutionProvider"]) + session = ort.InferenceSession( + result.SerializeToString(), providers=["CPUExecutionProvider"] + ) actual = session.run(None, feeds) for expected_output, actual_output in zip(expected, actual, strict=True): assert actual_output.dtype == expected_output.dtype