Name the producer key in rename_weight_keys collision error

Address PR review: when two source keys rename to the same target, the
ValueError now identifies the first producer key (not just "another key"),
which makes the collision actionable when it fires in a large checkpoint.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>

Author: justinchuby

src/mobius/_weight_utils.py

@@ -303,16 +303,18 @@ def rename_weight_keys(
             (a collision that would otherwise silently drop a tensor).
     """
     result: dict[str, torch.Tensor] = {}
+    producers: dict[str, str] = {}
     for name, tensor in state_dict.items():
         new_name = name
         for old, new in replacements:
             new_name = new_name.replace(old, new)
         if new_name in result:
             raise ValueError(
                 f"Weight key collision after rename: {name!r} -> {new_name!r} "
-                f"(already produced by another key)"
+                f"(already produced by {producers[new_name]!r})"
             )
         result[new_name] = tensor
+        producers[new_name] = name
     return result
 
 

src/mobius/_weight_utils_test.py

@@ -213,13 +213,22 @@ def test_shares_tensor_values(self):
         assert result["new.k"].data_ptr() == t.data_ptr()
 
     def test_collision_raises(self):
-        """Two source keys mapping to the same renamed key raises."""
+        """Two source keys mapping to the same renamed key raises.
+
+        The error message must name both the colliding key and the original
+        producer key to aid debugging in large checkpoints.
+        """
         state_dict = {
             "a.weight": torch.tensor(1.0),
             "b.weight": torch.tensor(2.0),
         }
-        with pytest.raises(ValueError, match="collision"):
+        with pytest.raises(ValueError, match="collision") as exc_info:
             rename_weight_keys(state_dict, [("a.", "x."), ("b.", "x.")])
+        message = str(exc_info.value)
+        # The producer ("a.weight", processed first) and the colliding key
+        # ("b.weight") must both appear.
+        assert "a.weight" in message
+        assert "b.weight" in message
 
     def test_empty_state_dict(self):
         """Empty input returns empty output."""