[`Blip2`] Fix Blip-2 multi gpu
Author: younesbelkada
Pull request: yes

# What does this PR do?

This PR should fix all the issues related to BLIP-2 and multi-gpu

Before this PR, BLIP2 had incorrect `set_input_embeddings` and `get_input_embeddings` functions, leading to unexpected behaviours when using it with `accelerate` since it `accelerate` gets confused when creating a device map with incorrect tied weights. 

Do not merge before I figure out why this does not fix the behaviour with `blip2-flan-t5`. EDIT: should work properly now, there are some corner cases that users can encounter but if they follow strictly the guidelines presented in https://github.com/huggingface/blog/blob/main/accelerate-large-models.md , it should be fine.

Fixes: https://github.com/TimDettmers/bitsandbytes/issues/153 & https://github.com/huggingface/transformers/pull/21441#issuecomment-1435370577

cc @sgugger 

Comments:
_The documentation is not available anymore as the PR was closed or merged._

Thanks for the suggestions! Adapted the code accordingly
However, there seems to be something off with the current `accelerate` integration w `blip2` - here since we are calling `generate` of a module of `Blip2ForConditionalGeneration`, I think that we are facing an edge case that I am not sure whether the fix should be here or on `accelerate`

I think there are 2 challenging points

1- if a users uses a device_map such as `auto` or `balanced` sometimes the `language_model` attribute will not have any `_hf_hook` attribute (from what I have got only the modules that are on the device_map and the parent module will have an `_hf_hook`), leading to the output of `language_model` not being on the correct device (as the output device will be determined by `language_model.lm_head._hf_hook`). 

2- If a well-educated user passes a custom device_map such as:
```python
device_map = {
    "query_tokens": 0,
    "vision_model":0,
    "language_model": 1,
    "language_projection": 0,
    "qformer": 0,
}
```
the `language_model` attribute will indeed have a `_hf_hook`, however the output of the module will not be set to the correct device as per my understanding, `_hf_hook.io_same_device` is set to `True` only on the parent class. I proposed a hacky solution at e0104f0 but I am not happy with it. 

So I think that maybe a fix should be upstreamed on `accelerate` to enable some child modules behave as the parent module (i.e. have the same `_hf_hook` behaviour), or maybe on `generate` to allow output on a different device that the input, meaning that we add multiple `.to` in `sample`, `greedy`, etc.
I am not sure what is the best solution here and I am sure that there is something simpler we can try!
Here is a script to reproduce the initial bug:

```python
import torch
from transformers import Blip2ForConditionalGeneration, Blip2Processor
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
from PIL import Image
import requests


url = "https://huggingface.co/hf-internal-testing/blip-test-image/resolve/main/demo.jpg"
image = Image.open(requests.get(url, stream=True).raw)


model_t5 = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-small", device_map="balanced", torch_dtype=torch.float16)
tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-small")

processor = Blip2Processor.from_pretrained("Salesforce/blip2-flan-t5-xl")

device_map = {
    "query_tokens": 0,
    "vision_model":0,
    "language_model": 1,
    "language_projection": 0,
    "qformer": 0,
}

# can also try with custom device_map
model = Blip2ForConditionalGeneration.from_pretrained("Salesforce/blip2-flan-t5-xl", device_map="balanced", torch_dtype=torch.float16)

inputs = processor(images=image, return_tensors="pt").to(0, torch.float16)

print(model_t5._hf_hook.execution_device)
print(model._hf_hook.execution_device)

print(model.language_model._hf_hook.execution_device) # this will fail

predictions = model.generate(**inputs, do_sample=True)
generated_text = processor.decode(predictions[0], skip_special_tokens=True)
print(generated_text)
```

I can confirm the current implementation works is a users passes a `device_map` that has `language_model` , let me know if you see anything else that needs to be addressed!

Thank you very much @akkikiki for the very useful feedback! So just to summarize, we need:
1- A more explicit warning pointing to the links you have shared so that users can understand how to use correct `device_map`
2- The fix you proposed in https://github.com/huggingface/transformers/pull/21707#discussion_r1119275324 to make it work for some edge cases where the masks are spread across different devices in the case when n_gpus > 2 (I can only test on a enviornment where I have 2 GPUs for now)
Is that correct? 

> Thank you very much @akkikiki for the very useful feedback! So just to summarize, we need: 1- A more explicit warning pointing to the links you have shared so that users can understand how to use correct `device_map` 2- The fix you proposed in [#21707 (comment)](https://github.com/huggingface/transformers/pull/21707#discussion_r1119275324) to make it work for some edge cases where the masks are spread across different devices in the case when n_gpus > 2 (I can only test on a enviornment where I have 2 GPUs for now) Is that correct?

Exactly! But the first one is just a suggestion so feel free to discard it :)

