enable static cache on TP model
Author: jiqing-feng
Pull request: yes

The TP model usually needs torch.compile to get speed-up, but the static cache failed on running TP model. This PR will check the tp size and then allocate the correct cache shape.

Reproduce error:
`torchrun --standalone --nproc-per-node 2 tp_hf.py`
```python
import torch
import torch.distributed as dist
from transformers import AutoTokenizer, AutoModelForCausalLM


model_id = "meta-llama/Llama-3.1-8B-Instruct"
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, tp_plan="auto")
if dist.is_initialized():
    print("Backend:", dist.get_backend())

# Prepare input tokens
tokenizer = AutoTokenizer.from_pretrained(model_id)
prompt = "It is done, and submitted. You can play 'Survival of the Tastiest' on Android, and on the web."
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

with torch.no_grad():
    outputs = model.generate(**inputs, max_new_tokens=10, cache_implementation="static")
``` 

error log:
```
[rank0]:   File "/workspace/jiqing/transformers/src/transformers/cache_utils.py", line 1197, in update
[rank0]:     return _static_cache_update(
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/workspace/jiqing/transformers/src/transformers/cache_utils.py", line 54, in _static_cache_update
[rank0]:     k_cache.index_copy_(2, cache_position, key_states)
[rank0]: RuntimeError: index_copy_(): Source/destination tensor must have same slice shapes. Destination slice shape: 1 8 128 at dimen
sion 2 and source slice shape: 1 4 128 at dimension 0.
```


Comments:
Hi @yao-matrix . Could you please review the PR? Thanks!

Hi @Rocketknight1 . Would you please review this PR? Thanks!

Hi @SunMarc . This PR changes the StaticCache api, needs your review. Thanks!

Hi @ArthurZucker . I didn't meet any compile issue. We usually only use compile on the model forward rather than the whole generation. The init cache process will not be traced in my script.

`TORCHINDUCTOR_CPP_WRAPPER=1 TORCHINDUCTOR_FREEZING=1 torchrun --standalone --nproc-per-node 2 tp_hf.py`
```python
import torch
import torch.distributed as dist
from transformers import AutoTokenizer, AutoModelForCausalLM


model_id = "meta-llama/Llama-3.1-8B-Instruct"
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, tp_plan="auto")
if dist.is_initialized():
    print("Backend:", dist.get_backend())

model.forward = torch.compile(model.forward)
# Prepare input tokens
tokenizer = AutoTokenizer.from_pretrained(model_id)
prompt = "It is done, and submitted. You can play 'Survival of the Tastiest' on Android, and on the web."
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
outputs = model.generate(**inputs, max_new_tokens=10, cache_implementation="static")
print(tokenizer.batch_decode(outputs, skip_special_tokens=True))
```

I've added the test, please run `CUDA_VISIBLE_DEVICES=0,1 RUN_SLOW=1 pytest -sv tests/tensor_parallel/test_tensor_parallel.py`.

I can get the 30% speed-up when tp=2 and 60% speed-up when tp=4 on Intel Xeon 6th Gen CPU, and 20%+ speed-up on Intel XPU. 

```
| TP size | CPU  |  XPU  | A100  |
| 2       | 1.3x | 1.24x | 1.07x |
| 4       | 1.6x | 1.54x | 1.07x |
```

The A100 has almost no speed-up because the communication is slow via PCIe. I don't have other NV devices other than A100 PCIe. You can test it on an NVlink device like H100 to see the speed-up.

Ah sorry, if you use static cache, generate will compile by default! 

```diff
import torch
import torch.distributed as dist
from transformers import AutoTokenizer, AutoModelForCausalLM


model_id = "meta-llama/Llama-3.1-8B-Instruct"
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, tp_plan="auto")
if dist.is_initialized():
    print("Backend:", dist.get_backend())

-model.forward = torch.compile(model.forward)
# Prepare input tokens
tokenizer = AutoTokenizer.from_pretrained(model_id)
prompt = "It is done, and submitted. You can play 'Survival of the Tastiest' on Android, and on the web."
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
outputs = model.generate(**inputs, max_new_tokens=10, cache_implementation="static")
print(tokenizer.batch_decode(outputs, skip_special_tokens=True))
```
should already use compile! 

> ```diff
> -model.forward = torch.compile(model.forward)
> outputs = model.generate(**inputs, max_new_tokens=10, cache_implementation="static")
> ```
> 
> should already use compile!

It seems that only cuda will trigger the compile when using static cache, cpu still needs to compile manually.


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

> Very nice! I am just wondering if we have a way to lazy init... https://docs.pytorch.org/docs/stable/generated/torch.compiler.allow_in_graph.html or [pytorch/pytorch#156377](https://github.com/pytorch/pytorch/issues/156377)
> 
> As it does not really seem possible, can you make sure the tp_size fix is propagated to other caches if possible?

Done.

Thanks! I am a bit sad that we are not able to hide the tensor parallel but hoping we can in the future!