device_map='sequential' does not utilize gpu devices other than the first when running in 8bit and 4bit
Author: Daryl149
Pull request: no

### System Info

transformers==4.31.0
python==3.10.6
bitsandbytes==0.40.2
torch==2.0.1

Whenever I set the parameter `device_map='sequential'`, only the first gpu device is taken into account. For models that do not fit on the first gpu, the model returns a cuda OOM, as if only running on the first gpu, instead of spilling over to the second gpu.
This is contrary to the expected behaviour as described in https://huggingface.co/docs/accelerate/usage_guides/big_modeling. It has happened for as long as I have been using it (last 3 months), so it is not just for the most recent version of transformers or torch, or bitsandbytes.
I have both a cuda:0 and cuda:1 gpu correctly installed and recognized by torch and transformers. For example, when setting `device_map= 'auto'`, the model is split on both gpus equally, as expected.




### Who can help?

@ArthurZucker @younesbelkada 

### 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

example code:
```
from transformers import LlamaForCausalLM, LlamaTokenizer
tokenizer = LlamaTokenizer.from_pretrained("meta-llama/Llama-2-70b-chat-hf", use_safetensors=True)
model = LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-70b-chat-hf", device_map="sequential", load_in_8bit=True, use_safetensors=True) 

```


### Expected behavior

model correcly loaded by spilling over remaining layers to second gpu, instead of OOM on just first gpu.

Comments:
Hi @Daryl149 
Thanks for your issue, per the documentation it says:

> "sequential" will fit what it can on GPU 0, then move on GPU 1 and so forth (so won’t use the last GPUs if it doesn’t need to).

Maybe there are some corner cases when using 8bit quantization, hence the OOM error. What is the GPU hardware you are using ? (with gpu vram)
Indeed the canonical way to load a model and split it across all GPU evenly is to use device_map=auto

Could be 8 bit related, also happens in 4bit. Will have to try unquantized.
I have unequal memory on this particular setup, so `device_map=auto` is
suboptimal:
First gpu is an A6000 (non Ada), 48GB.
Second gpu is an RTX 3090, 24GB. With `auto`, it only uses 24GB on both
cards as expected (48GB combined) . Whereas the expected `sequential`
behavior would be perfect for this situation by filling up to the combined
72GB.

On Fri, Jul 21, 2023, 08:46 Younes Belkada ***@***.***> wrote:

> Hi @Daryl149 <https://github.com/Daryl149>
> Thanks for your issue, per the documentation it says:
>
> "sequential" will fit what it can on GPU 0, then move on GPU 1 and so
> forth (so won’t use the last GPUs if it doesn’t need to).
>
> Maybe there are some corner cases when using 8bit quantization, hence the
> OOM error. What is the GPU hardware you are using ? (with gpu vram)
> Indeed the canonical way to load a model and split it across all GPU
> evenly is to use device_map=auto
>
> —
> Reply to this email directly, view it on GitHub
> <https://github.com/huggingface/transformers/issues/24965#issuecomment-1645064989>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/ABTMWHBJYDKPTVEFXZCMLLDXRIQV3ANCNFSM6AAAAAA2SC443A>
> .
> You are receiving this because you were mentioned.Message ID:
> ***@***.***>
>


Update:
It is definitely caused by setting the flag `load_in_8bit=True` (also occurs for `load_in_4bit=True`). When loading the model as such:
```
from transformers import LlamaForCausalLM, LlamaTokenizer
tokenizer = LlamaTokenizer.from_pretrained("meta-llama/Llama-2-70b-chat-hf", use_safetensors=True)
model = LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-70b-chat-hf", device_map="sequential", use_safetensors=True) 

```
It shows the expected behaviour for `sequential`. Unfortunately, this means that: 
- when running in 4 and 8 bit it goes OOM because it only utilizes the first gpu.
- when running the model in full, it goes OOM because it is too big anyway for 2 GPUs.

Hmmm, I'll report this with bitsandbytes as well. I would not consider the 8bit and 4bit flags as corner cases anymore for transformers though.


Any updates on this issue?

cc @muellerzr, you might have an idea as to what is going on here. When porting llama2 I had also a similar bug, using `device_map = 'auto'` just produces random outputs. Reproducer is [here](https://github.com/ArthurZucker/transformers/blob/98c869f16df2f23a2fffed102864c901fd3b2702/tests/models/llama/test_modeling_llama.py#L440)

cc @SunMarc for big model inference.

I think the key is that 'auto' works but 'sequential' does not. This means it is definitely capable of spreading weights across multiple gpus in 8 and 4-bit, but the part that handles that for 'auto' is probably missing from the logic for 'sequential'.

I was able to reproduce and I will work on fixing this. In the meantime, please use `max_memory` if you have a particular setup.