Revert "[Core] Allow disabling TP sharding for parallel Linear layer (vllm-project#23024)"

This reverts commit 53b19cc.

Author: minosfuture

vllm/model_executor/layers/linear.py

@@ -69,7 +69,6 @@ def __init__(self, load_config: LoadConfig):
         # Store all module names (from transformers) that support
         # BNB quantization.
         self.target_modules: list[str] = []
-        self.tp_disabled_modules: list[str] = []
         # Store the mapping of expert parameters for MoE models.
         self.expert_params_mapping: list[tuple[str, str, int, str]] = []
         # mapping weight names from transformers to vllm.
@@ -323,24 +322,14 @@ def _unquantized_generator(self, hf_weights_files, use_safetensors,
                                quant_state_dict) -> Generator:
         from bitsandbytes.functional import quantize_4bit
 
-        global_tp_size = get_tensor_model_parallel_world_size()
-        global_tp_rank = get_tensor_model_parallel_rank()
+        tp_size = get_tensor_model_parallel_world_size()
+        tp_rank = get_tensor_model_parallel_rank()
 
         for (
                 org_weight_name,
                 mapped_weight_name,
                 weight_tensor,
         ) in self._hf_weight_iter(hf_weights_files, use_safetensors):
-
-            # override tp_size and tp_rank if the module has disabled TP
-            if any(tp_disabled_module in mapped_weight_name
-                   for tp_disabled_module in self.tp_disabled_modules):
-                tp_size = 1
-                tp_rank = 0
-            else:
-                tp_size = global_tp_size
-                tp_rank = global_tp_rank
-
             if any(target_module in mapped_weight_name
                    for target_module in self.target_modules
                    ) and mapped_weight_name.endswith(".weight"):
@@ -429,16 +418,12 @@ def _get_bnb_target_modules(self, model: nn.Module) -> None:
                     # Map vllm's names to transformers's names.
                     rep_name, sub_modules = modules_info
                     for sub_name in sub_modules:
-                        new_name = name.replace(rep_name, sub_name)
-                        self.target_modules.append(new_name)
-                        if module.disable_tp:
-                            self.tp_disabled_modules.append(new_name)
+                        self.target_modules.append(
+                            name.replace(rep_name, sub_name))
                 # Add original module name even if the module has stacked map,
                 # in case model has a mixture of disk-merged and disk-split
                 # weights with same last name.
                 self.target_modules.append(name)
-                if module.disable_tp:
-                    self.tp_disabled_modules.append(name)
             elif isinstance(module, FusedMoE) and hasattr(
                     module.quant_method, "quant_config"):
                 # TODO: support FusedMoE with prequant and 8bit.

vllm/model_executor/model_loader/bitsandbytes_loader.py

