Fix the initialization of the cache when we have multi gpu
Author: SunMarc
Pull request: yes

# What does this PR do ?

Fixes https://github.com/huggingface/transformers/pull/33287#issuecomment-2328906406
This PR initializes the cache on the right device when we are in multi-gpu setup. Before this PR, we would move the tensors to the right device during `update()` which created a few issues with `export()`. This is also a cleaner solution in general. 

cc @gante I would love to have a quick feedback from you 
cc @ArthurZucker as you were also interested in

Tested with:

`RUN_SLOW=True CUDA_VISIBLE_DEVICES=0,1 pytest tests/generation/test_utils.py -k "test_generate_with_static_cache_multi_gpu"`

`RUN_SLOW=True CUDA_VISIBLE_DEVICES=0,1 pytest tests/generation/test_utils.py -k "test_init_static_cache_multi_gpu"`

`RUN_SLOW=1 pytest tests/utils/test_cache_utils.py -k test_static_cache_exportability`

and: 

```py
from transformers import LlamaTokenizer, LlamaForCausalLM
import torch

NUM_TOKENS_TO_GENERATE = 40
# Note on `EXPECTED_TEXT_COMPLETION`'s diff: the current value matches the original test if the original test
# was changed to have a cache of 53 tokens (as opposed to 4096), on Ampere GPUs.
EXPECTED_TEXT_COMPLETION = [
    "Simply put, the theory of relativity states that 1) the speed of light is constant in all inertial "
    "reference frames, and 2) the laws of physics are the same for all inertial reference frames.\nThe "
    "theory of relativ",
    "My favorite all time favorite condiment is ketchup. I love it on everything. I love it on my eggs, "
    "my fries, my chicken, my burgers, my hot dogs, my sandwiches, my salads, my p",
]

prompts = [
    "Simply put, the theory of relativity states that ",
    "My favorite all time favorite condiment is ketchup.",
]
tokenizer = LlamaTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf", pad_token="
", padding_side="right")
model = LlamaForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf", device_map="balanced", torch_dtype=torch.float16
)
inputs = tokenizer(prompts, return_tensors="pt", padding=True).to(model.device)
# Dynamic Cache
generated_ids = model.generate(**inputs, max_new_tokens=NUM_TOKENS_TO_GENERATE, do_sample=False)
dynamic_text = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
# self.assertEqual(EXPECTED_TEXT_COMPLETION, dynamic_text)
print(dynamic_text)

# Static Cache
generated_ids = model.generate(
    **inputs, max_new_tokens=NUM_TOKENS_TO_GENERATE, do_sample=False, cache_implementation="static"
)
static_text = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
# We get the same expected output ! 
print(static_text)

```

Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_33303). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

> make sure we throw an exception somewhere: if a user is using multi GPU, then a layer map has to be passed 

The only way to know if we are using multi-gpu is with the `device_map`. Since we are computing the `layer_device_map` using `device_map` in `get_layer_device_map`, that should be good enough no ? If we want to do the check inside the cache, we would have to pass the `device_map` arg. 

@SunMarc 

That's a good point... 👀 My question at the moment is the following: if a user decides to instantiate a cache manually, is using multi-gpu, and doesn't pass the new argument, how can we let the user know that they should have used the new argument?

Refactored a bit + added an integration test to check the device on the cache ! Note that we do the check after generation as we need to initialize the cache, meaning that we need the model. We already have tests for multi-gpu static cache so i don't think we need to add these. Let me know if there are other tests you would like to see. 

`RUN_SLOW=True CUDA_VISIBLE_DEVICES=0,1 pytest tests/generation/test_utils.py -k "test_generate_with_static_cache_multi_gpu"`



> That's a good point... 👀 My question at the moment is the following: if a user decides to instantiate a cache manually, is using multi-gpu, and doesn't pass the new argument, how can we let the user know that they should have used the new argument?

When passing the cache to the model in `generate`, is there a way to post initialize the cache again ? It seems to be the best angle to fix this issue. If one is using multi-gpu, we would have to initialize the model before creating the cache object if we want to pass a relevant args such as the device_map. However, I prefer not passing the `device_map` as we would have to recreate `layer_device_map` afterwards in the cache init + not friendly to the user. 


>  is there a way to post initialize the cache again ?

There is not atm, and I think we would hit a few issues if we do so 🤔  We would have to either:
1. check whether there is data to copy across devices (data checks = not compile-friendly)
2. never copy data when creating tensors in new devices (unexpected behavior from a user perspective)
3. always copy data (we would have to use `to` to copy data across devices, which is essentially the issue #33287 tries to solve)

________________________

What about a simple device check at update time? If cache tensor device != new cache data device then throw informative exception

> What about a simple device check at update time? If cache tensor device != new cache data device then throw informative exception

Sounds good  @gante ! I've updated the PR with the check ! Let me know if this is better ! I'll add quick test to check if the warning is raised correctly ! 

@SunMarc can you include `RUN_SLOW=1 pytest tests/utils/test_cache_utils.py -k test_static_cache_exportability` in the Test Plan to ensure the export test is fixed?

> @SunMarc can you include RUN_SLOW=1 pytest tests/utils/test_cache_utils.py -k test_static_cache_exportability in the Test Plan to ensure the export test is fixed?

Test passed ! 

@gante @LysandreJik can we prioritize to get this fix merged? I will need this one to unblock `ExecuTorch` integration. Thanks!

Thanks, for cache related PRs I recommend pinging @ArthurZucker for a review (pinging him now)

Let's remove it for now and merge

cc @SunMarc 

I ran a quick benchmark to see what the impact on generate with `NUM_TOKENS_TO_GENERATE = 2000`.
- without the check : ~145.5s
- with the check : ~145s
-> We have a very small overhead 

But yeah, let's remove that for the release and I can do a follow-up pr with either the same patch or use try except as joao suggested. Feel free to merge the PR for the release @ArthurZucker if you are fine with the modification!