GPTJ6-B python process seem to be stuck forever at ```from_pretrained``` method 
Author: kd303
Pull request: no

## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->

- `transformers` version: 4.14.1
- Platform: Linux-5.4.0-56-generic-x86_64-with-glibc2.10
- Python version: 3.8.8
- PyTorch version (GPU?): 1.10.0+cu113 (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, Nvidia A100, 3 * 40GB
- Using distributed or parallel set-up in script?: Yes

### Who can help
- GPT-Neo, GPT-J, CLIP: @patil-suraj
Models:
- GPT-J with DeepSpeed
- - Deepspeed: @stas00
Library:

Model hub:

- for issues with a model, report at https://discuss.huggingface.co/ and tag the model's creator.

Model I am using (Bert, XLNet ...): gpt-j-6B

The problem arises when using:
[* ] my own modified scripts: (give details below)
[gpt-neo-fine-tuning-example/gpt_j_deepspeed.py at main · dredwardhyde/gpt-neo-fine-tuning-example · GitHub](https://github.com/dredwardhyde/gpt-neo-fine-tuning-example/blob/main/gpt_j_deepspeed.py)

The tasks I am working on is:
 [ ] my own task or dataset: (give details below)

## To reproduce

Steps to reproduce the behavior:

1. Allocated 3 40GB A100 GPU, 48 CPU Cores, 128 GB RAM, running on Kubeflow notebook image, pip freeze environment is attached 
2.  Line at the dump is getting produced
 ```
model = AutoModelForCausalLM.from_pretrained("./data-gptj-vol/models/",local_files_only=True, 
 torch_dtype=torch.float32).cuda() 

```

3. ds_config_gpt_json
```
{
   "train_batch_size":21,
   "fp16":{
      "enabled":true,
      "min_loss_scale":1,
      "opt_level":"O3"
   },
   "zero_optimization":{
      "stage":3,
      "offload_param":{
         "device":"cpu"
      },
      "offload_optimizer":{
         "device":"cpu"
      },
      "allgather_partitions":true,
      "allgather_bucket_size":5e8,
      "contiguous_gradients":true
   },
   "optimizer":{
      "type":"AdamW",
      "params":{
         "lr":5e-05,
         "betas":[
            0.9,
            0.999
         ],
         "eps":1e-08
      }
   },
   "scheduler":{
      "type":"WarmupLR",
      "params":{
         "warmup_min_lr":0,
         "warmup_max_lr":5e-05,
         "warmup_num_steps":100
      }
   }
}

```

<!-- If you have code snippets, error messages, stack traces please provide them here as well.
     Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting
     Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.-->

```
import faulhandler
faulthandler.dump_traceback_later(20, repeat=True) 
```
**gives following trace** it seems it is stuck at loading the model
```
Thread 0x00007f2dbb2c4740 (most recent call first):
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/init.py", line 395 in kaiming_uniform_
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/linear.py", line 96 in reset_parameters
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/linear.py", line 90 in __init__
  File "/opt/conda/lib/python3.8/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 254 in wrapper
  File "/opt/conda/lib/python3.8/site-packages/transformers/models/gptj/modeling_gptj.py", line 97 in __init__
  File "/opt/conda/lib/python3.8/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 254 in wrapper
  File "/opt/conda/lib/python3.8/site-packages/transformers/models/gptj/modeling_gptj.py", line 266 in __init__
  File "/opt/conda/lib/python3.8/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 254 in wrapper
  File "/opt/conda/lib/python3.8/site-packages/transformers/models/gptj/modeling_gptj.py", line 450 in <listcomp>
  File "/opt/conda/lib/python3.8/site-packages/transformers/models/gptj/modeling_gptj.py", line 450 in __init__
  File "/opt/conda/lib/python3.8/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 254 in wrapper
  File "/opt/conda/lib/python3.8/site-packages/transformers/models/gptj/modeling_gptj.py", line 688 in __init__
  File "/opt/conda/lib/python3.8/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 254 in wrapper
  File "/opt/conda/lib/python3.8/site-packages/transformers/modeling_utils.py", line 1411 in from_pretrained
  File "/opt/conda/lib/python3.8/site-packages/transformers/models/auto/auto_factory.py", line 436 in from_pretrained
  File "gptj_train.py", line 19 in <module>

```

## Expected behavior

it should run :). However the processes seem to be stuck for minutes at ```from_pretrained``` method as evident from faulttraces. 

Is there anything obvious I am missing, please let us know if there are any clue:

It does not seem to proceed after few minutes as well [10-15 minutes]
CPU utilization goes up from top and GPU utilization does not seem to go up at all when the ```from_pretrained``` method is executing.


Comments:
Interesting. I have never seen this one before.

1. I don't suppose it fails when you try it alone as the hanging isn't inside deepspeed but pytorch according to your trace:

```
python -c "import AutoModelForCausalLM; AutoModelForCausalLM.from_pretrained("./data-gptj-vol/models/",local_files_only=True, torch_dtype=torch.float32)"
```

2. Why are you putting the model on `cuda()` btw? Deepspeed takes care of that for you automatically and in the correct way - we don't want the whole model on each gpu, which is what your current code does. Could you please try to remove `.cuda()` and try again? On the other hand the `.cuda` call is after `from_pretrained` so the issue is before that.

edit: I can see now that the code you linked to wasn't designed for distributed training, but you want 3 gpus. so no `.cuda()` please.

Not sure where you copied this code from https://github.com/dredwardhyde/gpt-neo-fine-tuning-example/blob/main/gpt_j_deepspeed.py this is definitely not the way to run a distributed application - the code is using a hack to run the code from a notebook or just `python` on a single gpu (not `torch.distributed`).

if you want distributed training remove all those env vars that are a hack to emulate `torch.distributed`:
```
os.environ['MASTER_ADDR'] = 'localhost'
os.environ['MASTER_PORT'] = '9994'
os.environ['RANK'] = "0"
os.environ['LOCAL_RANK'] = "0"
os.environ['WORLD_SIZE'] = "1"
os.environ["TOKENIZERS_PARALLELISM"] = "false"
```

and run with a proper command:

```
python -m torch.distributed.run --nproc_per_node=3 gpt_j_deepspeed.py 
```

3 gpus, right? adjust above if needed

Full docs are here: https://pytorch.org/docs/stable/elastic/run.html

----

But actually have you tried the original example first and validating it worked on your setup before applying your changes to it?

As I can't see your changes it's hard to tell if the issue is in your code or the environment or the original code.

And of course show us how you launch the script.


Hi @stas00 ,

Thanks for replying. removing ```cuda()``` resolves the original issue I reported

I also found another issue on how Deepspeed containers are built, I built a container on Intel machine and deployed on A100 AMD, it was giving issue. Just leaving it here, for anyone who may be struggling. It is really a documentation issue not necessarily a bug [issue : 1886#](https://github.com/microsoft/DeepSpeed/issues/1886)

Thanks for replying once again.


Great to hear that it solved your problem, @kd303