[`BC 4.37 -> 4.38`] for Llama family, memory and speed 
Author: ArthurZucker
Pull request: yes

# What does this PR do?
Fixes the BC issues between the two versions in term of memory consumption.
This fix is made a lot easier by all the tests, so thanks a lot @gante! 

fixes #29412, fixes #29484 , fixes #29644, fixes #29651

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

torch script can be fixed by @fxmarty in a follow up PR + patch IMO! 

@ArthurZucker that would be great if torchscript/fx tests pass &

`optimum-cli export onnx --model fxmarty/tiny-llama-fast-tokenizer llama_onnx`

does not break

@ArthurZucker the line `causal_mask *= torch.arange(target_length, device=device) > cache_position[0]` should become `causal_mask *= torch.arange(target_length, device=device) > cache_position.reshape(-1, 1)`

In the previous version, in a situation with 17 cached tokens and 4 new assistant tokens, the causal mask would become
```
(Pdb) causal_mask
tensor([[[[     0.,     -0.,     -0.,     -0.,     -0.,     -0.,     -0.,
               -0.,     -0.,     -0.,     -0.,     -0.,     -0.,     -0.,
               -0.,     -0.,     -0., -65504., -65504., -65504., -65504.],
          [     0.,      0.,     -0.,     -0.,     -0.,     -0.,     -0.,
               -0.,     -0.,     -0.,     -0.,     -0.,     -0.,     -0.,
               -0.,     -0.,     -0., -65504., -65504., -65504., -65504.],
          [     0.,      0.,      0.,     -0.,     -0.,     -0.,     -0.,
               -0.,     -0.,     -0.,     -0.,     -0.,     -0.,     -0.,
               -0.,     -0.,     -0., -65504., -65504., -65504., -65504.],
          [     0.,      0.,      0.,      0.,     -0.,     -0.,     -0.,
               -0.,     -0.,     -0.,     -0.,     -0.,     -0.,     -0.,
               -0.,     -0.,     -0., -65504., -65504., -65504., -65504.],
          [     0.,      0.,      0.,      0.,      0.,     -0.,     -0.,
               -0.,     -0.,     -0.,     -0.,     -0.,     -0.,     -0.,
               -0.,     -0.,     -0., -65504., -65504., -65504., -65504.]]]],
       device='cuda:0', dtype=torch.float16)
```

i.e. not upper triangular. After the suggested change, it becomes

```
(Pdb) causal_mask
tensor([[[[     0.,     -0.,     -0.,     -0.,     -0.,     -0.,     -0.,
               -0.,     -0.,     -0.,     -0.,     -0.,     -0.,     -0.,
               -0.,     -0.,     -0., -65504., -65504., -65504., -65504.],
          [     0.,      0.,     -0.,     -0.,     -0.,     -0.,     -0.,
               -0.,     -0.,     -0.,     -0.,     -0.,     -0.,     -0.,
               -0.,     -0.,     -0.,     -0., -65504., -65504., -65504.],
          [     0.,      0.,      0.,     -0.,     -0.,     -0.,     -0.,
               -0.,     -0.,     -0.,     -0.,     -0.,     -0.,     -0.,
               -0.,     -0.,     -0.,     -0.,     -0., -65504., -65504.],
          [     0.,      0.,      0.,      0.,     -0.,     -0.,     -0.,
               -0.,     -0.,     -0.,     -0.,     -0.,     -0.,     -0.,
               -0.,     -0.,     -0.,     -0.,     -0.,     -0., -65504.],
          [     0.,      0.,      0.,      0.,      0.,     -0.,     -0.,
               -0.,     -0.,     -0.,     -0.,     -0.,     -0.,     -0.,
               -0.,     -0.,     -0.,     -0.,     -0.,     -0.,     -0.]]]],
       device='cuda:0', dtype=torch.float16)
```

Any idea why our tests don't complain? 

> Any idea why our tests don't complain?

I think we don't have hard correctness checks for assisted generation 🙈 only API checks

congratulations with this PR fixing many important things!
Unfortunately, it broke ~~my heart~~ custom 4D masks support. 
try `RUN_SLOW=1 python -m pytest -v ./tests/test_modeling_utils.py::Mask4DTestHard::test_partial_stacked_causal_mask`

@gante and myself introduced this test recently in #29731 - please make it a part of the test suite for all things related to attention_masks and `StaticCache`

here is what happens (numbers based on `test_partial_stacked_causal_mask` after [this line 2170](https://github.com/gante/transformers/blob/a97a7e907978d1e08c45bdffbdab1b99358dc712/tests/test_modeling_utils.py#L2170):
in `modeling_llama.py::_update_causal_mask()`
the `causal_mask` has shape `(1, 1, sequence_length, target_length)` (1,1,9,12) (3 tokens in cache)
if the custom 4D mask enters with same shape, it triggers offset = 3 and then causes error when copied over causal mask.
if the custom 4D mask enters with shape (1, 1, 12, 12), offset stays at zero but then again this causes error when copied over causal mask, for mask_slice and causal_mask shapes don't match.
There are also changes [here](https://github.com/huggingface/transformers/blob/1abd0986ddde1b7235e691703c32cd8bda3ae069/src/transformers/models/llama/modeling_llama.py#L661) which may affect this test.

I hesitate to offer a PR because it may break other things that you try to do with this part of code. But, please, make the test work. BTW, it may be OK to change the test, for instance passing the whole bigger mask including cached items by editing [this line](https://github.com/gante/transformers/blob/a97a7e907978d1e08c45bdffbdab1b99358dc712/tests/test_modeling_utils.py#L2168) to `mask_1b = mask_1`

Special note on StaticCache: I like this feature and I want to use custom 4D masks with it. So far this is not tested. I'd be glad to contribute such test once this issue is fixed. It will look like `test_partial_stacked_causal_mask`, only with StaticCache.

cc @ArthurZucker @gante 

I can try to fix it, and yes I thought it would be tested automatically. It should be part of the `tokenization_common` or at least LlamaIntegrationTests or something. cc @gante if you can take a look I'll gladly review a PR

^ this PR should fix it 🤗