Failed to run `Qwen3-235B-A22B` with `tp_size=8`
Author: Hermit-w
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.56.1
- Platform: Linux-4.18.0-193.28.1.el8_2.x86_64-x86_64-with-glibc2.28
- Python version: 3.12.11
- Huggingface_hub version: 0.34.4
- Safetensors version: 0.5.3
- Accelerate version: 1.10.1
- Accelerate config:    not found
- DeepSpeed version: 0.17.5
- PyTorch version (accelerator?): 2.8.0+cu128 (CUDA)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: yes
- Using GPU in script?: yes
- GPU type: NVIDIA H800

### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [x] My own modified scripts

### Tasks

- [x] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

Use the following cmd to run the script.
```bash
torchrun --nproc-per-node=8 run_model.py
```
The script is:
```python
import torch

from transformers import AutoModelForCausalLM, AutoTokenizer

if __name__ == '__main__':
    torch.distributed.init_process_group("nccl")
    model_name = "Qwen/Qwen3-235B-A22B"
    model_path = model_name
    tp_size = 8
    model = AutoModelForCausalLM.from_pretrained(model_path, tp_plan="auto", tp_size=tp_size, dtype="auto")
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    prompt = "Hello, please introduce yourself.\n"
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    outputs = model(
        **inputs,
        use_cache=True,
    )
    torch.distributed.destroy_process_group()
```
And it will cause the following error, which seems to be the problem about shape related to tp_plan.

<img width="1524" height="149" alt="Image" src="https://github.com/user-attachments/assets/e7b46744-3780-4800-8f72-2288a3baf35c" />

### Expected behavior

Running model without error.

Comments:
I suppose the problem may come from that `Qwen3-235B-A22B` contains only 4 kv heads in attention parts, but existing tensor parallelism method slice the head into 8 pieces with `tp_size=8`, so that each device don't contain full attention head.
Therefore, each device will only produce hidden states with size of 64, which is half of the `head_dim=128`. And with the length of prompt being 6, we may get the size of `6*64=384` which is the number shown in the above fig.
However, it may be difficult to run this model on only 4 gpus due to limit of memory, and I have no idea to customize my own `tp_plan` and pass it to the loading function since it seems to only support `"auto"`.

cc @cyrilvallez @arthurzucker maybe? We may need to adapt `tp_plan` for these extreme GQA models where `num_kv_heads < tp_size`

Hey! This is similar to https://github.com/huggingface/transformers/issues/40953#issuecomment-3311635988, you can check my answers for more explanations! Unfortunately this is a limitation of how TP works for the attention matrices. We could maybe surface an error about it though!

> Hey! This is similar to [#40953 (comment)](https://github.com/huggingface/transformers/issues/40953#issuecomment-3311635988), you can check my answers for more explanations! Unfortunately this is a limitation of how TP works for the attention matrices. We could maybe surface an error about it though!

Yeah, I see. Thanks for your replying.
~~However, I'm wondering if simply duplicating the kv head across 2 adjacent gpus can be compatiable with `tp_size=8` and avoid being less efficient in the method of chaning q/k/v to colwise and o to rowwise you mentioned in another issue (e.g. put the first kv_head into gpu1 and gpu 2 without slicing and do so on for other kv heads)~~
I realize this is what you said. Sorry for bothering.
By the way, I also want to know is there any method provided for us to pass our own `tp_plan`?

For now, we don't support passing it directly an an argument to `from_pretrained`. You could however change the class attribute on the config directly, prior to loading the model for example:

```
from transformers import LlamaConfig, AutoModelForCausalLM

LlamaConfig.base_model_tp_plan = my_plan
model = AutoModelForCausalLM.from_pretrained("LlamaModel", tp_plan="auto")
```

We should allow passing it for v5 for sure!