Fix weight_id collision when materializing multiple source functions - #2828
Open
kasper0406 wants to merge 1 commit into
Open
Fix weight_id collision when materializing multiple source functions#2828kasper0406 wants to merge 1 commit into
kasper0406 wants to merge 1 commit into
Conversation
`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 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01SSuo4pE3Kr3GWNwcLKLn5v
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Bug
materialize_symbolic_shape_programassigns aweight_idto cloned consts that lack one, deriving the id from the const's name alone (symbol_transform.py):f"const_{source_input_var.name}_weight_id". Const names are only unique within a function, so when two source functions each contain a same-named const holding different values/shapes (common when functions are traced separately — e.g. both produce a const namedrange_1d_0, oneint32[8]and oneint32[512]), the materialized functions claim the same weight-blob entry (the backend caches blob values keyed byweight_idinbackend/mil/load.py). The serialized package is corrupt: loading fails withAttribute val has incompatible type with operation output, and the package cannot be re-loaded into pymil.The collision is invisible when the affected consts are small (<10 elements are stored inline and never reach the blob), which is why it lurks — we hit it in production the moment a shared const name crossed the inline threshold.
Fix
Scope the invented fallback id by the source function name:
const_{source_function_name}_{name}_weight_id. Consts that already carry aweight_id(e.g. assigned byconst_deduplication, which groups by dtype+shape+value) are untouched, so intentional cross-function weight sharing keeps working, and every function materialized from one source still shares that source's blob entry.Tests
test_weight_id_no_collision_across_source_functions: two source functions with same-named, different-valued blob-sized consts, materialized, checked in-memory (noweight_idmaps to two different weights) and end-to-end (mlprogram serialize + milproto reload, each function reads back its own values). Fails without the fix (shapes (4096,), (8192,) mismatch), passes with it.test_weight_id_pass_downupdated for the scoped id format (the only expectation change the fix forces).test_symbol_transform.py10/10,TestMaterializeSymbolicShapeMLModel17/17 (8 pre-existing xfails),test_cleanup_passes.pyconst-dedup coverage unaffected.🤖 Generated with Claude Code