[Core] Allow disabling TP sharding for parallel Linear layer

[Isotr0py]

Purpose

• Refer to #[Feature]: Generalized the DP feature for ViT and multimodal backbone for the benefit of all models #22743, we need to disable TP at ViT if enable DP for it.
• Currently, to achieve this, we need to replaced all parallel linear layers to replicated linear, which will make model implementation too complicated if it needs weights packing logic. Furthermore, to fuse qkv_proj and gate_up_proj, we need to implement extra layers like MergedReplicatedParallelLinear etc, increasing maintenance burden.
• This PR allows parallel linear layers falling back to replicated mode with disable_tp=True.
• Will update if [FEAT] [Performance] Enable DP for ViT in Qwen2.5VL #22742 merged.

Test Plan

python examples/offline_inference/basic/generate.py --model RedHatAI/DeepSeek-V2.5-1210-FP8 -tp 4 --max-model-len 4096 --enforce-eager

Test Result

• FP8 deepseek-v3 model can still work with fused_qkv_a_proj after reverting MergedReplicatedParallelLinear back to MergedColumnParallelLinear

(Optional) Documentation Update

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

[gemini-code-assist]

Code Review

This pull request introduces a disable_tp flag to parallel linear layers, allowing them to fall back to a replicated mode. This is a valuable feature for models that need to conditionally disable tensor parallelism, for instance when data parallelism is enabled. The implementation across the various linear layer classes in vllm/model_executor/layers/linear.py is clean and consistent. The refactoring in other model files to leverage this new flag simplifies the code and demonstrates its utility. I have identified a critical typo and a potential prefixing issue in vllm/model_executor/models/step3_vl.py that should be addressed to ensure correctness, particularly for quantized models.

[DarkLight1337]

cc @mgoin

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @Isotr0py.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[Isotr0py]

@mgoin Can you please take a look to this PR? It could help us simplify the work to finish #22743. Thanks!

[minosfuture]

I think this PR broke TP on trunk. I see errors like

