🚨All attention refactor🚨
Author: ArthurZucker
Pull request: yes

# What does this PR do?
Todo in this PR:

- [ ] Cohere
- [ ] Chameleon
- [ ] DBRX
- [x] Gemma
- [x] Gemma2
- [x] GLM (modular donc rien à faire je crois)
- [ ] gpt_neoX et GPT2
- [x] Granite
- [ ] Jamba
- [ ] JetMoe
- [ ] Mimi
- [x] Mistral
- [x] Mixtral
- [ ] Mllama
- [ ] Moshi
- [ ] Nemotron
- [ ] OPT
- [x] Phi
- [ ] Ph3
- [ ] PhiMoe
- [x] Qwen2
- [ ] qwen2Moe
- [ ] qwen2VL
- [ ] SableML
- [x] StartCoder2 -> Modular normalement oK
- [ ] Idefics1,2,3
- [x] Olmo
- [x] Olmo2
- [ ] Siglip
- [ ] Whisper


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

Confirmed slow tests with Llama, everything is similar to main!

run-slow: vit

(just a check unrelated to this PR)

Thanks a lot for the feedback @vasqu! Please have a look at the fixes in https://github.com/huggingface/transformers/pull/35342!

@ArthurZucker @Cyrilvallez do you plan to refactor `modeling_phi3.py` using the `ALL_ATTENTION_FUNCTIONS` you recently introduced ? I see Phi 3 is in the list shared by @ArthurZucker at the beginning of this PR. It would be very helpful for our [kvpress](https://github.com/NVIDIA/kvpress) package. We plan to update it to be compatible with the future v4.48, any idea of when it will be released ? (december ? january ?)

Release will happen on 🎅🏻 🎁 !
Yeah for sure. Unless you submit a PR first 👀 not sure it will be in this release as we are all going on holidays but in january's release it will be included


I'm going to holidays too so I'll wait for January ^^ Happy holidays!

Happy holidays!

I noticed that there is still torch.reshape used in quite a few places, for instance in `modeling_gpt2.py`  Won't it be an obstacle to compiling these models? Why not replacing it with einops and using `einops._torch_specific.allow_ops_in_compiled_graph()`  ?

Hey! Indeed for gpt2 I used a 'reshape' instead of the usual 'view' because I hit an edge case that wasn't compatible with viewing at some point (but I will recheck that it still appears with latest developments, might have been an artifact during the debugging process). Whenever possible (most of the time), 'reshape' is actually equivalent to 'view' so no worries there anyway 😉

Einops has never been used in our modeling codes, and we avoid using it at all costs! `Llama` uses `reshape` and pretty sur it's compile compatible! 🤗 

I didn't follow the whole discussion, so I'm not sure if this is intentional or not: Some modules have lost the `hidden_size` attribute after the refactor, e.g. `LlamaAttention` had this attribute:

https://github.com/huggingface/transformers/blob/5d7739f15a6e50de416977fe2cc9cb516d67edda/src/transformers/models/llama/modeling_llama.py#L273

and now it's gone. I know it can be worked around relatively easily by using `module.config.hidden_size`, but this may break some code out there (e.g. in PEFT).

@ArthurZucker Not sure how to interpret your reaction, will the `hidden_size` attribute be re-added in transformers or should I adjust the PEFT code?

I noticed another issue caused by this PR when using multiple GPUs. Here is the reproducer:

```python
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "hf-internal-testing/tiny-random-MistralForCausalLM"
tokenizer = AutoTokenizer.from_pretrained(model_id, padding="left")
inputs = tokenizer(["A list of colors: red, blue"], return_tensors="pt").to("cuda")

device_map = {
    "model.embed_tokens": 0,
    "model.layers.0": 0,
    "model.layers.1": 1,
    "model.norm": 1,
    "lm_head": 1,
}
model = AutoModelForCausalLM.from_pretrained(model_id, device_map=device_map)
assert len({p.device for p in model.parameters()}) > 1
model.generate(**inputs)  # fails
```

The error is:

```
  File "/home/name/work/forks/transformers/src/transformers/models/mistral/modeling_mistral.py", line 325, in forward
    freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
             ~~~~~~~~~~~~~~~~~~~~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0! (when checking argument for argument mat2 in method wrapper_CUDA_bmm)
```

<details>

<summary>Full error message</summary>

```
Traceback (most recent call last):
  File "/home/name/anaconda3/envs/peft/lib/python3.11/pdb.py", line 1777, in main
    pdb._run(target)
  File "/home/name/anaconda3/envs/peft/lib/python3.11/pdb.py", line 1643, in _run
    self.run(target.code)
  File "/home/name/anaconda3/envs/peft/lib/python3.11/bdb.py", line 600, in run
    exec(cmd, globals, locals)
  File "<string>", line 1, in <module>
  File "/home/name/work/forks/peft/foo.py", line 19, in <module>
    model.generate(**inputs)
  File "/home/name/anaconda3/envs/peft/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/name/work/forks/transformers/src/transformers/generation/utils.py", line 2253, in generate
    result = self._sample(
             ^^^^^^^^^^^^^
  File "/home/name/work/forks/transformers/src/transformers/generation/utils.py", line 3252, in _sample
    outputs = self(**model_inputs, return_dict=True)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/name/anaconda3/envs/peft/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/name/anaconda3/envs/peft/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/name/work/forks/accelerate/src/accelerate/hooks.py", line 170, in new_forward
    output = module._old_forward(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/name/work/forks/transformers/src/transformers/models/mistral/modeling_mistral.py", line 838, in forward
    outputs = self.model(
              ^^^^^^^^^^^
  File "/home/name/anaconda3/envs/peft/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/name/anaconda3/envs/peft/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/name/work/forks/transformers/src/transformers/models/mistral/modeling_mistral.py", line 544, in forward
    position_embeddings = self.rotary_emb(hidden_states, position_ids)
                          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/name/anaconda3/envs/peft/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/name/anaconda3/envs/peft/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/name/anaconda3/envs/peft/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/name/work/forks/transformers/src/transformers/models/mistral/modeling_mistral.py", line 325, in forward
    freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
             ~~~~~~~~~~~~~~~~~~~~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
```

</details>

The issue stems from the fact that on `MistralRotaryEmbedding`, the `inv_freq` is still on CPU for some reason. When I checkout the previous commit (`75be5a0a5b1898ee86e5e0c1f7b58b77bb105101`), it is on CUDA as it should be. The `torch.autocast` context manager doesn't help here, since `enabled=False` (although setting it to `True` does not help).

Hey @BenjaminBossan! The error arises because Mistral no longer has the `rotary_emb` in the Layers, it's in the Model directly! Thus your custom map does not correctly set it. Switching your device_map to:
```py
device_map = {
    "model.embed_tokens": 0,
    "model.layers.0": 0,
    "model.layers.1": 1,
    "model.norm": 1,
    "model.rotary_emb": 1,
    "lm_head": 1,
}
```
fixes it

BTW, maybe it would be safer to check `{p.device for p in model.parameters()} | {p.device for p in model.buffers()}`?

Thanks a lot @Cyrilvallez I did not notice that. Indeed, your change fixes the failure.

> BTW, maybe it would be safer to check `{p.device for p in model.parameters()} | {p.device for p in model.buffers()}`?

Thanks for the suggestion. In the original unit test, this is just a sanity check, the proper testing comes further down, I just removed that part for clarity.

Yeah sorry @BenjaminBossan, it's just that we don't use that attribute so not sure we are gonna add it back! THe PR is breaking as the name indicates! But I hope it was not too much trouble!

Thanks for letting me know, it should be an easy fix on the PEFT side.

query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)

    raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
AttributeError: 'MistralAttention' object has no attribute 'num_heads'

How can I fix this?

Hey! you should try to use the latest release of `transformers`! `query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)` is what's used now. 

Is this by any chance related to `AWQ` or another package? 


Is there any doc about how to migrate from previous version to this version, like the variable definition, the alias change?

Have you tested on several benchmarks about the performance? I knew that the [Longbench](https://github.com/THUDM/LongBench) score on transformer v4.47 vs v4.36 varies a lot on llama-3. Is it stable on this version?
I suggest adding some simple and small dataset tests.

Hey! Everything stays the same in terms of user experience/benchmark scores. If you used to hack into the different Layer classes however, it may have changed a bit. You can simply go and check-out the modeling code in this case (as was the case if you hacked into it in the first place I guess!)

My friends use a `GPT2Model` in production and want to compile it with `StaticCache`. With the maintainers blessing, I would try to create a PR with `DynamicCache` / `StaticCache` support in `GPT2Model`. 
I am quite familiar with `Cache` class, I already coded some and made the `DynamicCache` work.

Please let me know if there are any hidden obstacles in Cache implementation for GPT2? Which tests to run or add? 
@ArthurZucker

cc @gante to that question!

I've chatted to @poedator offline -- I couldn't think of any obstacle in particular, and suggested a) to ensure we leave a deprecation warning regarding the old cache format b) use `RUN_SLOW=1 py.test tests/models/gpt2/test_modeling_gpt2.py` as a correctness check (gpt2 is fairly well tested, especially wrt text generation)

It looks like `test_flash_attn_2_from_config` is broken - it expects attention layer to have flashattention in its name, 
` if "FlashAttention" in module.__class__.__name__:...`
but after this refactoring, the attention classes are named differently. 

ref https://github.com/huggingface/transformers/blob/5fa35344755d8d9c29610b57d175efd03776ae9e/tests/test_modeling_common.py#L4641

please fix or suspend the test.
@ArthurZucker 

indeed gimme a min!