A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
Support more parallel styles in Transformers backend TP (#22651)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Harry Mellor committed
Aug 11, 2025 at 18:42 UTC
458e74eb907f96069e6d8a4f3c9f457001fef2ea
1 file changed
+13
-5
vllm/model_executor/models/transformers.py
+13
-5
@@ -107,10 +107,17 @@ def replace_linear_class(
107
raise ValueError(
108
f"Unsupported parallel style type {type(style)}, expected str")
109
110
- vllm_linear_cls = {
111
- "colwise": ColumnParallelLinear,
112
- "rowwise": RowParallelLinear,
113
- }.get(style, ReplicatedLinear)
110
+ vllm_linear_cls, vllm_linear_kwargs = {
111
+ "colwise": (ColumnParallelLinear, {}),
112
+ "colwise_rep": (ColumnParallelLinear, {
113
+ "gather_output": True
114
+ }),
115
+ "rowwise": (RowParallelLinear, {}),
116
+ "rowwise_rep": (RowParallelLinear, {
117
+ "input_is_parallel": False
118
+ }),
119
+ "replicate": (ReplicatedLinear, {}),
120
+ }.get(style, (ReplicatedLinear, {}))
121
122
return vllm_linear_cls(
123
input_size=linear.in_features,
@@ -118,6 +125,7 @@ def replace_linear_class(
125
bias=linear.bias is not None,
126
quant_config=quant_config,
127
return_bias=False,
128
+ **vllm_linear_kwargs,
129
)
130
131
@@ -506,7 +514,7 @@ class TransformersBase(nn.Module, SupportsQuant, SupportsLoRA, SupportsPP):
514
# Some weight loaders expect linear layers to inherit from vLLM's
515
# LinearBase class, so we set a default style which causes any
516
# unspecified linear layers to be replaced with ReplicatedLinear
509
- tp_plan[".*"] = "replicated"
517
+ tp_plan[".*"] = "replicate"
518
519
def _tensor_parallel(module: nn.Module, prefix: str = ""):
520
for child_name, child_module in module.named_children():