Repository: huggingface/accelerate

Issue:
[Big model inference] ValueError: weight is on the meta device, we need a `value` to put in on 3.
### Information
Trying to dispatch a large language model's weights on multiple GPUs for inference following the [official user guide](https://huggingface.co/docs/accelerate/v0.17.1/en/usage_guides/big_modeling#handling-big-models-for-inference).

Everything works fine when I follow the example referenced in the `accelerate.load_checkpoint_and_dispatch` function documentation:
```python
from accelerate import init_empty_weights, load_checkpoint_and_dispatch
from huggingface_hub import hf_hub_download
from transformers import AutoConfig, AutoModelForCausalLM

# Download the Weights
checkpoint = 'EleutherAI/gpt-j-6B'
weights_location = hf_hub_download(checkpoint, "pytorch_model.bin")

# Create a model and initialize it with empty weights
config = AutoConfig.from_pretrained(checkpoint)
with init_empty_weights():
    model = AutoModelForCausalLM.from_config(config)

# Load the checkpoint and dispatch it to the right devices
model = load_checkpoint_and_dispatch(
    model, weights_location, device_map='auto', no_split_module_classes=['GPTJBlock']
)
```
But, I get the error mentioned in the issue title when loading other LLMs, including `EleutherAI/gpt-neo-1.3B`.

### System information
- Python version: 3.9.16
- Cuda version: Build cuda_11.8.r11.8/compiler.31833905_0
- GPUs: 4 x NVIDIA GeForce RTX 3090
- PyTorch 1.13.1
- Transformers 4.26.1
- Accelerate 0.17.1

**Accelerate config**
```
compute_environment: LOCAL_MACHINE
deepspeed_config: {}
distributed_type: MULTI_GPU
downcast_bf16: 'no'
dynamo_config: {}
fsdp_config: {}
gpu_ids: 0,1,2,3
machine_rank: 0
main_training_function: main
megatron_lm_config: {}
mixed_precision: 'no'
num_machines: 1
num_processes: 4
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false
```

### Code and error
```python
from accelerate import init_empty_weights, load_checkpoint_and_dispatch
from huggingface_hub import hf_hub_download
from transformers import AutoConfig, AutoModelForCausalLM

checkpoint = 'EleutherAI/gpt-neo-1.3B'
weights_location = hf_hub_download(checkpoint, 'pytorch_model.bin')

print(weights_location)
config = AutoConfig.from_pretrained(checkpoint, output_hidden_states=True)
with init_empty_weights():
    lmodel = AutoModelForCausalLM.from_config(config)
lmodel = load_checkpoint_and_dispatch(
    lmodel, weights_location, device_map='auto', no_split_module_classes=['GPTNeoBlock']
)
```
The line `print(weights_location)` outputs `/.cache/huggingface/hub/models--EleutherAI--gpt-neo-1.3B/snapshots/0f35a20339a9e208bc061570981180cfb87c9572/pytorch_model.bin` which is correct according to the documentation.

The code is basically identical to the one I used to dispatch the `EleutherAI/gpt-j-6B` model, except I specified `no_split_module_classes=['GPTNeoBlock']`. 

I am getting the following error:
```
Traceback (most recent call last):
  File "src/main.py", line 58, in <module>
    main(args)
  File "src/main.py", line 17, in main
    run_probing_train(args=args)
  File "src/run_probing.py", line 225, in run_probing_train
    lmodel = get_lmodel(args)
  File "src/run_probing.py", line 68, in get_lmodel
    lmodel = load_checkpoint_and_dispatch(
  File "/u/weyssowm/Research/AstProbing/venv/lib/python3.9/site-packages/accelerate/big_modeling.py", line 490, in load_checkpoint_and_dispatch
    return dispatch_model(
  File "/venv/lib/python3.9/site-packages/accelerate/big_modeling.py", line 370, in dispatch_model
    attach_align_device_hook_on_blocks(
  File "/venv/lib/python3.9/site-packages/accelerate/hooks.py", line 495, in attach_align_device_hook_on_blocks
    attach_align_device_hook_on_blocks(
  File "/venv/lib/python3.9/site-packages/accelerate/hooks.py", line 471, in attach_align_device_hook_on_blocks
    add_hook_to_module(module, hook)
  File "venv/lib/python3.9/site-packages/accelerate/hooks.py", line 155, in add_hook_to_module
    module = hook.init_hook(module)
  File "venv/lib/python3.9/site-packages/accelerate/hooks.py", line 244, in init_hook
    set_module_tensor_to_device(module, name, self.execution_device)
  File "venv/lib/python3.9/site-packages/accelerate/utils/modeling.py", line 136, in set_module_tensor_to_device
    raise ValueError(f"{tensor_name} is on the meta device, we need a `value` to put in on {device}.")
ValueError: weight is on the meta device, we need a `value` to put in on 3.
```



Discussion:
I can reproduce indeed, thanks for raising the issue! The problem comes from the tied weights being untied during the loading, and the state dict for this model does not have a key for the tied weights (to avoid the checkpoint being too large). I'll push a fix later today or tomorrow.

Resolving patch:
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
@@ -784,6 +784,7 @@ def load_checkpoint_in_model(
         offload_buffers (`bool`, *optional*, defaults to `False):
             Whether or not to include the buffers in the weights offloaded to disk.
     """
+    tied_params = find_tied_parameters(model)
     if offload_folder is None and device_map is not None and "disk" in device_map.values():
         raise ValueError(
             "At least one of the model submodule will be offloaded to disk, please pass along an `offload_folder`."
@@ -871,3 +872,5 @@ def load_checkpoint_in_model(
     if offload_state_dict:
         load_offloaded_weights(model, state_dict_index, state_dict_folder)
         shutil.rmtree(state_dict_folder)
+
+    retie_parameters(model, tied_params)


Test patch: