fix: flatten_tree fails loudly on flattened-key collisions instead of corrupting silently

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):
@@ -128,6 +132,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)

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)