Documentation says manual `tp_plan` can be set but it can't be set.

[st81]

System Info

• transformers version: 4.57.0.dev0
• platform: linux-5.10.102.1-microsoft-standard-wsl2-x86_64-with-glibc2.31
• python version: 3.12.11
• huggingface_hub version: 1.0.0.rc1
• safetensors version: 0.5.3
• accelerate version: 1.9.0
• accelerate config: not found
• deepspeed version: 0.17.6
• pytorch version (accelerator?): 2.8.0+cu128 (cuda)
• using distributed or parallel set-up in script?: yes
• using gpu in script?: yes
• gpu type: nvidia geforce rtx 2070 super

Who can help?

• distributed: @ArthurZucker
• documentation: @stevhliu

Information

• The official example scripts
• My own modified scripts

Tasks

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

Reproduction

Documentation says manual tp_plan can be set but if I set it, it raises error.

https://huggingface.co/docs/transformers/main/en/perf_infer_gpu_multi?sharding=manual+plan#partitioning-a-model

from transformers import AutoModelForCausalLM

model_id = "Qwen/Qwen2-1.5B-Instruct"
tp_plan = {
    "model.layers.*.self_attn.q_proj": "colwise",
    "model.layers.*.self_attn.k_proj": "colwise",
    "model.layers.*.self_attn.v_proj": "colwise",
    "model.layers.*.self_attn.o_proj": "rowwise",
}
model = AutoModelForCausalLM.from_pretrained(model_id, tp_plan=tp_plan)

torchrun --nproc-per-node=1 test.py 

Error:

Traceback (most recent call last):
  File "/home/shutotakahashi/projects/transformers-uv/transformers/test.py", line 11, in <module>
    model = AutoModelForCausalLM.from_pretrained(model_id, tp_plan=tp_plan)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/shutotakahashi/projects/transformers-uv/transformers/src/transformers/models/auto/auto_factory.py", line 389, in from_pretrained
    return model_class.from_pretrained(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/shutotakahashi/projects/transformers-uv/transformers/src/transformers/modeling_utils.py", line 281, in _wrapper
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/shutotakahashi/projects/transformers-uv/transformers/src/transformers/modeling_utils.py", line 4615, in from_pretrained
    raise ValueError(f"tp_plan supports 'auto' only for now but got {tp_plan}.")
ValueError: tp_plan supports 'auto' only for now but got {'model.layers.*.self_attn.q_proj': 'colwise', 'model.layers.*.self_attn.k_proj': 'colwise', 'model.layers.*.self_attn.v_proj': 'colwise', 'model.layers.*.self_attn.o_proj': 'rowwise'}.

Expected behavior

Looking at the source code, it appears that manual tp_plan is not supported yet, but the documentation says it is supported. I expect either:

• The documentation should indicate that manual tp_plan is not yet supported, or the manual tp_plan sections should be removed from the documentation
• If manual tp_plan is intended to be supported, it should work without error.

The error ValueError(f"tp_plan supports 'auto' only for now but got {tp_plan}.") was introduced in #34184, which was merged about 1 year ago.

[Rocketknight1]

cc @SunMarc @Cyrilvallez

[Cyrilvallez]

Indeed it's not supported at the moment, though we could easily make it work I believe. But in the meantime, we can indeed update the doc! cc @stevhliu if you have the time! 🤗