From b5e859dc61a3a6ad10daa4c669bdff7c195751e3 Mon Sep 17 00:00:00 2001 From: arham766 Date: Sun, 5 Jul 2026 18:48:08 -0700 Subject: [PATCH 1/2] fix: flatten_tree fails loudly on flattened-key collisions instead of corrupting Flattened names join nested keys with '.', and the dict() pass deduped colliding names AND their values while TreeSpec.spec kept both leaf slots - so {'a': {'b': 1}, 'a.b': 2} silently dropped a value and the later unflatten_tree crashed with IndexError far from the root cause (observable at inference time through the ONNX deploy paths, which flatten user-supplied input/output mappings where dotted state-dict style keys are idiomatic). Duplicate flattened names now raise a ValueError naming the colliding key; distinct keys are untouched. Part of the findings in #1902 (wave 2). Signed-off-by: arham766 --- modelopt/torch/utils/_pytree.py | 16 ++++++++++++--- tests/unit/torch/utils/test_pytree.py | 28 +++++++++++++++++++++++++++ 2 files changed, 41 insertions(+), 3 deletions(-) diff --git a/modelopt/torch/utils/_pytree.py b/modelopt/torch/utils/_pytree.py index 586c92b893e..4499081ffa0 100644 --- a/modelopt/torch/utils/_pytree.py +++ b/modelopt/torch/utils/_pytree.py @@ -128,6 +128,16 @@ def collect_spec(pytree, prefix): yield prefix, pytree # retrieve flattened values and names. Then initialize tree_spec with the flattened names. - flattened = dict(collect_spec(pytree, prefix)) - - return list(flattened.values()), TreeSpec(pytree, list(flattened.keys())) + seen_names: set[str] = set() + names: list[str] = [] + values: list[Any] = [] + for name, value in collect_spec(pytree, prefix): + if name in seen_names: + raise ValueError( + f"Cannot flatten pytree: multiple leaves map to the flattened key {name!r}!" + ) + seen_names.add(name) + names.append(name) + values.append(value) + + return values, TreeSpec(pytree, names) diff --git a/tests/unit/torch/utils/test_pytree.py b/tests/unit/torch/utils/test_pytree.py index 2dd95265528..4b6e7e9dee4 100644 --- a/tests/unit/torch/utils/test_pytree.py +++ b/tests/unit/torch/utils/test_pytree.py @@ -72,3 +72,31 @@ def test_flatten_tree(data, expected_keys): # try re-building to see if we get same structure tree_rebuilt = unflatten_tree(copy.deepcopy(values), tree_spec) assert data == tree_rebuilt + + +@pytest.mark.parametrize( + "data", + [ + {"a": {"b": 1}, "a.b": 2}, + {"a": [1], "a.0": 2}, + {"a.b": 1, "a": {"b": 2}}, + ], +) +def test_flatten_tree_key_collision(data): + """Colliding flattened keys must raise instead of silently dropping values.""" + with pytest.raises(ValueError, match=r"'a\.(b|0)'"): + flatten_tree(data) + + +@pytest.mark.parametrize( + ("data", "expected_keys"), + [ + ({"a": {"b": 1}, "ab": 2}, ["a.b", "ab"]), + ({"a": {"b": 1}, "a_b": 2}, ["a.b", "a_b"]), + ], +) +def test_flatten_tree_near_collision(data, expected_keys): + """Distinct flattened keys that merely look similar must still round-trip.""" + values, tree_spec = flatten_tree(data) + assert expected_keys == tree_spec.names + assert data == unflatten_tree(copy.deepcopy(values), tree_spec) From b663bed895657302a2e509e0ae7980bd58cbeeef Mon Sep 17 00:00:00 2001 From: arham766 Date: Sun, 5 Jul 2026 19:14:23 -0700 Subject: [PATCH 2/2] docs: add Raises section for the flattened-key collision ValueError Signed-off-by: arham766 --- modelopt/torch/utils/_pytree.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/modelopt/torch/utils/_pytree.py b/modelopt/torch/utils/_pytree.py index 4499081ffa0..7ff4078f4f7 100644 --- a/modelopt/torch/utils/_pytree.py +++ b/modelopt/torch/utils/_pytree.py @@ -114,6 +114,10 @@ def flatten_tree(pytree: Any, prefix: str = "") -> tuple[list[Any], TreeSpec]: Returns: A tuple (values, pytree) where values is a list of values flattened from the provided pytree, and tree_spec is the pytree spec describing the structure of the pytree. + + Raises: + ValueError: If two leaves map to the same flattened key (nested keys + are joined with ".", so ``{"a": {"b": 1}, "a.b": 2}`` collides). """ def collect_spec(pytree, prefix):