Fix FP16 external data relocation

Author: ssss141414

src/winml/modelkit/quant/fp16.py

@@ -17,17 +17,17 @@
 
 
 if TYPE_CHECKING:
-    from onnx import ModelProto
+    import onnx
 
 logger = logging.getLogger(__name__)
 
 
 def convert_to_fp16(
-    model: ModelProto | str | Path,
+    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.
@@ -49,24 +49,29 @@ def convert_to_fp16(
     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}
+    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:
-        import onnx
-
         inspection_model = onnx.load(str(model_path), load_external_data=False)
     else:
-        inspection_model = cast("ModelProto", model)
+        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.")
+            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(inspection_model.graph.node)
@@ -78,8 +83,6 @@ def convert_to_fp16(
         logger.info("  Keeping ops in FP32: %s", op_block_list)
 
     if model_path is not None:
-        from onnx.shape_inference import infer_shapes_path
-
         # 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
@@ -89,12 +92,12 @@ def convert_to_fp16(
         ) as temporary:
             inferred_path = Path(temporary.name)
         try:
-            infer_shapes_path(str(model_path), str(inferred_path))
+            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(
-            "ModelProto",
+            "onnx.ModelProto",
             convert_float_to_float16(
                 inferred_model,
                 keep_io_types=keep_io_types,
@@ -104,7 +107,7 @@ def convert_to_fp16(
         )
     else:
         converted = cast(
-            "ModelProto",
+            "onnx.ModelProto",
             convert_float_to_float16(
                 model,
                 keep_io_types=keep_io_types,

tests/unit/optim/test_fp16.py

@@ -15,10 +15,12 @@
 
 from __future__ import annotations
 
+import shutil
+
 import numpy as np
 import onnx
-from onnx import ModelProto, TensorProto, helper, numpy_helper
 
+import winml.modelkit.onnx
 from winml.modelkit.quant.fp16 import convert_to_fp16
 
 
@@ -27,25 +29,29 @@
 # =============================================================================
 
 
-def _build_simple_fp32_model() -> ModelProto:
+def _build_simple_fp32_model() -> onnx.ModelProto:
     """Build a simple FP32 model: out = x + weight."""
-    x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 4])
-    out = helper.make_tensor_value_info("out", TensorProto.FLOAT, [1, 4])
-    weight = numpy_helper.from_array(np.array([[1.0, 2.0, 3.0, 4.0]], dtype=np.float32), "weight")
-    add = helper.make_node("Add", ["x", "weight"], ["out"], name="add")
-    graph = helper.make_graph([add], "simple", [x], [out], [weight])
-    return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)])
+    x = onnx.helper.make_tensor_value_info("x", onnx.TensorProto.FLOAT, [1, 4])
+    out = onnx.helper.make_tensor_value_info("out", onnx.TensorProto.FLOAT, [1, 4])
+    weight = onnx.numpy_helper.from_array(
+        np.array([[1.0, 2.0, 3.0, 4.0]], dtype=np.float32), "weight"
+    )
+    add = onnx.helper.make_node("Add", ["x", "weight"], ["out"], name="add")
+    graph = onnx.helper.make_graph([add], "simple", [x], [out], [weight])
+    return onnx.helper.make_model(graph, opset_imports=[onnx.helper.make_opsetid("", 17)])
 
 
-def _build_multi_op_fp32_model() -> ModelProto:
+def _build_multi_op_fp32_model() -> onnx.ModelProto:
     """Build a model with multiple ops: out = Relu(x + weight)."""
-    x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 4])
-    out = helper.make_tensor_value_info("out", TensorProto.FLOAT, [1, 4])
-    weight = numpy_helper.from_array(np.array([[1.0, 2.0, 3.0, 4.0]], dtype=np.float32), "weight")
-    add = helper.make_node("Add", ["x", "weight"], ["add_out"], name="add")
-    relu = helper.make_node("Relu", ["add_out"], ["out"], name="relu")
-    graph = helper.make_graph([add, relu], "multi_op", [x], [out], [weight])
-    return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)])
+    x = onnx.helper.make_tensor_value_info("x", onnx.TensorProto.FLOAT, [1, 4])
+    out = onnx.helper.make_tensor_value_info("out", onnx.TensorProto.FLOAT, [1, 4])
+    weight = onnx.numpy_helper.from_array(
+        np.array([[1.0, 2.0, 3.0, 4.0]], dtype=np.float32), "weight"
+    )
+    add = onnx.helper.make_node("Add", ["x", "weight"], ["add_out"], name="add")
+    relu = onnx.helper.make_node("Relu", ["add_out"], ["out"], name="relu")
+    graph = onnx.helper.make_graph([add, relu], "multi_op", [x], [out], [weight])
+    return onnx.helper.make_model(graph, opset_imports=[onnx.helper.make_opsetid("", 17)])
 
 
 # =============================================================================