Thanks a mile @akkikiki , may I ask you to run the latest changes that I made on your side to confirm everything works as expected? Then we can merge I think! 

Re-installed your latest branch and works perfectly fine! Thanks a lot @younesbelkada!!

Hi, I'm still getting errors with this weirdly enough (am on 4.29.0 which I believe should have this fix included). 

I'm running the code from the model tutorial copy pasted:
```python
import torch
import requests
from PIL import Image
from transformers import Blip2Processor, Blip2ForConditionalGeneration

processor = Blip2Processor.from_pretrained("Salesforce/blip2-flan-t5-xxl")

model = Blip2ForConditionalGeneration.from_pretrained("Salesforce/blip2-flan-t5-xxl", 
                                                      load_in_8bit=True, 
                                                      device_map="auto",
                                                      )

img_url = 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/demo.jpg' 
raw_image = Image.open(requests.get(img_url, stream=True).raw).convert('RGB')

question = "how many dogs are in the picture?"
inputs = processor(raw_image, question, return_tensors="pt").to("cuda", torch.float16)

out = model.generate(**inputs)
print(processor.decode(out[0], skip_special_tokens=True))
```
which gives
```
/proj/vondrick4/sachit/miniconda3/envs/pyt13/lib/python3.10/site-packages/torch/utils/_contextli │
│ b.py:115 in decorate_context                                                                     │
│                                                                                                  │
│   112 │   @functools.wraps(func)                                                                 │
│   113 │   def decorate_context(*args, **kwargs):                                                 │
│   114 │   │   with ctx_factory():                                                                │
│ ❱ 115 │   │   │   return func(*args, **kwargs)                                                   │
│   116 │                                                                                          │
│   117 │   return decorate_context                                                                │
│   118                                                                                            │
│                                                                                                  │
│ /proj/vondrick4/sachit/miniconda3/envs/pyt13/lib/python3.10/site-packages/transformers/models/bl │
│ ip_2/modeling_blip_2.py:1854 in generate                                                         │
│                                                                                                  │
│   1851 │   │   inputs_embeds = self.get_input_embeddings()(input_ids)                            │
│   1852 │   │   inputs_embeds = torch.cat([language_model_inputs, inputs_embeds.to(language_mode  │
│   1853 │   │                                                                                     │
│ ❱ 1854 │   │   outputs = self.language_model.generate(                                           │
│   1855 │   │   │   inputs_embeds=inputs_embeds,                                                  │
│   1856 │   │   │   attention_mask=attention_mask,                                                │
│   1857 │   │   │   **generate_kwargs,                                                            │
│                                                                                                  │
│ /proj/vondrick4/sachit/miniconda3/envs/pyt13/lib/python3.10/site-packages/torch/utils/_contextli │
│ b.py:115 in decorate_context                                                                     │
│                                                                                                  │
│   112 │   @functools.wraps(func)                                                                 │
│   113 │   def decorate_context(*args, **kwargs):                                                 │
│   114 │   │   with ctx_factory():                                                                │
│ ❱ 115 │   │   │   return func(*args, **kwargs)                                                   │
│   116 │                                                                                          │
│   117 │   return decorate_context                                                                │
│   118                                                                                            │
│                                                                                                  │
│ /proj/vondrick4/sachit/miniconda3/envs/pyt13/lib/python3.10/site-packages/transformers/generatio │
│ n/utils.py:1515 in generate                                                                      │
│                                                                                                  │
│   1512 │   │   │   │   )                                                                         │
│   1513 │   │   │                                                                                 │
│   1514 │   │   │   # 11. run greedy search                                                       │
│ ❱ 1515 │   │   │   return self.greedy_search(                                                    │
│   1516 │   │   │   │   input_ids,                                                                │
│   1517 │   │   │   │   logits_processor=logits_processor,                                        │
│   1518 │   │   │   │   stopping_criteria=stopping_criteria,                                      │
│                                                                                                  │
│ /proj/vondrick4/sachit/miniconda3/envs/pyt13/lib/python3.10/site-packages/transformers/generatio │
│ n/utils.py:2372 in greedy_search                                                                 │
│                                                                                                  │
│   2369 │   │   │   if eos_token_id is not None:                                                  │
│   2370 │   │   │   │   if pad_token_id is None:                                                  │
│   2371 │   │   │   │   │   raise ValueError("If `eos_token_id` is defined, make sure that `pad_  │
│ ❱ 2372 │   │   │   │   next_tokens = next_tokens * unfinished_sequences + pad_token_id * (1 - u  │
│   2373 │   │   │                                                                                 │
│   2374 │   │   │   # update generated ids, model inputs, and length for next step                │
│   2375 │   │   │   input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1)              │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:7 and cuda:0!
```
Any ideas?

@sachit-menon 
The issue is that `device_map=auto` (at least for the recent HF versions that I touched) places the language model head at the later in GPUs (in this case `cuda:7`), but the original inputs are placed in the first GPU i.e., `cuda:0`

Can you run 
```
from accelerate import init_empty_weights, infer_auto_device_map
from transformers import Blip2Processor, Blip2ForConditionalGeneration

