GPTJForCausalLM with instruction provided on tutorial doesn't load on 4090
Author: keyuchen21
Pull request: no

### System Info

GPU 4090 

- `transformers` version: 4.30.2
- Platform: Linux-5.10.16.3-microsoft-standard-WSL2-x86_64-with-glibc2.35
- Python version: 3.9.16
- Huggingface_hub version: 0.15.1
- Safetensors version: 0.3.1
- PyTorch version (GPU?): GPU 2.0.1+cu117 (True)
- Using GPU in script?: <4090>
- Using distributed or parallel set-up in script?: <no>







### Who can help?

@ArthurZucker @younesbelkada 

### Information

- [X] The official example scripts
- [ ] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction


![image](https://github.com/huggingface/transformers/assets/54015474/74568eac-b7c5-40c3-ba4e-b118758d4761)
![image](https://github.com/huggingface/transformers/assets/54015474/530a6f53-a379-4acb-a6c9-72bf1381be30)
![image](https://github.com/huggingface/transformers/assets/54015474/f8aa219f-00e9-47b6-ad38-35da7ba8c7b5)

### Expected behavior

from the official tutorial, it said it should work with 16 GB of Ram, but 4090 has 24G of ram.
the kernel is died.

Can anyone help me out?

https://huggingface.co/docs/transformers/model_doc/gptj

error:
Kernel Restarting
The kernel for gpt-j-6b/1.ipynb appears to have died. It will restart automatically.

Comments:
Hi @km5ar 
Thanks for the issue

This is because google colab instances have a relatively low CPU RAM (<24GB) and the GPT-J model stored on the Hub at `EleutherAI/gpt-j-6b` are actually in float32 (24GB). Therefore from_pretrained will try to download that large file and crashed. Moreover, for large models it is recommended to load a model either with `low_cpu_mem_usage=True` or `device_map="auto"` as by default from_pretrained will initialize a random model with the same number of paramaters then try to populate the model with the dowloaded weights. Using `low_cpu_mem_usage=True` will avoid that step and not create a dummy random model at the beginning of the from_pretrained call.

To run gpt-j 6B on google colab consider using this repo: [`ybelkada/gpt-j-6b-sharded-bf16`](https://huggingface.co/ybelkada/gpt-j-6b-sharded-bf16) if you want to load the model in bf16 (by passing `torch_dtype=torch.bfloat16` ) or this repo: [`philschmid/gpt-j-6B-fp16-sharded`](https://huggingface.co/philschmid/gpt-j-6B-fp16-sharded) if you want to run the model in `float16` (by passing `torch_dtype=torch.float16`).

@younesbelkada 
Thanks for the answer

However, my question is 
if you see code, I did use torch.float16.
and please see the following screenshot,
I was copy the code from the tutorial from official guide, which clearly said "The model should fit on 16GB GPU for inference."

I was using a 4090 in my local machine.
![image](https://github.com/huggingface/transformers/assets/54015474/57cdf8b4-0281-4afa-a106-46852d880fbf)
![1687434430321](https://github.com/huggingface/transformers/assets/54015474/e88bd673-671c-4461-a9e0-3a6ef70ec560)



@km5ar 
Thanks ! 
The statement "The model should fit on 16GB GPU" is still true. As explained above, the culprit is that `low_cpu_mem_uage` is set by default to `False` therefore blows up the Google Colab's CPU memory that is relatively low for that model due to the reasons I have detailed. Also, loading a checkpoint that is sharded helps to not getting those errors as the shards are processed one by one and deleted afterwards. 

<img width="560" alt="Screenshot 2023-06-22 at 14 06 53" src="https://github.com/huggingface/transformers/assets/49240599/d919490b-e915-431b-8d64-08c78762adb9">

As you can see from the screen shot above, `low_cpu_mem_usage=False` combined with that checkpoint will force the program to allocate 12+12=24GB CPU memory before moving the weights on GPU, hence your error.

Can you try to call `from_pretrained` with `low_cpu_mem_usage=True`, and use [philschmid/gpt-j-6B-fp16-sharded](https://huggingface.co/philschmid/gpt-j-6B-fp16-sharded) instead of the original repository?

Thanks

Thank you!!! that's very helpful!