AutoModelForCausalLM.from_pretrained is killed by Loading checkpoint shards:
Author: 50516017
Pull request: no

### System Info

`- `transformers` version: 4.34.0.dev0
- Platform: Linux-5.15.90.1-microsoft-standard-WSL2-x86_64-with-glibc2.35
- Python version: 3.10.12
- Huggingface_hub version: 0.16.4
- Safetensors version: 0.3.3
- Accelerate version: 0.20.3
- Accelerate config:    not found
- PyTorch version (GPU?): 2.0.1+cu117 (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?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>`

### Who can help?

@younesbelkada,@ArthurZucker

### 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

I would like to fine tune AIBunCho/japanese-novel-gpt-j-6b using QLora. When I executed AutoModelForCausalLM.from_pretrained, it was killed by the python function and execution stopped. I was looking at the task manager and found that it was caused by CPU usage, but is it possible to load pretrained on the GPU? I have been able to fine tune other smaller models with Lora without any problems.
code is 
````
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
model_name = "AIBunCho/japanese-novel-gpt-j-6b"
config = AutoConfig.from_pretrained(model_name,use_fast=False)
bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_use_double_quant=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16
    )
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    config=config,
    device_map="cuda",
    load_in_8bit=True,
    #quantization_config=bnb_config
    
)
`

### Expected behavior

```
 Found duplicate ['libcudart.so', 'libcudart.so.11.0', 'libcudart.so.12.0'] files: {PosixPath('/usr/local/cuda/lib64/libcudart.so'), PosixPath('/usr/local/cuda/lib64/libcudart.so.11.0')}.. We'll flip a coin and try one of these, in order to fail forward.
Either way, this might cause trouble in the future:
If you get `CUDA error: invalid device function` errors, the above might be the cause and the solution is to make sure only one ['libcudart.so', 'libcudart.so.11.0', 'libcudart.so.12.0'] in the paths that we search based on your env.
  warn(msg)
CUDA exception! Error code: no CUDA-capable device is detected
CUDA exception! Error code: initialization error
CUDA SETUP: CUDA runtime path found: /usr/local/cuda/lib64/libcudart.so
/home/shimizu/create_LLM/macin/work2/venv/lib/python3.10/site-packages/bitsandbytes/cuda_setup/main.py:149: UserWarning: WARNING: No GPU detected! Check your CUDA paths. Proceeding to load CPU-only library...
  warn(msg)
CUDA SETUP: Detected CUDA version 118
CUDA SETUP: Loading binary /home/shimizu/create_LLM/macin/work2/venv/lib/python3.10/site-packages/bitsandbytes/libbitsandbytes_cpu.so...
None
Loading checkpoint shards:   0%|                                                                  | 0/2 [00:00<?, ?it/s]Killed
```

Comments:
Having similar lssue

Hi @50516017 @jmanhype 
This is because the first shard of the model is ~10GB, what I usually do is to push a smaller sharded version on the Hub.
I have just pushed a version with smaller shards under: `ybelkada/japanese-novel-gpt-j-6b-sharded`: https://huggingface.co/ybelkada/japanese-novel-gpt-j-6b-sharded - can you try that out instead?
In the future, to create these repos you can run:

```python
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "AIBunCho/japanese-novel-gpt-j-6b"
target_model_id = "yourusername/japanese-novel-gpt-j-6b-sharded"

model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True)
tokenizer = AutoTokenizer.from_pretrained(model_id)

model.push_to_hub(target_model_id, max_shard_size="2GB")
tokenizer.push_to_hub(target_model_id)
```

When this happened to me, I could fix it by just adding the arguments to load in half precision (`torch_dtype=torch.bfloat16`) and to initialize without duplicating the parameters in memory (`low_cpu_mem_usage=True`).

```
### The model that you want to train from the Hugging Face hub
model_name = "mistralai/Mixtral-8x7B-Instruct-v0.1"

###The instruction dataset to use
dataset_name = "vwxyzjn/openhermes-dev__mistralai_Mixtral-8x7B-Instruct-v0.1__1706886961"

### Fine-tuned model name
new_model = "Mixtral-8x7B-Instruct-finetune"

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=bnb_config,
    device_map=device_map
)
```
i am using this mixtral instruct model and i am getting a KeyError: 'mixtral'

@gautam-17-2003 can you try `pip install -U transformers`

@younesbelkada thanks it's working now


🤗 

It's true. Set `torch_dtype=torch.bfloat16` and `low_cpu_mem_usage=True` together.

I'm trying to load a 30B (65GB) model onto an 80GB A100, and it hits OOM only halfway through. I only need to run inference, and this PyTorch functionality seems to work for that:
```python
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH, weights_only=True))
```
But I can't find anything in a HF documentation that supports loading weights only and not gradients or all the other stuff that comes with a typical model. Am I missing something?

Few things, 1. the dtype you are using by default is gonna be `float32` 2. transformers does not load gradients either so this should not be the problem!