From 26f89fb436e0edca68a947e48e38721da48eedf8 Mon Sep 17 00:00:00 2001 From: Kasper Nielsen Date: Sun, 23 Aug 2026 08:34:56 +0100 Subject: [PATCH] Fix weight_id collision when materializing multiple source functions `materialize_symbolic_shape_program` assigns a `weight_id` to every const it clones that does not have one yet, so that the materialized functions share the weight file value of the const they were cloned from. That `weight_id` was derived from the const name alone, but a const name is only unique within a function, not within a program. So materializing 2 source functions of a multifunction program that each own a same-named const holding a different value (e.g. 2 separately traced functions that both name their leading `arange` `range_1d_0`, one of length 128 and one of length 512) makes those 2 unrelated consts claim the same weight blob entry: only one value is written to the blob, both functions point to it, and the model becomes corrupted. Loading it fails with "Attribute val has incompatible type with operation output", and the saved package cannot be loaded back into pymil either. The collision is invisible for small consts, since those are stored immediately in the proto and never reach the weight blob. Include the source function name in the invented `weight_id`. All the functions materialized from one source function still share the weight file value of that source, and consts that got their `weight_id` elsewhere (`const_deduplication` groups by dtype + shape + value) are untouched, so deduplication across functions keeps working. Tests: - New `test_weight_id_no_collision_across_source_functions` in test_symbol_transform.py: materializes 2 source functions holding same-named consts that are large enough to be blob stored, lowers the program to mlprogram, loads it back into pymil, and checks that every materialized function reads back its own weight. It fails before this change, with `func2_materialized` reading `main`'s weight. - `test_weight_id_pass_down` now expects the source-function-scoped id. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01SSuo4pE3Kr3GWNwcLKLn5v --- .../mil/mil/passes/defs/symbol_transform.py | 11 ++- .../mil/passes/tests/test_symbol_transform.py | 91 ++++++++++++++++++- 2 files changed, 99 insertions(+), 3 deletions(-) diff --git a/coremltools/converters/mil/mil/passes/defs/symbol_transform.py b/coremltools/converters/mil/mil/passes/defs/symbol_transform.py index c81c76f79..a4f5e8314 100644 --- a/coremltools/converters/mil/mil/passes/defs/symbol_transform.py +++ b/coremltools/converters/mil/mil/passes/defs/symbol_transform.py @@ -307,9 +307,18 @@ def apply(self, prog: Program) -> None: ), "Only const may be absent from context" # The consts across function should share the same file value while lowering into milproto, # so we assign the weight_id, if not presented. + # A const name is only unique within a function, so the + # source function name has to be part of the weight_id: + # 2 source functions may own same-named consts that hold + # different values, and sharing a weight file value among + # those would corrupt the model. Note that this only affects + # the weight_id we invent here, so consts that got a weight_id + # elsewhere (e.g. by ``const_deduplication``) keep sharing + # their file value across functions. if source_input_var.op.weight_id is None: source_input_var.op.weight_id = ( - f"const_{source_input_var.name}_weight_id" + f"const_{self.source_function_name}_" + f"{source_input_var.name}_weight_id" ) context[source_input_var.name] = self._copy_construct_const_var( source_input_var diff --git a/coremltools/converters/mil/mil/passes/tests/test_symbol_transform.py b/coremltools/converters/mil/mil/passes/tests/test_symbol_transform.py index df0b73382..a512a0900 100644 --- a/coremltools/converters/mil/mil/passes/tests/test_symbol_transform.py +++ b/coremltools/converters/mil/mil/passes/tests/test_symbol_transform.py @@ -9,8 +9,11 @@ import numpy as np import pytest +import coremltools as ct +from coremltools.converters.mil.converter import mil_convert +from coremltools.converters.mil.frontend.milproto.load import load as _milproto_to_pymil from coremltools.converters.mil.mil import Builder as mb -from coremltools.converters.mil.mil import get_new_symbol +from coremltools.converters.mil.mil import get_new_symbol, types from coremltools.converters.mil.mil.passes.pass_registry import PASS_REGISTRY from coremltools.converters.mil.testing_utils import ( apply_pass_and_basic_check, @@ -45,7 +48,91 @@ def prog(x): for block in prog.functions.values(): const_ops = block.find_ops(op_type="const") assert len(const_ops) == 1 - assert const_ops[0].weight_id == "const_y_weight_id" + assert const_ops[0].weight_id == "const_main_y_weight_id" + + @staticmethod + def test_weight_id_no_collision_across_source_functions(): + """ + A const name is unique only within a function, so 2 source functions may each own + a same-named const that holds a different value. Since the auto assigned weight_id + determines which consts share a weight blob entry, deriving it from the const name + alone would make those 2 unrelated consts claim the same entry, and the resulting + model would be corrupted: one value gets written to the blob, while both functions + point to it, so the function whose shape disagrees with the blob content fails to load + + Note that a const is only written to the weight blob if it is large enough + (small ones are stored immediately in the proto), so this test uses large consts + """ + symbolic_shape = (get_new_symbol(),) + fixed_shape = (1,) + source_function_name_to_weight = { + "main": np.arange(4096, dtype=np.float32), + "func2": np.arange(8192, dtype=np.float32) + 0.5, + } + + @mb.program( + input_specs=[mb.TensorSpec(shape=symbolic_shape, dtype=types.int32)], + opset_version=ct.target.iOS18, + ) + def prog(x): + # gather from a runtime index, so the const cannot be folded away + weight = mb.const(val=source_function_name_to_weight["main"], name="shared_name") + return mb.gather(x=weight, indices=x, axis=0) + + @mb.function( + input_specs=[mb.TensorSpec(shape=symbolic_shape, dtype=types.int32)], + opset_version=ct.target.iOS18, + ) + def func2(x): + weight = mb.const(val=source_function_name_to_weight["func2"], name="shared_name") + return mb.gather(x=weight, indices=x, axis=0) + + prog.add_function("func2", func2) + + for source_function_name in source_function_name_to_weight: + graph_pass = PASS_REGISTRY["common::materialize_symbolic_shape_program"] + graph_pass.source_function_name = source_function_name + graph_pass.function_name_to_materialization_map = { + f"{source_function_name}_materialized": {"x": fixed_shape} + } + graph_pass.apply(prog) + + # the same-named consts of the 2 source functions must not claim the same weight_id + weight_id_to_weight = {} + for function in prog.functions.values(): + for const_op in function.find_ops(prefix="shared_name", op_type="const"): + assert const_op.weight_id is not None + weight = const_op.outputs[0].val + # a materialized const shares the weight blob entry of the source const + # it is cloned from, but consts of different source functions must not + np.testing.assert_array_equal( + weight_id_to_weight.setdefault(const_op.weight_id, weight), weight + ) + assert len(weight_id_to_weight) == len(source_function_name_to_weight) + + # end-to-end: every materialized function has to read back its own weight + prog.export_as_multifunction = True + prog.skip_all_passes = True + mlmodel = mil_convert( + prog, + convert_from="milinternal", + convert_to="mlprogram", + specification_version=ct.target.iOS18, + compute_units=ct.ComputeUnit.CPU_ONLY, + skip_model_load=True, + ) + spec = mlmodel.get_spec() + reloaded_prog = _milproto_to_pymil( + model_spec=spec, + specification_version=spec.specificationVersion, + file_weights_dir=mlmodel.weights_dir, + ) + for source_function_name, weight in source_function_name_to_weight.items(): + const_ops = reloaded_prog.functions[f"{source_function_name}_materialized"].find_ops( + prefix="shared_name", op_type="const" + ) + assert len(const_ops) == 1 + np.testing.assert_array_equal(const_ops[0].outputs[0].val, weight) @pytest.mark.parametrize("override_main_function", (True, False)) def test_simple(self, override_main_function):