Llama: fix batched generation
Author: gante
Pull request: yes

# What does this PR do?

Fixes batched inference on llama, after the static cache changes were added. For instance, `RUN_SLOW=1 py.test tests/test_cache_utils.py::CacheIntegrationTest::test_dynamic_cache_beam_search` now passes.

### What was wrong?
`position_ids` has shape `[bsz, seq_len]`. The line computing `freqs` was correct for batch size = 1, but incorrect for larger batch sizes: it was summing the values for the different batch members. Therefore, we need to create another dimension to prevent this sum from happening, which is what this PR does.

### Throughput impact of changes
None 🙌  [Measured on my end, RTX3090 + `TinyLlama/TinyLlama-1.1B-Chat-v1.0`]

Before this PR
![Screenshot 2024-02-19 at 13 10 54](https://github.com/huggingface/transformers/assets/12240844/5b622062-a01d-4408-b81d-6e492e8f74e7)

After this PR
![Screenshot 2024-02-19 at 13 43 29](https://github.com/huggingface/transformers/assets/12240844/2bdc9c25-fba7-43ae-affc-751467962b14)


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

I'll have to run the benchmark on the A100 to make sure everything is alright but otherwise should be good 

Alright, no significant slow downs so 🟢 but I can't do naive Dynamic generation with the same script as before: 
Probably because I gave ` position_ids = torch.arange(seq_length, device=device)` and they are not unsqueezed

```python3
  File "/home/arthur/transformers/../static-kv-cache/clean_bench.py", line 147, in <module>
    outputs = model(input_ids, past_key_values=past_key_values,position_ids=position_ids,cache_position=cache_position, return_dict=False, use_cache = True)
  File "/home/arthur/miniconda3/envs/py310/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1536, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/arthur/miniconda3/envs/py310/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1545, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/arthur/transformers/src/transformers/models/llama/modeling_llama.py", line 1155, in forward
    outputs = self.model(
  File "/home/arthur/miniconda3/envs/py310/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1536, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/arthur/miniconda3/envs/py310/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1545, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/arthur/transformers/src/transformers/models/llama/modeling_llama.py", line 995, in forward
    layer_outputs = decoder_layer(
  File "/home/arthur/miniconda3/envs/py310/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1536, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/arthur/miniconda3/envs/py310/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1545, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/arthur/transformers/src/transformers/models/llama/modeling_llama.py", line 721, in forward
    hidden_states, self_attn_weights, present_key_value = self.self_attn(
  File "/home/arthur/miniconda3/envs/py310/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1536, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/arthur/miniconda3/envs/py310/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1545, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/arthur/transformers/src/transformers/models/llama/modeling_llama.py", line 628, in forward
    cos, sin = self.rotary_emb(value_states, position_ids, seq_len=None)
  File "/home/arthur/miniconda3/envs/py310/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1536, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/arthur/miniconda3/envs/py310/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1545, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/arthur/transformers/src/transformers/models/llama/modeling_llama.py", line 107, in forward
    position_ids[:, None, :].float()
IndexError: too many indices for tensor of dimension 1
```

@ArthurZucker regarding the benchmark error: position ids should be a 2D tensor, just like the input ids :D  I also had to adapt it on my end

Alright if passing a 1d before was erroring out! 

@gante thanks a lot for this