[Cache] Don't initialize the cache on `meta` device
Author: gante
Pull request: yes

# What does this PR do?

### Context:
1. On #35164, we started initializing the cache on `meta` device by default. This was done to fix multi-device inference with cache, to postpone device allocation;
2. On #35873, we removed the `torch.nn.Module` inheritance from `Cache`, to enable tracing (e.g. for `optimum`)

Sadly, these two changes conflict. On main, we see graph breaks, because we assign new tensors on `.update()` when the device is `meta`:
```py
from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig

model_id = "meta-llama/Llama-3.2-1B-Instruct"

device = "cuda"
model = AutoModelForCausalLM.from_pretrained(model_id).to(device)
tokenizer = AutoTokenizer.from_pretrained(model_id)

generation_config = GenerationConfig(use_cache=True, cache_implementation="static")

processed = tokenizer('Anything', return_tensors="pt")
input_ids = processed["input_ids"].to(device)
output = model.generate(input_ids, generation_config=generation_config)
```

### The fix

Somewhere between removing the `torch.nn.Module` inheritance and other fixes merged to `main`, reverting most of #35164 solves both the graph breaks and the issue that PR is meant to work. In other words, the following example works with this PR. 

Note that the last line was not working with #35164, but is working with this PR. It gets a graph break, but it is expected ([we can't compile multi-device operations](https://pytorch.org/docs/stable/torch.compiler_cudagraph_trees.html#reasons-for-skipping-cudagraph))

(run with 2 GPUs)
```py
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model_id = 'google/gemma-2-2b-it'
tokenizer = AutoTokenizer.from_pretrained(model_id)

device_map = {"model.embed_tokens": 0, "model.norm": 1, "model.rotary_emb": 1, "lm_head": 0}
num_hidden_layers = 26
for i in range(num_hidden_layers):
    device_map[f"model.layers.{i}"] = 0 if i < 13 else 1

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype="bfloat16",
    device_map=device_map,
)

inputs = tokenizer("Today is a beautiful day!", return_tensors='pt').to(0)
out = model(**inputs)

out_2 = model.generate(**inputs, max_new_tokens=2, cache_implementation="static")
```

### Benchmarks/Tests

✅ slow llama tests (vs main)
✅ slow cache tests (vs main)
✅ no slowdowns (vs main)

(main)
![Screenshot 2025-03-04 at 17 56 09](https://github.com/user-attachments/assets/1ffd96be-e0e9-4917-a054-465f69a98462)

(this PR)
![Screenshot 2025-03-04 at 17 47 50](https://github.com/user-attachments/assets/188acf8a-6104-4732-8f04-9a9960e89cd6)




Comments:
Hi 👋, thank you for opening this pull request! The pull request is converted to draft by default. When it is ready for review, please click the `Ready for review` button (at the bottom of the PR page).

Are there concrete plans on landing this PR and if so which release will that get into? #35164 has caused some troubles for our torch.compile backend, which currently prevents us from using v4.49.0 transformers

Btw, @gante this issue was originally fixing a case in Gemma2 like models when we init cache within model's forward. In multiGPU setting we don't do any `layer_device_map` so all cache layers get assigned to `self.device`

I didn't add test probably, but we need to make sure that case doesn't fail before merging, causing all forward call to model fail in multiGPU

@zucchini-nlp ready for review :)

I've added regression tests, including for the gemma 2 case you mentioned!

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

@zucchini-nlp 
We can't have both automatic devices and compilation: if we don't do explicit device allocation, we have to create and assign new non-`meta` tensors in the first `forward` (or move devices, which ultimately leads to the same problem). This breaks compilation (graph break), because the address of the cache tensors was tagged as static. 

Why must the address of the cache tensors be static? It's a requirement for high throughput when we have the cache as a input (as opposed to an internal parameter of the `nn.Module`, our first implementation of the static cache). You can read more about it [here](https://dev-discuss.pytorch.org/t/compiling-the-optimizer-with-pt2/1669) (read from `In order to mitigate this problem, the mark_static_address` until the end of that paragraph).

~To alleviate the difficulty of having to find the right devices, I've exposed the function to find the cache devices :)~
EDIT: I've added a `.to` operation instead in the classes that strictly need them.

@zucchini-nlp the multigpu test is not yet corrected, will tag you for review when it's ready :)

@zucchini-nlp ready now :)

The current solution simultaneously:
- is compatible with multi-gpu (even when initializing the cache inside the model, although this approach will never be compatible with torch.compile, so I've added a note there); Added tests to prevent regressions.
- has no cuda graph breaks; Added a test to prevent regressions
- is compatible with torch.export; Confirmed by `RUN_SLOW=1 py.test tests/utils/test_cache_utils.py -k test_static_cache_exportability`

> we just remove warning and no logging needed from us?

Yes, because no action is needed from the user 

(this specific case was about removing an argument that will no longer be removed. If it was a rename, we would need to deprecate the new name)

@ZolotukhinM I missed your comment :D `v4.50` should include this fix, but I don't know when will we release it