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

[arham766]

What does this PR do?

Type of change: Bug fix

Silent value loss + delayed crash in flatten_tree (found writing #1916): dotted-key collisions ({'a': {'b': 1}, 'a.b': 2}) were deduped by a dict() pass that dropped a VALUE while TreeSpec kept both leaf slots — unflatten_tree then died with IndexError far from the cause. The ONNX deploy paths flatten user-supplied input/output mappings where state-dict-style dotted keys are idiomatic, so this was reachable at inference time. Duplicate flattened names now raise ValueError naming the colliding key. Caller audit and near-collision round-trip tests included; collision tests verified to fail pre-fix (silent loss reproduced).

Usage

N/A

Testing

Regression tests included, each verified to FAIL with the fix reverted. Combined battery with all sibling wave fixes: 381 passed across tests/unit/torch/opt, tests/unit/torch/utils, tests/unit/recipe, and quantization test_mode. Note: the unmerged suite PRs #1913-#1916 contain behavior-documenting NOTE tests pinning the OLD behavior fixed here — whichever lands second gets the NOTE tests flipped (happy to rebase either way).

Before your PR is "Ready for review"

• Is this change backward compatible?: ✅
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A
• Did you write any new necessary tests?: ✅
• Did you update Changelog?: N/A
• Did you get Claude approval on this PR?: N/A (external contributor)

Additional Information

Issue: #1902 (wave-2 findings)

Summary by CodeRabbit

• Bug Fixes
  • Fixed tree flattening to raise an error when flattened keys collide, preventing silent overwrites.
  • Improved key-collision handling to reliably reject true flattened-key conflicts while allowing near-matching keys to round-trip correctly.
• Tests
  • Added unit tests covering flattened key collisions and non-colliding “near-looking” keys.
• Documentation
  • Clarified flattening behavior when duplicate flattened keys would be produced.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 6938875f-826c-4d50-a99e-be396b7aefc9

📥 Commits

Reviewing files that changed from the base of the PR and between b5e859d and b663bed.

📒 Files selected for processing (1)

• modelopt/torch/utils/_pytree.py

🚧 Files skipped from review as they are similar to previous changes (1)

• modelopt/torch/utils/_pytree.py

---

📝 Walkthrough

Walkthrough

flatten_tree now detects duplicate flattened keys and raises ValueError instead of overwriting them. The tests cover both colliding key inputs and near-collision inputs that should still round-trip correctly.

Changes

flatten_tree duplicate key detection

Layer / File(s)	Summary
Duplicate key detection logic and tests modelopt/torch/utils/_pytree.py, tests/unit/torch/utils/test_pytree.py	flatten_tree now tracks flattened key names during traversal, raises ValueError on duplicates, and the new unit tests verify collision failure and near-collision round-tripping.

Estimated code review effort: 2 (Simple) | ~10 minutes

Related PRs: None specified.
Suggested labels: bug, tests
Suggested reviewers: None specified.

A flattened key should stay unique,
Collisions now fail, not overwrite sneak.
Near-matches pass through, still intact,
Round-trips return the original stack.

🚥 Pre-merge checks | ✅ 6

✅ Passed checks (6 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately summarizes the main change: collision detection in flatten_tree with an explicit failure instead of silent corruption.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Security Anti-Patterns	✅ Passed	The PR only changes flatten_tree and its tests; no new torch.load/numpy.load/trust_remote_code/eval-exec/nosec/dependency violations were added.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

🧹 Nitpick comments (1)

modelopt/torch/utils/_pytree.py (1)

107-117: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Consider documenting the new ValueError behavior.

flatten_tree can now raise ValueError on colliding flattened keys, but the docstring has no Raises section.

📝 Suggested docstring addition

     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 distinct leaves map to the same flattened key (e.g. ``{"a": {"b": 1}, "a.b": 2}``).
     """

As per coding guidelines, "Document public APIs with docstrings, including examples when useful."

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@modelopt/torch/utils/_pytree.py` around lines 107 - 117, Update the public
API docstring for flatten_tree in _pytree.py to document its new ValueError
behavior. Add a Raises section that mentions the collision case for flattened
keys, and keep the existing description of the function’s return values and
prefix handling aligned with the implementation. Ensure the docstring for
flatten_tree clearly tells callers that colliding flattened keys will fail with
ValueError.

Source: Coding guidelines

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@modelopt/torch/utils/_pytree.py`:
- Around line 107-117: Update the public API docstring for flatten_tree in
_pytree.py to document its new ValueError behavior. Add a Raises section that
mentions the collision case for flattened keys, and keep the existing
description of the function’s return values and prefix handling aligned with the
implementation. Ensure the docstring for flatten_tree clearly tells callers that
colliding flattened keys will fail with ValueError.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: ca2d9879-a098-4734-814e-1c4929002a2a

📥 Commits

Reviewing files that changed from the base of the PR and between b96a785 and b5e859d.

📒 Files selected for processing (2)

• modelopt/torch/utils/_pytree.py
• tests/unit/torch/utils/test_pytree.py