Repository: huggingface/accelerate

Issue:
`dispatch_model` seems to untie weights (typically the embedding matrix and the vocab projection `nn.Linear`)
### System Info

```Shell
- `Accelerate` version: 0.15.0.dev0
- Platform: Linux-5.4.0-1098-gcp-x86_64-with-glibc2.17
- Python version: 3.8.13
- Numpy version: 1.23.3
- PyTorch version (GPU?): 1.12.1+cu102 (True)
```


### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] One of the scripts in the examples/ folder of Accelerate or an officially supported `no_trainer` script in the `examples` folder of the `transformers` repo (such as `run_no_trainer_glue.py`)
- [X] My own task or dataset (give details below)

### Reproduction

The function `dispatch_model` seems to break the tie-ing of weights performed at initialization of a model:

```
from transformers import OPTForCausalLM
from accelerate import dispatch_model, infer_auto_device_map
from accelerate.utils import get_balanced_memory

# This model has the embedding matrix and the vocab projection tied
model = OPTForCausalLM.from_pretrained("facebook/opt-125m")

# This should start with `model.decoder.embed_tokens.weight cpu, model.decoder.embed_positions.weight cpu, ...`
# and should end with `model.decoder.layers.11.final_layer_norm.bias cpu`
for n, p in model.named_parameters():
    print(n, p.device)

# Dispatching - for simplicity, i am using a single GPU V100 16gb
max_memory = get_balanced_memory(model)
device_map = infer_auto_device_map(
    model, max_memory=max_memory
)
model = dispatch_model(model, device_map=device_map)

# This starts with 
# model.decoder.embed_tokens.weight cuda:0
# model.decoder.embed_positions.weight cuda:0
# but ends with an additional guest: `lm_head.weight cuda:0`
for n, p in model.named_parameters():
    print(n, p.device)

# tie-ing the weights again solve it
model.tie_weights()
for n, p in model.named_parameters():
    print(n, p.device)
```

### Expected behavior

No `lm_head.weight` in the call to `named_parameters()`



Discussion:
cc @SaulLu @leot13 

Resolving patch:
diff --git a/src/accelerate/big_modeling.py b/src/accelerate/big_modeling.py
--- a/src/accelerate/big_modeling.py
+++ b/src/accelerate/big_modeling.py
@@ -24,10 +24,12 @@
     OffloadedWeightsLoader,
     check_device_map,
     extract_submodules_state_dict,
+    find_tied_parameters,
     get_balanced_memory,
     infer_auto_device_map,
     load_checkpoint_in_model,
     offload_state_dict,
+    retie_parameters,
 )
 from .utils.versions import is_torch_version
 
@@ -312,6 +314,7 @@ def dispatch_model(
     else:
         weights_map = None
 
+    tied_params = find_tied_parameters(model)
     attach_align_device_hook_on_blocks(
         model,
         execution_device=execution_device,
@@ -320,6 +323,8 @@ def dispatch_model(
         weights_map=weights_map,
         preload_module_classes=preload_module_classes,
     )
+    # Attaching the hook may break tied weights, so we retie them
+    retie_parameters(model, tied_params)
     model.hf_device_map = device_map
     return model
 
diff --git a/src/accelerate/utils/__init__.py b/src/accelerate/utils/__init__.py
--- a/src/accelerate/utils/__init__.py
+++ b/src/accelerate/utils/__init__.py
@@ -55,6 +55,7 @@
     load_checkpoint_in_model,
     load_offloaded_weights,
     named_module_tensors,
+    retie_parameters,
     set_module_tensor_to_device,
 )
 from .offload import (
diff --git a/src/accelerate/utils/modeling.py b/src/accelerate/utils/modeling.py
--- a/src/accelerate/utils/modeling.py
+++ b/src/accelerate/utils/modeling.py
@@ -223,6 +223,25 @@ def find_tied_parameters(model: nn.Module, **kwargs):
     return result
 
 
+def retie_parameters(model, tied_params):
+    """
+    Reties tied parameters in a given model if the link was broken (for instance when adding hooks).
+
+    Args:
+        model (`torch.nn.Module`): The model in which to retie parameters.
+        tied_params (`Dict[str, str]`):
+            A mapping parameter name to tied parameter name as obtained by `find_tied_parameters`.
+    """
+    for param_name, tied_param_name in tied_params.items():
+        param = model
+        for split in param_name.split("."):
+            param = getattr(param, split)
+        tied_module = model
+        for split in tied_param_name.split(".")[:-1]:
+            tied_module = getattr(tied_module, split)
+        setattr(tied_module, tied_param_name.split(".")[-1], param)
+
+
 def compute_module_sizes(model: nn.Module, dtype: Optional[Union[str, torch.device]] = None):
     """
     Compute the size of each submodule of a given model.


Test patch:
diff --git a/tests/test_big_modeling.py b/tests/test_big_modeling.py
--- a/tests/test_big_modeling.py
+++ b/tests/test_big_modeling.py
@@ -44,6 +44,17 @@ def forward(self, x):
         return self.linear2(self.batchnorm(self.linear1(x)))
 
 
+class ModelForTestTiedWeights(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.linear1 = nn.Linear(4, 4)
+        self.batchnorm = nn.BatchNorm1d(4)
+        self.linear2 = nn.Linear(4, 4)
+
+    def forward(self, x):
+        return self.linear2(self.batchnorm(self.linear1(x)))
+
+
 class BiggerModelForTest(nn.Module):
     def __init__(self):
         super().__init__()
@@ -273,6 +284,15 @@ def test_dispatch_model(self):
             output = model(x)
             self.assertTrue(torch.allclose(expected, output.cpu(), atol=1e-5))
 
+    @require_cuda
+    def test_dispatch_model_tied_weights(self):
+        model = ModelForTestTiedWeights()
+        model.linear1.weight = model.linear2.weight
+        device_map = {"linear1": 0, "batchnorm": 0, "linear2": 0}
+
+        dispatch_model(model, device_map)
+        self.assertIs(model.linear2.weight, model.linear1.weight)
+
     @require_multi_gpu
     def test_dispatch_model_multi_gpu(self):
         model = BiggerModelForTest()