@@ -43,6 +43,7 @@
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (ColumnParallelLinear,
                                                MergedColumnParallelLinear,
+                                               MergedReplicatedLinear,
                                                ReplicatedLinear,
                                                RowParallelLinear)
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -434,13 +435,12 @@ def __init__(
         self.max_position_embeddings = max_position_embeddings
 
         if self.q_lora_rank is not None:
-            self.fused_qkv_a_proj = MergedColumnParallelLinear(
+            self.fused_qkv_a_proj = MergedReplicatedLinear(
                 self.hidden_size,
                 [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],
                 bias=False,
                 quant_config=quant_config,
-                prefix=f"{prefix}.fused_qkv_a_proj",
-                disable_tp=True)
+                prefix=f"{prefix}.fused_qkv_a_proj")
         else:
             self.kv_a_proj_with_mqa = ReplicatedLinear(
                 self.hidden_size,

vllm/model_executor/models/deepseek_v2.py

@@ -51,10 +51,14 @@
 from vllm.logger import init_logger
 from vllm.model_executor import SamplingMetadata
 from vllm.model_executor.layers.layernorm import RMSNorm
+# yapf: disable
 from vllm.model_executor.layers.linear import (ColumnParallelLinear,
                                                MergedColumnParallelLinear,
+                                               MergedReplicatedLinear,
                                                QKVParallelLinear,
+                                               ReplicatedLinear,
                                                RowParallelLinear)
+# yapf: enable
 from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm.model_executor.models.module_mapping import MultiModelKeys
@@ -170,22 +174,20 @@ def __init__(
         use_data_parallel: bool = False,
     ):
         super().__init__()
-        self.gate_up_proj = MergedColumnParallelLinear(
-            input_size=in_features,
-            output_sizes=[hidden_features] * 2,
-            bias=bias,
-            quant_config=quant_config,
-            prefix=f"{prefix}.gate_up_proj",
-            disable_tp=use_data_parallel,
-        )
-        self.down_proj = RowParallelLinear(
-            hidden_features,
-            in_features,
-            bias=bias,
-            quant_config=quant_config,
-            prefix=f"{prefix}.down_proj",
-            disable_tp=use_data_parallel,
-        )
+        cls_gate_up = (MergedReplicatedLinear
+                       if use_data_parallel else MergedColumnParallelLinear)
+        self.gate_up_proj = cls_gate_up(input_size=in_features,
+                                        output_sizes=[hidden_features] * 2,
+                                        bias=bias,
+                                        quant_config=quant_config,
+                                        prefix=f"{prefix}.gate_up_proj")
+        cls_down = (ReplicatedLinear
+                    if use_data_parallel else RowParallelLinear)
+        self.down_proj = cls_down(hidden_features,
+                                  in_features,
+                                  bias=bias,
+                                  quant_config=quant_config,
+                                  prefix=f"{prefix}.down_proj")
         self.act_fn = SiluAndMul()
 
     def forward(self, x: torch.Tensor):
@@ -232,32 +234,48 @@ def __init__(
         # Per attention head and per partition values.
         self.tp_size = (1 if use_data_parallel else
                         get_tensor_model_parallel_world_size())
-        self.tp_rank = (0 if use_data_parallel else
-                        parallel_state.get_tensor_model_parallel_rank())
+        self.tp_rank = parallel_state.get_tensor_model_parallel_rank()
         self.hidden_size_per_attention_head = dist_utils.divide(
             projection_size, num_heads)
         self.num_attention_heads_per_partition = dist_utils.divide(
             num_heads, self.tp_size)
 
-        self.qkv = QKVParallelLinear(
-            hidden_size=embed_dim,
-            head_size=self.hidden_size_per_attention_head,
-            total_num_heads=num_heads,
-            total_num_kv_heads=num_heads,
-            bias=False,
-            quant_config=quant_config,
-            # Change qkv prefix to align with GLM-4.5V-FP8 quantization cfg
-            prefix=f"{prefix}.qkv_proj" if quant_config else f"{prefix}.qkv",
-            disable_tp=use_data_parallel,
-        )
-        self.proj = RowParallelLinear(
-            input_size=projection_size,
-            output_size=embed_dim,
-            quant_config=quant_config,
-            prefix=f"{prefix}.proj",
-            bias=False,
-            disable_tp=use_data_parallel,
-        )
+        if use_data_parallel:
+            self.qkv = ReplicatedLinear(
+                input_size=embed_dim,
+                output_size=3 * projection_size,
+                bias=False,
+                quant_config=quant_config,
+                # Change qkv prefix to align with GLM-4.5V-FP8 quantization cfg
+                prefix=f"{prefix}.qkv_proj"
+                if quant_config else f"{prefix}.qkv",
+            )
+            self.proj = ReplicatedLinear(
+                input_size=projection_size,
+                output_size=embed_dim,
+                quant_config=quant_config,
+                prefix=f"{prefix}.proj",
+                bias=False,
+            )
+        else:
+            self.qkv = QKVParallelLinear(
+                hidden_size=embed_dim,
+                head_size=self.hidden_size_per_attention_head,
+                total_num_heads=num_heads,
+                total_num_kv_heads=num_heads,
+                bias=False,
+                quant_config=quant_config,
+                # Change qkv prefix to align with GLM-4.5V-FP8 quantization cfg
+                prefix=f"{prefix}.qkv_proj"
+                if quant_config else f"{prefix}.qkv",
+            )
+            self.proj = RowParallelLinear(
+                input_size=projection_size,
+                output_size=embed_dim,
+                quant_config=quant_config,
+                prefix=f"{prefix}.proj",
+                bias=False,
+            )
 
         # Detect attention implementation.
         self.attn_backend: _Backend = get_vit_attn_backend(support_fa=True)
@@ -476,31 +494,41 @@ def __init__(
     ) -> None:
         super().__init__()
         self.hidden_size = d_model
-        self.proj = ColumnParallelLinear(
-            self.hidden_size,
-            self.hidden_size,
-            bias=bias,
-            gather_output=True,
-            quant_config=quant_config,
-            prefix=f"{prefix}.proj",
-            disable_tp=use_data_parallel,
-        )
+        if use_data_parallel:
+            self.proj = ReplicatedLinear(
+                input_size=self.hidden_size,
+                output_size=self.hidden_size,
+                bias=bias,
+                quant_config=quant_config,
+                prefix=f"{prefix}.proj",
+            )
+        else:
+            self.proj = ColumnParallelLinear(
+                self.hidden_size,
+                self.hidden_size,
+                bias=bias,
+                gather_output=True,
+                quant_config=quant_config,
+                prefix=f"{prefix}.proj",
+            )
         self.post_projection_norm = nn.LayerNorm(self.hidden_size)
-        self.gate_up_proj = MergedColumnParallelLinear(
+        cls_gate_up = (MergedReplicatedLinear
+                       if use_data_parallel else MergedColumnParallelLinear)
+        self.gate_up_proj = cls_gate_up(
             input_size=self.hidden_size,
             output_sizes=[context_dim] * 2,
             bias=bias,
             quant_config=quant_config,
             prefix=f"{prefix}.gate_up_proj",
-            disable_tp=use_data_parallel,
         )
-        self.down_proj = RowParallelLinear(
+        cls_down = (ReplicatedLinear
+                    if use_data_parallel else RowParallelLinear)
+        self.down_proj = cls_down(
             context_dim,
             self.hidden_size,
             bias=bias,
             quant_config=quant_config,
             prefix=f"{prefix}.down_proj",
-            disable_tp=use_data_parallel,
         )
         self.act_fn = SiluAndMul()
         self.extra_activation_func = nn.GELU()

vllm/model_executor/models/glm4_1v.py

@@ -48,6 +48,7 @@
 # yapf: disable
 from vllm.model_executor.layers.linear import (ColumnParallelLinear,
                                                MergedColumnParallelLinear,
+                                               MergedReplicatedLinear,
                                                QKVParallelLinear,
                                                ReplicatedLinear,
                                                RowParallelLinear)
@@ -177,20 +178,22 @@ def __init__(self,
                  prefix: str = "",
                  use_data_parallel: bool = False):
         super().__init__()
-        self.gate_up_proj = MergedColumnParallelLinear(
+        cls_gate_up_proj = (MergedReplicatedLinear if use_data_parallel else
+                            MergedColumnParallelLinear)
+        self.gate_up_proj = cls_gate_up_proj(
             input_size=in_features,
             output_sizes=[hidden_features] * 2,  # [gate_proj, up_proj]
             bias=bias,
             quant_config=quant_config,
-            prefix=f"{prefix}.gate_up_proj",
-            disable_tp=use_data_parallel)
-
-        self.down_proj = RowParallelLinear(hidden_features,
-                                           in_features,
-                                           bias=bias,
-                                           quant_config=quant_config,
-                                           prefix=f"{prefix}.down_proj",
-                                           disable_tp=use_data_parallel)
+            prefix=f"{prefix}.gate_up_proj")
+
+        cls_down_proj = (ReplicatedLinear
+                         if use_data_parallel else RowParallelLinear)
+        self.down_proj = cls_down_proj(hidden_features,
+                                       in_features,
+                                       bias=bias,
+                                       quant_config=quant_config,
+                                       prefix=f"{prefix}.down_proj")
         self.act_fn = act_fn
 
     def forward(self, x: torch.Tensor):
@@ -240,21 +243,30 @@ def __init__(
         self.num_attention_heads_per_partition = dist_utils.divide(
             num_heads, self.tp_size)
 
-        self.qkv = QKVParallelLinear(
-            hidden_size=embed_dim,
-            head_size=self.hidden_size_per_attention_head,
-            total_num_heads=num_heads,
-            total_num_kv_heads=num_heads,
-            bias=True,
-            quant_config=quant_config,
-            prefix=f"{prefix}.qkv",
-            disable_tp=use_data_parallel)
-
-        self.proj = RowParallelLinear(input_size=projection_size,
-                                      output_size=embed_dim,
-                                      quant_config=quant_config,
-                                      prefix=f"{prefix}.proj",
-                                      disable_tp=use_data_parallel)
+        if use_data_parallel:
+            self.qkv = ReplicatedLinear(embed_dim,
+                                        self.hidden_size_per_attention_head *
+                                        3 * num_heads,
+                                        bias=True,
+                                        quant_config=quant_config,
+                                        prefix=f"{prefix}.qkv")
+
+        else:
+            self.qkv = QKVParallelLinear(
+                hidden_size=embed_dim,
+                head_size=self.hidden_size_per_attention_head,
+                total_num_heads=num_heads,
+                total_num_kv_heads=num_heads,
+                bias=True,
+                quant_config=quant_config,
+                prefix=f"{prefix}.qkv")
+
+        cls_proj = (ReplicatedLinear
+                    if use_data_parallel else RowParallelLinear)
+        self.proj = cls_proj(input_size=projection_size,
+                             output_size=embed_dim,
+                             quant_config=quant_config,
+                             prefix=f"{prefix}.proj")
 
         # Detect attention implementation.
         self.attn_backend: _Backend = get_vit_attn_backend(support_fa=True)