@@ -61,7 +67,9 @@ def test_converts_weights_to_fp16(self) -> None:
         model = _build_simple_fp32_model()
         result = convert_to_fp16(model)
 
-        has_fp16 = any(init.data_type == TensorProto.FLOAT16 for init in result.graph.initializer)
+        has_fp16 = any(
+            init.data_type == onnx.TensorProto.FLOAT16 for init in result.graph.initializer
+        )
         assert has_fp16, "Expected at least one FP16 initializer after conversion"
 
     def test_path_conversion_uses_external_data_safe_shape_inference(self, tmp_path) -> None:
@@ -79,27 +87,66 @@ def test_path_conversion_uses_external_data_safe_shape_inference(self, tmp_path)
 
         result = convert_to_fp16(model_path)
 
-        assert any(init.data_type == TensorProto.FLOAT16 for init in result.graph.initializer)
+        assert any(init.data_type == onnx.TensorProto.FLOAT16 for init in result.graph.initializer)
+
+    def test_already_fp16_external_data_path_is_relocatable(self, tmp_path) -> None:
+        """Already-FP16 path input can be saved independently of its source sidecar."""
+        source_dir = tmp_path / "source"
+        source_dir.mkdir()
+        source_path = source_dir / "source.onnx"
+        expected_weight = np.arange(1024, dtype=np.float16).reshape(1, 1024)
+        x = onnx.helper.make_tensor_value_info("x", onnx.TensorProto.FLOAT16, [1, 1024])
+        out = onnx.helper.make_tensor_value_info("out", onnx.TensorProto.FLOAT16, [1, 1024])
+        weight = onnx.numpy_helper.from_array(expected_weight, "weight")
+        graph = onnx.helper.make_graph(
+            [onnx.helper.make_node("Add", ["x", "weight"], ["out"])],
+            "external_fp16",
+            [x],
+            [out],
+            [weight],
+        )
+        model = onnx.helper.make_model(graph, opset_imports=[onnx.helper.make_opsetid("", 17)])
+        onnx.save_model(
+            model,
+            str(source_path),
+            save_as_external_data=True,
+            all_tensors_to_one_file=True,
+            location="source.onnx.data",
+            size_threshold=0,
+        )
+
+        converted = convert_to_fp16(source_path)
+        destination_dir = tmp_path / "destination"
+        destination_path = destination_dir / "model.onnx"
+        winml.modelkit.onnx.save_onnx(converted, destination_path, threshold_size=0)
+
+        assert (destination_dir / "model.onnx.data").is_file()
+        shutil.rmtree(source_dir)
+        onnx.checker.check_model(str(destination_path))
+        relocated = onnx.load(str(destination_path))
+        np.testing.assert_array_equal(
+            onnx.numpy_helper.to_array(relocated.graph.initializer[0]), expected_weight
+        )
 
     def test_default_keeps_io_types(self) -> None:
         """Default keep_io_types=True preserves FP32 model I/O."""
         model = _build_simple_fp32_model()
         result = convert_to_fp16(model, keep_io_types=True)
 
         for inp in result.graph.input:
-            assert inp.type.tensor_type.elem_type == TensorProto.FLOAT
+            assert inp.type.tensor_type.elem_type == onnx.TensorProto.FLOAT
         for outp in result.graph.output:
