Loading LLaMA hf format from local folder is not using GPU in Google Colab
Author: algiraldohe
Pull request: no

### System Info

Using Google colab:

NAME="Ubuntu"
VERSION="20.04.5 LTS (Focal Fossa)"

GPU = T4
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2022 NVIDIA Corporation
Built on Wed_Sep_21_10:33:58_PDT_2022
Cuda compilation tools, release 11.8, V11.8.89
Build cuda_11.8.r11.8/compiler.31833905_0

RAM 12.6 GB
GPU RAM 15GB
Disk 78.2 GB

Packages: (From running !pip list)
torch == 2.0.0+cu118
transformers == 4.29.2

### Who can help?

@ArthurZucker 
@younesbelka

### Information

- [x] 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 followed the steps provided here: https://huggingface.co/docs/transformers/main/en/model_doc/llama#overview to get an hf format of the LLaMA model that I can use. When I load the model from the output path in my local computer with CPU is working fine, VERY SLOW but fine so I moved to Google Colab in order to use GPU cause I need to fine tune the model after loading it, but when I monitor the resources while loading the model I can see that the GPU is not being used.


from transformers import LlamaForCausalLM, LlamaTokenizer
import torch

#Check if GPU is available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("Device:", device)

Output:
> Device: cuda

base_model_path = "/content/drive/MyDrive/Qbot-gpt/LLaMA-HF"

model = LlamaForCausalLM.from_pretrained(base_model_path ).to(device)
print_gpu_utilization()

tokenizer = LlamaTokenizer.from_pretrained(base_model_path)


The above code crashes in colab cause is not using the GPU but only the RAM.

This is the content of the folder LLaMA-HF:
![image](https://github.com/huggingface/transformers/assets/55255205/d99a8956-0f58-4f10-855c-0e41715892b0)


The function being used to print the GPU usage:
from pynvml import *


def print_gpu_utilization():
    nvmlInit()
    handle = nvmlDeviceGetHandleByIndex(0)
    info = nvmlDeviceGetMemoryInfo(handle)
    print(f"GPU memory occupied: {info.used//1024**2} MB.")


def print_summary(result):
    print(f"Time: {result.metrics['train_runtime']:.2f}")
    print(f"Samples/second: {result.metrics['train_samples_per_second']:.2f}")
    print_gpu_utilization()


### Expected behavior

It should output a loading bar with all the shards loaded:

Loading checkpoint shards: 100%.  ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||| 2/2 [00:32<00:00, 14.09s/it]


And more importantly it should print the amount of GPU that is being used when loading the model:
"GPU memory occupied: 1343 MB." 

Comments:
Hi @algiraldohe, thanks for raising this issue. 

To load the model directly onto the available GPUs, you should pass `device_map='auto'` when loading the model:

```python
model = LlamaForCausalLM.from_pretrained(base_model_path, device_map='auto')
```

Thank you for your prompt response I did as you mentioned and I received the following error:

╭─────────────────────────────── Traceback (most recent call last) ────────────────────────────────╮
│ in <cell line: 5>:5                                                                              │
│                                                                                                  │
│ /usr/local/lib/python3.10/dist-packages/transformers/modeling_utils.py:2777 in from_pretrained   │
│                                                                                                  │
│   2774 │   │   │   │   mismatched_keys,                                                          │
│   2775 │   │   │   │   offload_index,                                                            │
│   2776 │   │   │   │   error_msgs,                                                               │
│ ❱ 2777 │   │   │   ) = cls._load_pretrained_model(                                               │
│   2778 │   │   │   │   model,                                                                    │
│   2779 │   │   │   │   state_dict,                                                               │
│   2780 │   │   │   │   loaded_state_dict_keys,  # XXX: rename?                                   │
│                                                                                                  │
│ /usr/local/lib/python3.10/dist-packages/transformers/modeling_utils.py:2871 in                   │
│ _load_pretrained_model                                                                           │
│                                                                                                  │
│   2868 │   │   │   )                                                                             │
│   2869 │   │   │   is_safetensors = archive_file.endswith(".safetensors")                        │
│   2870 │   │   │   if offload_folder is None and not is_safetensors:                             │
│ ❱ 2871 │   │   │   │   raise ValueError(                                                         │
│   2872 │   │   │   │   │   "The current `device_map` had weights offloaded to the disk. Please   │
│   2873 │   │   │   │   │   " for them. Alternatively, make sure you have `safetensors` installe  │
│   2874 │   │   │   │   │   " offers the weights in this format."                                 │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
ValueError: The current `device_map` had weights offloaded to the disk. Please provide an `offload_folder` for 
them. Alternatively, make sure you have `safetensors` installed if the model you are using offers the weights in 
this format.

Not quite sure what is requesting for. I tried to set the offload_folder = "/content/drive/MyDrive/Qbot-gpt/LLaMA-HF" (Same path of the hf model) like this:

model = LlamaForCausalLM.from_pretrained(base_model_path, device_map='auto' , offload_folder=offload_folder)

But no luck yet, still no using the GPU.

@algiraldohe Passing the `device_map` argument is actually using the accelerate library to smartly load the model weights  to maximise GPU usage. To understand more about how it works, there's a [great doc page here](https://huggingface.co/docs/accelerate/usage_guides/big_modeling) describing how to load large models. There's also [this blog](https://huggingface.co/blog/accelerate-large-models). 

Following these, specifying an offload folder should work. Could you try specifying a different folder than the one the model weights are stored in e.g. 

```python
model = LlamaForCausalLM.from_pretrained(base_model_path, device_map='auto', offload_folder="offload")
```

The GPU RAM is 15 GB, but for e.g. 7B Llama, the weights alone are ~13.5 GB, so offload might be necessary. Let us know if there's still an issue. 

Once the model can be loaded without error, then we can try and diagnose if the GPU is being properly utilized or not. 





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.