with init_empty_weights():
    model = Blip2ForConditionalGeneration(config)
    device_map = infer_auto_device_map(model, no_split_module_classes=["T5Block"])
device_map['language_model.lm_head'] = device_map["language_model.decoder.embed_tokens"]  # to make the genearted tokens and input_ids to be on the same device
```
and use the created `device_map` and set it as follows
```
model = Blip2ForConditionalGeneration.from_pretrained("Salesforce/blip2-flan-t5-xxl", 
                                                      load_in_8bit=True, 
                                                      device_map=device_map,
                                                      )
```
instead of `device_map="auto"` when initializing the model?


Thanks for the quick response. Now I get in the last line of the first block: `KeyError: 'language_model.decoder.embed_tokens'` 

@sachit-menon  Ops, sorry. Try out 
```
max_memory={i: "10GiB" for i in range(8)}
config = Blip2Config.from_pretrained(model_id)
with init_empty_weights():
    model = Blip2ForConditionalGeneration(config)
    device_map = infer_auto_device_map(model, no_split_module_classes=["T5Block"], dtype=torch.float16, max_memory=max_memory)
device_map['language_model.lm_head'] = device_map["language_model.encoder.embed_tokens"]
```
or tweak `10GiB` to be adjusted to your GPU memory you have (the above worked with 16GB GPU).

is this supposed to work on t5-large?

```
config = AutoConfig.from_pretrained("t5-large")
with init_empty_weights():
      model = AutoModelForSeq2SeqLM.from_config(config)
                
device_map = infer_auto_device_map(model, no_split_module_classes=["T5Block"], max_memory={i:1 for i in range(4)})
print(device_map)
device_map['lm_head'] = device_map['encoder.embed_tokens']

self.model = AutoModelForSeq2SeqLM.from_pretrained(self.hparams.model_name_or_path, load_in_8bit=True, device_map='auto', cache_dir="model_cache2")
```

leads to this error in the training step:

```
  File "/home/ec2-user/anaconda3/envs/JupyterSystemEnv/lib/python3.10/site-packages/transformers/models/t5/modeling_t5.py", line 260, in forward
    return self.weight * hidden_states
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1!
```

@cassianlewis Yes, replace
`self.model = AutoModelForSeq2SeqLM.from_pretrained(self.hparams.model_name_or_path, load_in_8bit=True, device_map='auto', cache_dir="model_cache2")` 
with 
`self.model = AutoModelForSeq2SeqLM.from_pretrained(self.hparams.model_name_or_path, load_in_8bit=True, device_map=device_map cache_dir="model_cache2")`

Hey @akkikiki 
Sorry, that was a typo in my original comment - I already tried using `device_map = device_map`

I tried `device_map['lm_head'] = device_map['encoder.embed_tokens']` and separately`device_map['lm_head'] = device_map['decoder.embed_tokens']` as suggested in https://github.com/akkikiki/huggingface_examples/blob/main/examples/load_flan_ul2.py

None of this worked...



> device_map['lm_head'] = device_map['encoder.embed_tokens']

oh~i got the same error ,could you tell me how to fit it

> @sachit-menon Ops, sorry. Try out
> 
> ```
> max_memory={i: "10GiB" for i in range(8)}
> config = Blip2Config.from_pretrained(model_id)
> with init_empty_weights():
>     model = Blip2ForConditionalGeneration(config)
>     device_map = infer_auto_device_map(model, no_split_module_classes=["T5Block"], dtype=torch.float16, max_memory=max_memory)
> device_map['language_model.lm_head'] = device_map["language_model.encoder.embed_tokens"]
> ```
> 
> or tweak `10GiB` to be adjusted to your GPU memory you have (the above worked with 16GB GPU).

@akkikiki i tried this way ,but it did not work. Can you tell me what the model_id are,in

__ config = Blip2Config.from_pretrained(model_id)
the error is:
device_map['language_model.lm_head'] = device_map["language_model.encoder.embed_tokens"]
KeyError: 'language_model.encoder.embed_tokens'


I noticed the `infer_auto_device_map()` returns an empty dict. The only solution that worked was to manually set the device map keys:

```
if torch.cuda.device_count() > 1:
    device_map = {
        "query_tokens": 0,
        "vision_model":0,
        "language_model": 1,
        "language_projection": 0,
        "qformer": 0,
    }
else:
    device_map = "auto"
```