-            assert outp.type.tensor_type.elem_type == TensorProto.FLOAT
+            assert outp.type.tensor_type.elem_type == onnx.TensorProto.FLOAT
 
     def test_keep_io_types_false_converts_io(self) -> None:
         """With keep_io_types=False, model I/O becomes FP16."""
         model = _build_simple_fp32_model()
         result = convert_to_fp16(model, keep_io_types=False)
 
         for inp in result.graph.input:
-            assert inp.type.tensor_type.elem_type == TensorProto.FLOAT16
+            assert inp.type.tensor_type.elem_type == onnx.TensorProto.FLOAT16
         for outp in result.graph.output:
-            assert outp.type.tensor_type.elem_type == TensorProto.FLOAT16
+            assert outp.type.tensor_type.elem_type == onnx.TensorProto.FLOAT16
 
     def test_preserves_model_structure(self) -> None:
         """FP16 conversion preserves graph structure (node count diff ≤ 2)."""
@@ -131,13 +178,13 @@ def test_none_op_block_list_uses_ort_defaults(self) -> None:
     def test_skips_already_fp16_model(self) -> None:
         """If all floating-point initializers are already FP16, conversion is skipped."""
         # Build a model with FP16 initializers directly
-        x = helper.make_tensor_value_info("x", TensorProto.FLOAT16, [1, 4])
-        out = helper.make_tensor_value_info("out", TensorProto.FLOAT16, [1, 4])
+        x = onnx.helper.make_tensor_value_info("x", onnx.TensorProto.FLOAT16, [1, 4])
+        out = onnx.helper.make_tensor_value_info("out", onnx.TensorProto.FLOAT16, [1, 4])
         weight_data = np.array([[1.0, 2.0, 3.0, 4.0]], dtype=np.float16)
-        weight = numpy_helper.from_array(weight_data, "weight")
-        add = helper.make_node("Add", ["x", "weight"], ["out"], name="add")
-        graph = helper.make_graph([add], "fp16_model", [x], [out], [weight])
-        model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)])
+        weight = onnx.numpy_helper.from_array(weight_data, "weight")
+        add = onnx.helper.make_node("Add", ["x", "weight"], ["out"], name="add")
+        graph = onnx.helper.make_graph([add], "fp16_model", [x], [out], [weight])
+        model = onnx.helper.make_model(graph, opset_imports=[onnx.helper.make_opsetid("", 17)])
 
         original_nodes = len(model.graph.node)
         result = convert_to_fp16(model)
@@ -148,15 +195,15 @@ def test_skips_already_fp16_model(self) -> None:
 
     def test_skips_fp16_model_with_int_initializers(self) -> None:
         """FP16 model with non-float initializers (e.g. INT64 shapes) should still skip."""
-        x = helper.make_tensor_value_info("x", TensorProto.FLOAT16, [1, 4])
-        out = helper.make_tensor_value_info("out", TensorProto.FLOAT16, [1, 4])
+        x = onnx.helper.make_tensor_value_info("x", onnx.TensorProto.FLOAT16, [1, 4])
+        out = onnx.helper.make_tensor_value_info("out", onnx.TensorProto.FLOAT16, [1, 4])
         weight_data = np.array([[1.0, 2.0, 3.0, 4.0]], dtype=np.float16)
-        weight = numpy_helper.from_array(weight_data, "weight")
+        weight = onnx.numpy_helper.from_array(weight_data, "weight")
         # INT64 initializer (e.g., shape tensor) — should be ignored by skip logic
-        shape_tensor = numpy_helper.from_array(np.array([1, 4], dtype=np.int64), "shape")
-        add = helper.make_node("Add", ["x", "weight"], ["out"], name="add")
-        graph = helper.make_graph([add], "fp16_mixed", [x], [out], [weight, shape_tensor])
-        model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)])
+        shape_tensor = onnx.numpy_helper.from_array(np.array([1, 4], dtype=np.int64), "shape")
+        add = onnx.helper.make_node("Add", ["x", "weight"], ["out"], name="add")
+        graph = onnx.helper.make_graph([add], "fp16_mixed", [x], [out], [weight, shape_tensor])
+        model = onnx.helper.make_model(graph, opset_imports=[onnx.helper.make_opsetid("", 17)])
 
         original_nodes = len(model.graph.node)
         result = convert_to_fp16(model)