from_pretrained()'s load() blocks forever in subprocess
Author: levon003
Pull request: no

## Environment info
- `transformers` version: 3.5.1
- Platform: Linux-5.4.58 x86_64-with-glibc2.10
- Python version: 3.8.3
- PyTorch version (GPU?): 1.6.0 (False)
- Tensorflow version (GPU?): not installed (NA)
- Using GPU in script?: No
- Using distributed or parallel set-up in script?: Yes


### Who can help
Anyone familiar with the from_pretrained() code path. Perhaps @sgugger? Thank you! 

## Information

[from_pretrained()'s load()](https://github.com/huggingface/transformers/blob/master/src/transformers/modeling_utils.py#L1004) blocks forever loading roberta-base, due specifically to the call to `nn.Module._load_from_state_dict` that would load the "embeddings.word_embeddings".  Occurs when loading the model in both the first process and a subprocess started via `multiprocessing`.

I observe the same behavior when loading via keyword vs loading local files cached via `save_pretrained`.

Model I am using (Bert, XLNet ...): roberta-base

The problem arises when using:
* [ ] the official example scripts: 
* [x] my own modified scripts: see sample script below.

## To reproduce

Steps to reproduce the behavior:

```python
import torch
import transformers
import multiprocessing as mp

def load_model_in_subprocess():
    print("Started subprocess.")
    model2 = transformers.RobertaModel.from_pretrained('roberta-base')
    print("Model loaded in subprocess.")

def main():
    model1 = transformers.RobertaModel.from_pretrained('roberta-base')
    print("Model loaded in main process.")

    p = mp.Process(target=load_model_in_subprocess, daemon=True)
    p.start()
    p.join()
    print("Main thread terminating.")
    
if __name__ == "__main__":
    main()
```

Output:
```
Model loaded in main process.
Started subprocess.
<never terminates>
```

## Expected behavior

Model loads and is functional in both main process and subprocess. 


Comments:
If you've identified the issue to be coming from `nn.Module._load_from_state_dict`, then I guess this is more of a PyTorch issue than a `transformers` one? Do you have an idea what might cause this hang with that method?

Well, I don't know enough about torch state_dict behavior to understand why `transformers` would be directly calling the underscored "internal use" method `_load_from_state_dict` in the first place, but it strikes me that transformers is making assumptions about the functioning of this internal method that may not hold in practice; I don't see anything obvious in `_load_from_state_dict` that would cause it to lock up under these (or any) conditions, but we may be violating a usage assumption (e.g. providing a bad pre-load hook).  

Oh, I see. Looking at the `torch.load_state_dict` however, it doesn't seem to be doing something very differently to what we do. Have you managed to load several models using `torch.load()` with the same multiprocessing approach you have used?

Well, a fair test would be to load the _same_ (roBERTa-base) model, but I'm not sure how to write the code to do that... that's why I'm using `transformers`! But it's easy to verify that there's no problem with multi-process loading of PyTorch models:

```python
import torch
import torch.nn as nn
import multiprocessing as mp

USE_STATE_DICT = True

class SimpleNet(nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.fc1 = nn.Linear(768, 1)
        
    def forward(self, x):
        x = self.fc1(x)
        return x


def save_model():
    model = SimpleNet()
    torch.save(model, './full_model.pt')
    torch.save(model.state_dict(), './model_state_dict.pt')

def load_model_in_subprocess():
    print("Started subprocess.")
    if USE_STATE_DICT:
        model = SimpleNet()
        model.load_state_dict(torch.load('./model_state_dict.pt'))
    else:
        model = torch.load('./full_model.pt')
    print(f"Model loaded in subprocess: {model}")

def main():
    save_model()
    print("Saved model.")
    
    if USE_STATE_DICT:
        model = SimpleNet()
        model.load_state_dict(torch.load('./model_state_dict.pt'))
    else:
        model = torch.load('./full_model.pt')
    print(f"Model loaded in main process: {model}")

    p = mp.Process(target=load_model_in_subprocess, daemon=True)
    p.start()
    p.join()
    print("Main thread terminating.")
    
if __name__ == "__main__":
    main()
```

This script terminates fine when loading from state dict or a pickled model file.

Adding some debug prints to `transformers` load in modeling_utils.py, I can confirm that it is the call to `nn.Module._load_from_state_dict` when:

```
prefix = roberta.embeddings.word_embeddings.
local_metadata = {'version': 1}
missing_keys = ['roberta.embeddings.position_ids']
unexpected_keys = []
error_msgs = []
strict = True
```

The keys in the state_dict are:
```
roberta.embeddings.word_embeddings.weight
roberta.embeddings.position_embeddings.weight
roberta.embeddings.token_type_embeddings.weight
roberta.embeddings.LayerNorm.weight
roberta.embeddings.LayerNorm.bias
<snipping all of the individual layer keys e.g. roberta.encoder.layer.0.attention.self.query.weight>
roberta.pooler.dense.weight
roberta.pooler.dense.bias
lm_head.bias
lm_head.dense.weight
lm_head.dense.bias
lm_head.layer_norm.weight
lm_head.layer_norm.bias
lm_head.decoder.weight
```

The shape of the `roberta.embeddings.word_embeddings.weight` tensor is [50265,768].

(note: same blocking behavior when loading bert-base-uncased)



Okay, I did more investigation, and the problem is a blocking call to `Tensor.copy_` that copies the Parameter in the state_dict into the Parameter in the Module (in this case, the `Embedding(50265, 768, padding_idx=1)` parameter in the roBERTa model).  

The [documentation](https://pytorch.org/docs/stable/tensors.html#torch.Tensor.copy_) indicates a non_blocking parameter that can be used when copying between CPU and GPU, but we are copying between CPU and CPU. I confirmed that non_blocking does nothing, and that the `device` of both Parameters is `cpu`.  

That's where I'm going to stop pursuing this bug.  I don't know the structure of the C++ code, but it seems likely that this is an issue with the PyTorch CPU copy implementation and the idiosyncrasies of the specific OS I'm using.  If this problem can be reproduced on others systems it may be worth investigating further, but it does seem like the fault probably lies with `PyTorch` and not with `transformers`.  Hopefully this affects only a small set of OSes.

@LysandreJik, you may want to close this issue?

Thank you very much for your deep investigation of this issue. Unfortunately I don't see how we could change that on our front to make it work, so we'll close this for now.

If we get other reports of this we'll investigate further. 

Got the same issue in this environment : 

    Platform: Linux clem-MacBookAir 5.13.0-40-generic #45~20.04.1-Ubuntu x86_64 x86_64 x86_64 GNU/Linux
    Python version: 3.9.7
    PyTorch version (GPU?): 1.11.0+cu102 (False)
    Tensorflow version (GPU?): not installed (NA)
    Using GPU in script?: No
    Using distributed or parallel set-up in script?: Yes

Experiencing the same issue as well with a `torchvision.models` which seems to be coming from `nn.Module._load_from_state_dict` running as subprocess on CPU, unsure why this has just started to happen.

moving the model to the GPU before loading works as a workaround
```
model.to(device)
model.load_from_state_dict(ckpt)
```