(VllmWorker TP5 pid=969158) ERROR 09-06 00:42:40 [multiproc_executor.py:588] Traceback (most recent call last):
(VllmWorker TP5 pid=969158) ERROR 09-06 00:42:40 [multiproc_executor.py:588]   File "/data/users/yming/gitrepos/vllm/vllm/v1/executor/multiproc_executor.py", line 562, in worker_main
(VllmWorker TP5 pid=969158) ERROR 09-06 00:42:40 [multiproc_executor.py:588]     worker = WorkerProc(*args, **kwargs)
(VllmWorker TP5 pid=969158) ERROR 09-06 00:42:40 [multiproc_executor.py:588]              ^^^^^^^^^^^^^^^^^^^^^^^^^^^
(VllmWorker TP5 pid=969158) ERROR 09-06 00:42:40 [multiproc_executor.py:588]   File "/data/users/yming/gitrepos/vllm/vllm/v1/executor/multiproc_executor.py", line 431, in __init__
(VllmWorker TP5 pid=969158) ERROR 09-06 00:42:40 [multiproc_executor.py:588]     self.worker.load_model()
(VllmWorker TP5 pid=969158) ERROR 09-06 00:42:40 [multiproc_executor.py:588]   File "/data/users/yming/gitrepos/vllm/vllm/v1/worker/gpu_worker.py", line 213, in load_model
(VllmWorker TP5 pid=969158) ERROR 09-06 00:42:40 [multiproc_executor.py:588]     self.model_runner.load_model(eep_scale_up=eep_scale_up)
(VllmWorker TP5 pid=969158) ERROR 09-06 00:42:40 [multiproc_executor.py:588]   File "/data/users/yming/gitrepos/vllm/vllm/v1/worker/gpu_model_runner.py", line 2131, in load_model
(VllmWorker TP5 pid=969158) ERROR 09-06 00:42:40 [multiproc_executor.py:588]     self.model = model_loader.load_model(
(VllmWorker TP5 pid=969158) ERROR 09-06 00:42:40 [multiproc_executor.py:588]                  ^^^^^^^^^^^^^^^^^^^^^^^^
(VllmWorker TP5 pid=969158) ERROR 09-06 00:42:40 [multiproc_executor.py:588]   File "/data/users/yming/gitrepos/vllm/vllm/model_executor/model_loader/base_loader.py", line 49, in load_model
(VllmWorker TP5 pid=969158) ERROR 09-06 00:42:40 [multiproc_executor.py:588]     self.load_weights(model, model_config)
(VllmWorker TP5 pid=969158) ERROR 09-06 00:42:40 [multiproc_executor.py:588]   File "/data/users/yming/gitrepos/vllm/vllm/model_executor/model_loader/default_loader.py", line 264, in load_weights
(VllmWorker TP5 pid=969158) ERROR 09-06 00:42:40 [multiproc_executor.py:588]     loaded_weights = model.load_weights(
(VllmWorker TP5 pid=969158) ERROR 09-06 00:42:40 [multiproc_executor.py:588]                      ^^^^^^^^^^^^^^^^^^^
(VllmWorker TP5 pid=969158) ERROR 09-06 00:42:40 [multiproc_executor.py:588]   File "/data/users/yming/gitrepos/vllm/vllm/model_executor/models/deepseek_v2.py", line 906, in load_weights
(VllmWorker TP5 pid=969158) ERROR 09-06 00:42:40 [multiproc_executor.py:588]     weight_loader(param, loaded_weight, shard_id)
(VllmWorker TP5 pid=969158) ERROR 09-06 00:42:40 [multiproc_executor.py:588]   File "/data/users/yming/gitrepos/vllm/vllm/model_executor/layers/linear.py", line 807, in weight_loader_v2
(VllmWorker TP5 pid=969158) ERROR 09-06 00:42:40 [multiproc_executor.py:588]     param.load_merged_column_weight(loaded_weight=loaded_weight,
(VllmWorker TP5 pid=969158) ERROR 09-06 00:42:40 [multiproc_executor.py:588]   File "/data/users/yming/gitrepos/vllm/vllm/model_executor/parameter.py", line 145, in load_merged_column_weight
(VllmWorker TP5 pid=969158) ERROR 09-06 00:42:40 [multiproc_executor.py:588]     loaded_weight = loaded_weight.narrow(self.output_dim,
(VllmWorker TP5 pid=969158) ERROR 09-06 00:42:40 [multiproc_executor.py:588]                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(VllmWorker TP5 pid=969158) ERROR 09-06 00:42:40 [multiproc_executor.py:588] IndexError: start out of range (expected to be in range of [-576, 576], but got 2880)

reverting this fixed it. Lemme know if you need reproducing steps. I think kimi TP16 or deepseek DP2TP8EP would reproduce it.

[Isotr0py]

I think kimi TP16 or deepseek DP2TP8EP would reproduce it.

Ooops, can you provide the code for reproduction?

[minosfuture]

#node 0
vllm serve deepseek-ai/DeepSeek-V3-0324   --tensor-parallel-size 8   --enable-expert-parallel   --data-parallel-size 2   --data-parallel-size-local 1   --data-parallel-address $MASTER_IP  --data-parallel-rpc-port 13345

#node 1
vllm serve deepseek-ai/DeepSeek-V3-0324   --tensor-parallel-size 8   --enable-expert-parallel   --data-parallel-size 2   --data-parallel-size-local 1   --data-parallel-start-rank 1   --data-parallel-address  $MASTER_IP   --data-parallel-rpc-port 13345   --headless

I think single node test of a smaller model should also be able to reproduce it but I haven't tried.

[Isotr0py]

Hi @minosfuture, I think #24367 should fix this issue. Can you have a look? Thanks!