Loading LongT5 model with `device_map` causes some layers to not be loaded
Author: GiantTreeLP
Pull request: no

### System Info

- `transformers` version: 4.34.1
- Platform: Windows-10-10.0.22621-SP0
- Python version: 3.11.5
- Huggingface_hub version: 0.17.3
- Safetensors version: 0.3.3
- Accelerate version: 0.23.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.1.0 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: no

### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

Trying to load a [LongT5](https://huggingface.co/docs/transformers/model_doc/longt5) model using `AutoModelForSeq2SeqLM.from_pretrained` and passing `device_map` causes the model to be loaded with the following error message:

```
Some weights of LongT5ForConditionalGeneration were not initialized from the model checkpoint at ./ and are newly initialized: ['encoder.embed_tokens.weight', 'decoder.embed_tokens.weight']
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
```

Running inference then results in bogus output.

If the `device_map` parameter is not passed, the model is loaded correctly.

Code to load the model:
```python
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, AutoConfig

config = AutoConfig.from_pretrained("./")
tokenizer = AutoTokenizer.from_pretrained("./")
model = AutoModelForSeq2SeqLM.from_pretrained("./", config=config, device_map=0)

tokenized_input = tokenizer.encode(
    "Van der Valk ist eine britische Fernseh-Krimiserie aus den Jahren 1972 bis 1977 und wieder 1991/1992 mit <hl>Barry Foster<hl> in der Rolle des Piet van der Valk.",
    return_tensors="pt")
tokenized_input = tokenized_input.to(model.device)
generated_tokens = model.generate(tokenized_input, max_length=1024)
decoded_tokens = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
print(decoded_tokens)
```

The model has been fine-tuned on a QG dataset in the German language but that does not change the issue.

Output:

```
Some weights of LongT5ForConditionalGeneration were not initialized from the model checkpoint at ./ and are newly initialized: ['decoder.embed_tokens.weight', 'encoder.embed_tokens.weight']
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
G:\dev\anaconda3\envs\transformers-26989\Lib\site-packages\transformers\generation\utils.py:1421: UserWarning: You have modified the pretrained model configuration to control generation. This is a deprecated strategy to control generation and will be removed soon, in a future version. Please use and modify the model generation configuration (see https://huggingface.co/docs/transformers/generation_strategies#default-text-generation-configuration )
  warnings.warn(
G:\dev\anaconda3\envs\transformers-26989\Lib\site-packages\transformers\models\longt5\modeling_longt5.py:171: UserWarning: An output with one or more elements was resized since it had shape [42], which does not match the required output shape [1, 42]. This behavior is deprecated, and in a future PyTorch release outputs will not be resized unless they have zero elements. You can explicitly reuse an out tensor t by resizing it, inplace, to zero elements with t.resize_(0). (Triggered internally at C:\cb\pytorch_1000000000000\work\aten\src\ATen\native\Resize.cpp:35.)
  true_block_ends = torch.logical_and(block_ends, block_ids >= 0)
G:\dev\anaconda3\envs\transformers-26989\Lib\site-packages\transformers\models\longt5\modeling_longt5.py:147: UserWarning: An output with one or more elements was resized since it had shape [1, 1, 128, 1], which does not match the required output shape [1, 1, 128, 384]. This behavior is deprecated, and in a future PyTorch release outputs will not be resized unless they have zero elements. You can explicitly reuse an out tensor t by resizing it, inplace, to zero elements with t.resize_(0). (Triggered internally at C:\cb\pytorch_1000000000000\work\aten\src\ATen\native\Resize.cpp:35.)
  local_attention_mask = torch.logical_and(_blocked_attention_mask, _3blocked_attention_mask)
G:\dev\anaconda3\envs\transformers-26989\Lib\site-packages\transformers\modeling_utils.py:859: FutureWarning: The `device` argument is deprecated and will be removed in v5 of Transformers.
  warnings.warn(
['ssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssss.']
```

### Expected behavior

The model should be loaded to the device given the `device_map` parameter and all layers and weights should be loaded correctly.

Expected output:
```
G:\dev\anaconda3\envs\transformers-26989\Lib\site-packages\transformers\generation\utils.py:1421: UserWarning: You have modified the pretrained model configuration to control generation. This is a deprecated strategy to control generation and will be removed soon, in a future version. Please use and modify the model generation configuration (see https://huggingface.co/docs/transformers/generation_strategies#default-text-generation-configuration )
  warnings.warn(
G:\dev\anaconda3\envs\transformers-26989\Lib\site-packages\transformers\models\longt5\modeling_longt5.py:171: UserWarning: An output with one or more elements was resized since it had shape [42], which does not match the required output shape [1, 42]. This behavior is deprecated, and in a future PyTorch release outputs will not be resized unless they have zero elements. You can explicitly reuse an out tensor t by resizing it, inplace, to zero elements with t.resize_(0). (Triggered internally at C:\cb\pytorch_1000000000000\work\aten\src\ATen\native\Resize.cpp:35.)
  true_block_ends = torch.logical_and(block_ends, block_ids >= 0)
G:\dev\anaconda3\envs\transformers-26989\Lib\site-packages\transformers\models\longt5\modeling_longt5.py:147: UserWarning: An output with one or more elements was resized since it had shape [1, 1, 128, 1], which does not match the required output shape [1, 1, 128, 384]. This behavior is deprecated, and in a future PyTorch release outputs will not be resized unless they have zero elements. You can explicitly reuse an out tensor t by resizing it, inplace, to zero elements with t.resize_(0). (Triggered internally at C:\cb\pytorch_1000000000000\work\aten\src\ATen\native\Resize.cpp:35.)
  local_attention_mask = torch.logical_and(_blocked_attention_mask, _3blocked_attention_mask)
G:\dev\anaconda3\envs\transformers-26989\Lib\site-packages\transformers\modeling_utils.py:859: FutureWarning: The `device` argument is deprecated and will be removed in v5 of Transformers.
  warnings.warn(
['Wer ist der Producer von Van der Valk?']
```

Comments:
Hey! Before going further could you make sur to run this code on the latest release of transfomers? (4.34.1) this was probably fixed!

Hi,
I have updated my issue text. I have added the actual code to test the issue and the actual as well as expected output.

Thanks, @SunMarc let's try this maybe with a public model and see if we can reproduce this anyway

Sorry for the delay @GiantTreeLP, this happens because you are using safetensors files + `device_map`. This [PR](https://github.com/huggingface/transformers/pull/27204/files) makes sure that the weights will be tied if `self.config.tie_word_embeddings` is `True` (set to `True` by default when creating a config). However, there is a good chance that in your config, it is set to False if you finetuned a model similar to`Stancld/longt5-tglobal-large-16384-pubmed-3k_steps`. 
To solve this, you can do the following: 

```py
input_ids = tokenizer(LONG_ARTICLE, return_tensors="pt").input_ids.to("cuda")
config = AutoConfig.from_pretrained("Stancld/longt5-tglobal-large-16384-pubmed-3k_steps", tie_word_embeddings=True)
model = LongT5ForConditionalGeneration.from_pretrained("Stancld/longt5-tglobal-large-16384-pubmed-3k_steps", config=config, device_map=0).to("cuda")
```



Hi @SunMarc,

I have tried to use `tie_word_embeddings=True` with and without `device_map=0`.
Without setting a `device_map`, the model gets loaded but produces wrong results, even whilst not warning about initialized weights.

So your provided solution does not fix my issue.

For LongT5, the model contains a `shared.weight` layer/weight that is used by both, the encoder and the decoder embeddings.
Whilst debugging I found an inconsistency:
When trying to determine the tied parameters, without a `device_map`, the following code runs and returns a list of lists with an entry containing three elements (shared.weight, encoder.embed_tokens.weight and decoder.embed_tokens.weight):
https://github.com/huggingface/transformers/blob/514de24abfd4416aeba6a6455ad5920f57f3567d/src/transformers/modeling_utils.py#L3662-L3669

In contrast with `device_map` set, the `find_tied_parameters` function is called which does not return the same list but an empty list:
https://github.com/huggingface/transformers/blob/514de24abfd4416aeba6a6455ad5920f57f3567d/src/transformers/modeling_utils.py#L3672

This causes the warning to be displayed.
As far as I can tell, this is only regarding the symptom but not the actual cause of the untied weights.

The code that is supposed to move the missing weights back from the meta device to the cpu does not move back the shared weights, those get left on the meta device when initializing:
https://github.com/huggingface/transformers/blob/514de24abfd4416aeba6a6455ad5920f57f3567d/src/transformers/modeling_utils.py#L3693-L3722

The safetensors stored model does, however, contain all the necessary weights.

I will look further into why the weights don't get properly loaded/tied in the next days.

Hi @GiantTreeLP, were you able to fix your problem ? 

Hi @SunMarc, I haven't had much time in December and January to investigate this issue further.

As I am not using the `device_map` argument this issue doesn't affect any of my projects.
I haven't found anything else that is worth mentioning.

Due to other projects and interests, I probably won't have the time to investigate this issue further.
I also believe this issue is above my current skill level to solve.

I can close this issue, if necessary but it's not fixed.

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.