Repository: huggingface/transformers
Pull request: [Core] [Offloading] Enable saving offloaded models with multiple shared tensor groups

## Purpose ##
* Fixes https://github.com/huggingface/transformers/issues/39097
* Enables models with multiple groups of shared tensors to be saved
  * For example, the token_embeds and lm_head weights are one group of shared tensors
  * Right now, only one group can be shared at a time, as indicated by `tied_names = tied_params[0]`
  * This enablement is useful for registering additional modules with shared tensors, such as quip/spinquant-like transforms

## Changes ##
* Instead of building a separate pathway for calculating shared tensors when offloaded, instead update `ptrs` calculation logic to be capable of handling meta tensors
  * This is equivalent to the `find_tied_parameters` method, since in the meta device case, both check for equality of the python object itself
* Misc, update tied tensors runtime error to suggest using `_dynamic_tied_weights_keys`, which is also a valid way of saving tied tensors

## Testing ##
```python3
import torch
from transformers import AutoModelForCausalLM

device_map = {
    "model.embed_tokens": "cuda:0",
    "model.layers": "cpu",
    "model.norm": "cpu",
    "model.rotary_emb": "cpu",
    "lm_head": "cuda:0"
}
model = AutoModelForCausalLM.from_pretrained("nm-testing/llama2.c-stories15M", device_map=device_map)

transform_a = torch.nn.Linear(1, 1, bias=False)
transform_a._dynamic_tied_weights_keys = ["weight"]
transform_b = torch.nn.Linear(1, 1, bias=False)
transform_b._dynamic_tied_weights_keys = ["weight"]

model.model.layers[0].self_attn.q_proj.register_module("transform", transform_a)
model.model.layers[1].self_attn.q_proj.register_module("transform", transform_a)
model.model.layers[2].self_attn.q_proj.register_module("transform", transform_b)
model.model.layers[3].self_attn.q_proj.register_module("transform", transform_b)

model.save_pretrained("tmp")
```

* Added test fails on main but succeeds with these changes

Changed files:
filename: src/transformers/modeling_utils.py
patch: @@ -3835,27 +3835,23 @@ def save_pretrained(
             # We're going to remove aliases before saving
             ptrs = collections.defaultdict(list)
             for name, tensor in state_dict.items():
-                # Sometimes in the state_dict we have non-tensor objects.
-                # e.g. in bitsandbytes we have some `str` objects in the state_dict
-                if isinstance(tensor, torch.Tensor):
-                    ptrs[id_tensor_storage(tensor)].append(name)
-                else:
+                if not isinstance(tensor, torch.Tensor):
+                    # Sometimes in the state_dict we have non-tensor objects.
+                    # e.g. in bitsandbytes we have some `str` objects in the state_dict
                     # In the non-tensor case, fall back to the pointer of the object itself
                     ptrs[id(tensor)].append(name)
 
-            # These are all the pointers of shared tensors
-            if hasattr(self, "hf_device_map"):
-                # if the model has offloaded parameters, we must check using find_tied_parameters()
-                tied_params = find_tied_parameters(self)
-                if tied_params:
-                    tied_names = tied_params[0]
-                    shared_ptrs = {
-                        ptr: names for ptr, names in ptrs.items() if any(name in tied_names for name in names)
-                    }
+                elif tensor.device.type == "meta":
+                    # In offloaded cases, there may be meta tensors in the state_dict.
+                    # For these cases, key by the pointer of the original tensor object
+                    # (state_dict tensors are detached and therefore no longer shared)
+                    tensor = self.get_parameter(name)
+                    ptrs[id(tensor)].append(name)
+
                 else:
-                    shared_ptrs = {}
-            else:
-                shared_ptrs = {ptr: names for ptr, names in ptrs.items() if len(names) > 1}
+                    ptrs[id_tensor_storage(tensor)].append(name)
+
+            shared_ptrs = {ptr: names for ptr, names in ptrs.items() if len(names) > 1}
 
             # Recursively descend to find tied weight keys
             _tied_weights_keys = _get_tied_weight_keys(self)
@@ -3899,7 +3895,9 @@ def save_pretrained(
 
             if len(error_names) > 0:
                 raise RuntimeError(
-                    f"The weights trying to be saved contained shared tensors {error_names} that are mismatching the transformers base configuration. Try saving using `safe_serialization=False` or remove this tensor sharing.",
+                    f"The weights trying to be saved contained shared tensors {error_names} that are mismatching "
+                    "the transformers base configuration. Try saving using `safe_serialization=False`, setting the "
+                    "`_dynamic_tied_weights_keys` attribute for affected modules, or remove this tensor sharing.",
                 )
 
         # Shard the model if it is too big.

filename: tests/utils/test_modeling_utils.py
patch: @@ -1187,6 +1187,29 @@ def test_save_offloaded_model(self):
         torch.testing.assert_close(output, presaved_output, rtol=1e-4, atol=1e-4)
         torch.testing.assert_close(presaved_output, postsaved_output)
 
+    @require_accelerate
+    @mark.accelerate_tests
+    @require_torch_accelerator
+    def test_save_offloaded_model_dynamic_tied_weights_keys(self):
+        from accelerate import dispatch_model
+
+        device_map = {"base": f"{torch_device}:0", "linear": "cpu", "linear2": "cpu"}
+        model = ModelWithHead(PretrainedConfig())
+        dispatch_model(model, device_map)
+
+        transform_a = torch.nn.Linear(1, 1, bias=False)
+        transform_a._dynamic_tied_weights_keys = ["weight"]
+        transform_b = torch.nn.Linear(1, 1, bias=False)
+        transform_b._dynamic_tied_weights_keys = ["weight"]
+
+        model.linear.register_module("transform_a", transform_a)
+        model.linear.register_module("transform_b", transform_b)
+        model.linear2.register_module("transform_a", transform_a)
+        model.linear2.register_module("transform_b", transform_b)
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            model.save_pretrained(tmp_dir)
+
     @require_safetensors
     def test_use_safetensors(self):
         # Should not raise anymore