@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 53b19ccdd5

Fork 0

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

[Core] Allow disabling TP sharding for parallel Linear layer (#23024)

Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn> Signed-off-by: Isotr0py <2037008807@qq.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Isotr0py committed Sep 6, 2025 at 13:53 UTC 53b19ccdd5a04f442283a9b077b78da3b0df79d4

7 files changed +205 -282

vllm/model_executor/layers/linear.py

+70 -105

@@ -223,6 +223,7 @@ class LinearBase(CustomOp):

223 quant_config: Quantization configure.

224 prefix: Prefix for parameter names.

225 return_bias: If true, return bias together with outputs in forward pass.

226 + disable_tp: If true, tensor parallelism will be disabled for this layer.

227 """

228

229 def __init__(

@@ -235,6 +236,7 @@ class LinearBase(CustomOp):

236 prefix: str = "",

237 *,

238 return_bias: bool = True,

239 + disable_tp: bool = False,

240 ):

241 super().__init__()

242

@@ -254,6 +256,17 @@ class LinearBase(CustomOp):

256 self.quant_method = quant_config.get_quant_method(self,

257 prefix=prefix)

258 self.return_bias = return_bias

259 + self.disable_tp = disable_tp

260 + self.tp_rank = (get_tensor_model_parallel_rank()

261 + if not disable_tp else 0)

262 + self.tp_size = (get_tensor_model_parallel_world_size()

263 + if not disable_tp else 1)

264 +

265 + def __post_init__(self):

266 + for param in self.parameters():

267 + if isinstance(param, BasevLLMParameter):

268 + param.tp_rank = self.tp_rank

269 + param.tp_size = self.tp_size

270

271

272 @CustomOp.register("replicated_linear")

@@ -270,6 +283,7 @@ class ReplicatedLinear(LinearBase):

283 prefix: The name of the layer in the state dict, including all parents

284 (e.g. model.layers.0.qkv_proj)

285 return_bias: If true, return bias together with outputs in forward pass.

286 + disable_tp: Take no effect for replicated linear layers.

287 """

288

289 def __init__(

@@ -283,26 +297,21 @@ class ReplicatedLinear(LinearBase):

297 prefix: str = "",

298 *,

299 return_bias: bool = True,

300 + disable_tp: bool = False,

301 ):

287 - # If MergedReplicatedLinear, use output size of each partition.

288 - if hasattr(self, "output_sizes"):

289 - self.output_partition_sizes = self.output_sizes

290 - else:

291 - self.output_partition_sizes = [output_size]

292 -

302 super().__init__(input_size,

303 output_size,

304 skip_bias_add,

305 params_dtype,

306 quant_config,

307 prefix=prefix,

299 - return_bias=return_bias)

308 + return_bias=return_bias,

309 + disable_tp=disable_tp)

310

311 # All the linear layer supports quant method.

312 assert self.quant_method is not None

313 self.quant_method.create_weights(self,

304 - self.input_size,

305 - self.output_partition_sizes,

314 + self.input_size, [self.output_size],

315 self.input_size,

316 self.output_size,

317 self.params_dtype,

@@ -358,74 +367,6 @@ class ReplicatedLinear(LinearBase):

367 return s

368

369

361 -class MergedReplicatedLinear(ReplicatedLinear):

362 - """Replicated linear layer.

363 -

364 - Args:

365 - input_size: input dimension of the linear layer.

366 - output_sizes: list of output dimensions of the linear layer.

367 - bias: If true, add bias.

368 - skip_bias_add: If true, skip adding bias but instead return it.

369 - params_dtype: Data type for the parameters.

370 - quant_config: Quantization configure.

371 - prefix: The name of the layer in the state dict, including all parents

372 - (e.g. model.layers.0.qkv_proj)

373 - return_bias: If true, return bias together with outputs in forward pass.

374 - """

375 -

376 - def __init__(

377 - self,

378 - input_size: int,

379 - output_sizes: list[int],

380 - bias: bool = True,

381 - skip_bias_add: bool = False,

382 - params_dtype: Optional[torch.dtype] = None,

383 - quant_config: Optional[QuantizationConfig] = None,

384 - prefix: str = "",

385 - *,

386 - return_bias: bool = True,

387 - ):

388 - self.output_sizes = output_sizes

389 - super().__init__(input_size,

390 - sum(output_sizes),

391 - bias,

392 - skip_bias_add,

393 - params_dtype,

394 - quant_config,

395 - prefix=prefix,

396 - return_bias=return_bias)

397 -

398 - def weight_loader(self,

399 - param: Union[Parameter, BasevLLMParameter],

400 - loaded_weight: torch.Tensor,

401 - loaded_shard_id: Optional[int] = None):

402 - assert loaded_shard_id is not None

403 - assert loaded_shard_id < len(self.output_sizes)

404 -

405 - if isinstance(param, BlockQuantScaleParameter):

406 - from vllm.model_executor.layers.quantization.fp8 import (

407 - Fp8LinearMethod, Fp8MoEMethod)

408 - assert self.quant_method is not None

409 - assert isinstance(self.quant_method,

410 - (Fp8LinearMethod, Fp8MoEMethod))

411 - weight_block_size = self.quant_method.quant_config.weight_block_size

412 - assert weight_block_size is not None

413 - block_n, _ = weight_block_size[0], weight_block_size[1]

414 - shard_offset = (

415 - (sum(self.output_sizes[:loaded_shard_id]) + block_n - 1) //

416 - block_n)

417 - shard_size = ((self.output_sizes[loaded_shard_id] + block_n - 1) //

418 - block_n)

419 - elif isinstance(param, PerTensorScaleParameter):

420 - shard_offset = loaded_shard_id

421 - shard_size = 1

422 - else:

423 - shard_offset = sum(self.output_sizes[:loaded_shard_id])

424 - shard_size = self.output_sizes[loaded_shard_id]

425 -

426 - param.data[shard_offset:shard_offset + shard_size] = loaded_weight

427 -

428 -

370 @CustomOp.register("column_parallel_linear")

371 class ColumnParallelLinear(LinearBase):

372 """Linear layer with column parallelism.

@@ -448,7 +389,9 @@ class ColumnParallelLinear(LinearBase):

389 output_sizes: list of output sizes packed into one output, like for QKV

390 the list would be size 3.

391 prefix: The name of the layer in the state dict, including all parents

451 - (e.g. model.layers.0.qkv_proj)

392 + (e.g. model.layers.0.qkv_proj)

393 + return_bias: If true, return bias together with outputs in forward pass.

394 + disable_tp: If true, weights matrix won't be sharded through tp rank.

395 """

396

397 def __init__(

@@ -464,9 +407,13 @@ class ColumnParallelLinear(LinearBase):

407 prefix: str = "",

408 *,

409 return_bias: bool = True,

410 + disable_tp: bool = False,

411 ):

412 # Divide the weight matrix along the last dimension.

469 - self.tp_size = get_tensor_model_parallel_world_size()

413 + self.tp_rank = (get_tensor_model_parallel_rank()

414 + if not disable_tp else 0)

415 + self.tp_size = (get_tensor_model_parallel_world_size()

416 + if not disable_tp else 1)

417 self.input_size_per_partition = input_size

418 self.output_size_per_partition = divide(output_size, self.tp_size)

419 self.output_partition_sizes = [self.output_size_per_partition]

@@ -483,7 +430,8 @@ class ColumnParallelLinear(LinearBase):

430 params_dtype,

431 quant_config,

432 prefix,

486 - return_bias=return_bias)

433 + return_bias=return_bias,

434 + disable_tp=disable_tp)

435

436 self.gather_output = gather_output

437

@@ -512,8 +460,6 @@ class ColumnParallelLinear(LinearBase):

460 else:

461 self.register_parameter("bias", None)

462

515 - self.tp_rank = get_tensor_model_parallel_rank()

516 -

463 def weight_loader(self, param: Parameter, loaded_weight: torch.Tensor):

464

465 output_dim = getattr(param, "output_dim", None)

@@ -554,7 +500,8 @@ class ColumnParallelLinear(LinearBase):

500 assert param_data.shape == loaded_weight.shape

501 param_data.copy_(loaded_weight)

502

557 - def weight_loader_v2(self, param: Parameter, loaded_weight: torch.Tensor):

503 + def weight_loader_v2(self, param: BasevLLMParameter,

504 + loaded_weight: torch.Tensor):

505 # Special case for loading scales off disk, which often do not

506 # have a shape (such as in the case of AutoFP8).

507 if len(loaded_weight.shape) == 0:

@@ -570,7 +517,7 @@ class ColumnParallelLinear(LinearBase):

517 # Matrix multiply.

518 assert self.quant_method is not None

519 output_parallel = self.quant_method.apply(self, input_, bias)

573 - if self.gather_output:

520 + if self.gather_output and self.tp_size > 1:

521 # All-gather across the partitions.

522 output = tensor_model_parallel_all_gather(output_parallel)

523 else:

@@ -584,7 +531,7 @@ class ColumnParallelLinear(LinearBase):

531 s = f"in_features={self.input_size}"

532 s += f", output_features={self.output_size_per_partition}"

533 s += f", bias={self.bias is not None}"

587 - s += f", tp_size={get_tensor_model_parallel_world_size()}"

534 + s += f", tp_size={self.tp_size}"

535 s += f", gather_output={self.gather_output}"

536 return s

537

@@ -611,6 +558,8 @@ class MergedColumnParallelLinear(ColumnParallelLinear):

558 prefix: The name of the layer in the state dict, including all parents

559 (e.g. model.layers.0.qkv_proj)

560 return_bias: If true, return bias together with outputs in forward pass.

561 + disable_tp: If true, all weights matrix won't be sharded, this layer

562 + will be treated as a "Replicated" MergedLinear.

563 """

564

565 def __init__(

@@ -625,10 +574,13 @@ class MergedColumnParallelLinear(ColumnParallelLinear):

574 prefix: str = "",

575 *,

576 return_bias: bool = True,

577 + disable_tp: bool = False,

578 ):

579 self.output_sizes = output_sizes

630 - self.tp_size = get_tensor_model_parallel_world_size()

631 - self.tp_rank = get_tensor_model_parallel_rank()

580 + self.tp_size = (get_tensor_model_parallel_world_size()

581 + if not disable_tp else 1)

582 + self.tp_rank = (get_tensor_model_parallel_rank()

583 + if not disable_tp else 0)

584

585 assert all(output_size % self.tp_size == 0

586 for output_size in output_sizes)

@@ -640,7 +592,8 @@ class MergedColumnParallelLinear(ColumnParallelLinear):

592 params_dtype=params_dtype,

593 quant_config=quant_config,

594 prefix=prefix,

643 - return_bias=return_bias)

595 + return_bias=return_bias,

596 + disable_tp=disable_tp)

597

598 def weight_loader(self,

599 param: Parameter,

@@ -832,8 +785,6 @@ class MergedColumnParallelLinear(ColumnParallelLinear):

785

786 assert loaded_shard_id < len(self.output_sizes)

787

835 - tp_size = get_tensor_model_parallel_world_size()

836 -

788 if isinstance(param, BlockQuantScaleParameter):

789 from vllm.model_executor.layers.quantization.fp8 import (

790 Fp8LinearMethod, Fp8MoEMethod)

@@ -845,17 +796,19 @@ class MergedColumnParallelLinear(ColumnParallelLinear):

796 block_n, _ = weight_block_size[0], weight_block_size[1]

797 shard_offset = (

798 (sum(self.output_sizes[:loaded_shard_id]) + block_n - 1) //

848 - block_n) // tp_size

799 + block_n) // self.tp_size

800 shard_size = ((self.output_sizes[loaded_shard_id] + block_n - 1) //

850 - block_n // tp_size)

801 + block_n // self.tp_size)

802 else:

852 - shard_offset = sum(self.output_sizes[:loaded_shard_id]) // tp_size

853 - shard_size = self.output_sizes[loaded_shard_id] // tp_size

803 + shard_offset = sum(

804 + self.output_sizes[:loaded_shard_id]) // self.tp_size

805 + shard_size = self.output_sizes[loaded_shard_id] // self.tp_size

806

807 param.load_merged_column_weight(loaded_weight=loaded_weight,

808 shard_id=loaded_shard_id,

809 shard_offset=shard_offset,

858 - shard_size=shard_size)

810 + shard_size=shard_size,

811 + tp_rank=self.tp_rank)

812

813

814 class QKVParallelLinear(ColumnParallelLinear):

@@ -883,6 +836,7 @@ class QKVParallelLinear(ColumnParallelLinear):

836 prefix: The name of the layer in the state dict, including all parents

837 (e.g. model.layers.0.qkv_proj)

838 return_bias: If true, return bias together with outputs in forward pass.

839 + disable_tp: If true, weights matrix won't be sharded through tp rank.

840 """

841

842 def __init__(

@@ -898,6 +852,7 @@ class QKVParallelLinear(ColumnParallelLinear):

852 prefix: str = "",

853 *,

854 return_bias: bool = True,

855 + disable_tp: bool = False,

856 ):

857 self.hidden_size = hidden_size

858 self.head_size = head_size

@@ -906,7 +861,8 @@ class QKVParallelLinear(ColumnParallelLinear):

861 total_num_kv_heads = total_num_heads

862 self.total_num_kv_heads = total_num_kv_heads

863 # Divide the weight matrix along the last dimension.

909 - tp_size = get_tensor_model_parallel_world_size()

864 + tp_size = (get_tensor_model_parallel_world_size()

865 + if not disable_tp else 1)

866 self.num_heads = divide(self.total_num_heads, tp_size)

867 if tp_size >= self.total_num_kv_heads:

868 self.num_kv_heads = 1

@@ -932,7 +888,8 @@ class QKVParallelLinear(ColumnParallelLinear):

888 params_dtype=params_dtype,

889 quant_config=quant_config,

890 prefix=prefix,

935 - return_bias=return_bias)

891 + return_bias=return_bias,

892 + disable_tp=disable_tp)

893

894 def _get_shard_offset_mapping(self, loaded_shard_id: str):

895 shard_offset_mapping = {

@@ -993,10 +950,13 @@ class QKVParallelLinear(ColumnParallelLinear):

950 loaded_shard_id: Optional[str] = None):

951 if loaded_shard_id is None: # special case for certain models

952 if isinstance(param, PerTensorScaleParameter):

996 - param.load_qkv_weight(loaded_weight=loaded_weight, shard_id=0)

953 + param.load_qkv_weight(loaded_weight=loaded_weight,

954 + shard_id=0,

955 + tp_rank=self.tp_rank)

956 return

957 elif type(param) in (RowvLLMParameter, BasevLLMParameter):

999 - param.load_qkv_weight(loaded_weight=loaded_weight)

958 + param.load_qkv_weight(loaded_weight=loaded_weight,

959 + tp_rank=self.tp_rank)

960 return

961 # TODO: @dsikka - move to parameter.py

962 self._load_fused_module_from_checkpoint(param, loaded_weight)

@@ -1020,7 +980,8 @@ class QKVParallelLinear(ColumnParallelLinear):

980 num_heads=self.num_kv_head_replicas,

981 shard_id=loaded_shard_id,

982 shard_offset=shard_offset,

1023 - shard_size=shard_size)

983 + shard_size=shard_size,

984 + tp_rank=self.tp_rank)

985

986 def weight_loader(self,

987 param: Parameter,

@@ -1226,6 +1187,7 @@ class RowParallelLinear(LinearBase):

1187 prefix: The name of the layer in the state dict, including all parents

1188 (e.g. model.layers.0.down_proj)

1189 return_bias: If true, return bias together with outputs in forward pass.

1190 + disable_tp: If true, weights matrix won't be sharded through tp rank.

1191 """

1192

1193 def __init__(

@@ -1241,10 +1203,13 @@ class RowParallelLinear(LinearBase):

1203 prefix: str = "",

1204 *,

1205 return_bias: bool = True,

1206 + disable_tp: bool = False,

1207 ):

1208 # Divide the weight matrix along the first dimension.

1246 - self.tp_rank = get_tensor_model_parallel_rank()

1247 - self.tp_size = get_tensor_model_parallel_world_size()

1209 + self.tp_rank = (get_tensor_model_parallel_rank()

1210 + if not disable_tp else 0)

1211 + self.tp_size = (get_tensor_model_parallel_world_size()

1212 + if not disable_tp else 1)

1213 self.input_size_per_partition = divide(input_size, self.tp_size)

1214 self.output_size_per_partition = output_size

1215 self.output_partition_sizes = [output_size]

@@ -1255,7 +1220,8 @@ class RowParallelLinear(LinearBase):

1220 params_dtype,

1221 quant_config,

1222 prefix,

1258 - return_bias=return_bias)

1223 + return_bias=return_bias,

1224 + disable_tp=disable_tp)

1225

1226 self.input_is_parallel = input_is_parallel

1227 self.reduce_results = reduce_results

@@ -1339,10 +1305,9 @@ class RowParallelLinear(LinearBase):

1305 if self.input_is_parallel:

1306 input_parallel = input_

1307 else:

1342 - tp_rank = get_tensor_model_parallel_rank()

1308 splitted_input = split_tensor_along_last_dim(

1309 input_, num_partitions=self.tp_size)

1345 - input_parallel = splitted_input[tp_rank].contiguous()

1310 + input_parallel = splitted_input[self.tp_rank].contiguous()

1311

1312 # Matrix multiply.

1313 assert self.quant_method is not None

vllm/model_executor/model_loader/bitsandbytes_loader.py

+19 -4

@@ -69,6 +69,7 @@ class BitsAndBytesModelLoader(BaseModelLoader):

69 # Store all module names (from transformers) that support

70 # BNB quantization.

71 self.target_modules: list[str] = []

72 + self.tp_disabled_modules: list[str] = []

73 # Store the mapping of expert parameters for MoE models.

74 self.expert_params_mapping: list[tuple[str, str, int, str]] = []

75 # mapping weight names from transformers to vllm.

@@ -322,14 +323,24 @@ class BitsAndBytesModelLoader(BaseModelLoader):

323 quant_state_dict) -> Generator:

324 from bitsandbytes.functional import quantize_4bit

325

325 - tp_size = get_tensor_model_parallel_world_size()

326 - tp_rank = get_tensor_model_parallel_rank()

326 + global_tp_size = get_tensor_model_parallel_world_size()

327 + global_tp_rank = get_tensor_model_parallel_rank()

328

329 for (

330 org_weight_name,

331 mapped_weight_name,

332 weight_tensor,

333 ) in self._hf_weight_iter(hf_weights_files, use_safetensors):

334 +

335 + # override tp_size and tp_rank if the module has disabled TP

336 + if any(tp_disabled_module in mapped_weight_name

337 + for tp_disabled_module in self.tp_disabled_modules):

338 + tp_size = 1

339 + tp_rank = 0

340 + else:

341 + tp_size = global_tp_size

342 + tp_rank = global_tp_rank

343 +

344 if any(target_module in mapped_weight_name

345 for target_module in self.target_modules

346 ) and mapped_weight_name.endswith(".weight"):

@@ -418,12 +429,16 @@ class BitsAndBytesModelLoader(BaseModelLoader):

429 # Map vllm's names to transformers's names.

430 rep_name, sub_modules = modules_info

431 for sub_name in sub_modules:

421 - self.target_modules.append(

422 - name.replace(rep_name, sub_name))

432 + new_name = name.replace(rep_name, sub_name)

433 + self.target_modules.append(new_name)

434 + if module.disable_tp:

435 + self.tp_disabled_modules.append(new_name)

436 # Add original module name even if the module has stacked map,

437 # in case model has a mixture of disk-merged and disk-split

438 # weights with same last name.

439 self.target_modules.append(name)

440 + if module.disable_tp:

441 + self.tp_disabled_modules.append(name)

442 elif isinstance(module, FusedMoE) and hasattr(

443 module.quant_method, "quant_config"):

444 # TODO: support FusedMoE with prequant and 8bit.

vllm/model_executor/models/deepseek_v2.py

+3 -3

@@ -43,7 +43,6 @@ from vllm.model_executor.layers.fused_moe import FusedMoE

43 from vllm.model_executor.layers.layernorm import RMSNorm

44 from vllm.model_executor.layers.linear import (ColumnParallelLinear,

45 MergedColumnParallelLinear,

46 - MergedReplicatedLinear,

46 ReplicatedLinear,

47 RowParallelLinear)

48 from vllm.model_executor.layers.logits_processor import LogitsProcessor

@@ -435,12 +434,13 @@ class DeepseekV2MLAAttention(nn.Module):

434 self.max_position_embeddings = max_position_embeddings

435

436 if self.q_lora_rank is not None:

438 - self.fused_qkv_a_proj = MergedReplicatedLinear(

437 + self.fused_qkv_a_proj = MergedColumnParallelLinear(

438 self.hidden_size,

439 [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],

440 bias=False,

441 quant_config=quant_config,

443 - prefix=f"{prefix}.fused_qkv_a_proj")

442 + prefix=f"{prefix}.fused_qkv_a_proj",

443 + disable_tp=True)

444 else:

445 self.kv_a_proj_with_mqa = ReplicatedLinear(

446 self.hidden_size,

vllm/model_executor/models/glm4_1v.py

+50 -78

@@ -51,14 +51,10 @@ from vllm.distributed import utils as dist_utils

51 from vllm.logger import init_logger

52 from vllm.model_executor import SamplingMetadata

53 from vllm.model_executor.layers.layernorm import RMSNorm

54 -# yapf: disable

54 from vllm.model_executor.layers.linear import (ColumnParallelLinear,

55 MergedColumnParallelLinear,

57 - MergedReplicatedLinear,

56 QKVParallelLinear,

59 - ReplicatedLinear,

57 RowParallelLinear)

61 -# yapf: enable

58 from vllm.model_executor.layers.quantization import QuantizationConfig

59 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

60 from vllm.model_executor.models.module_mapping import MultiModelKeys

@@ -174,20 +170,22 @@ class Glm4vVisionMLP(nn.Module):

170 use_data_parallel: bool = False,

171 ):

172 super().__init__()

177 - cls_gate_up = (MergedReplicatedLinear

178 - if use_data_parallel else MergedColumnParallelLinear)

179 - self.gate_up_proj = cls_gate_up(input_size=in_features,

180 - output_sizes=[hidden_features] * 2,

181 - bias=bias,

182 - quant_config=quant_config,

183 - prefix=f"{prefix}.gate_up_proj")

184 - cls_down = (ReplicatedLinear

185 - if use_data_parallel else RowParallelLinear)

186 - self.down_proj = cls_down(hidden_features,

187 - in_features,

188 - bias=bias,

189 - quant_config=quant_config,

190 - prefix=f"{prefix}.down_proj")

173 + self.gate_up_proj = MergedColumnParallelLinear(

174 + input_size=in_features,

175 + output_sizes=[hidden_features] * 2,

176 + bias=bias,

177 + quant_config=quant_config,

178 + prefix=f"{prefix}.gate_up_proj",

179 + disable_tp=use_data_parallel,

180 + )

181 + self.down_proj = RowParallelLinear(

182 + hidden_features,

183 + in_features,

184 + bias=bias,

185 + quant_config=quant_config,

186 + prefix=f"{prefix}.down_proj",

187 + disable_tp=use_data_parallel,

188 + )

189 self.act_fn = SiluAndMul()

190

191 def forward(self, x: torch.Tensor):

@@ -234,48 +232,32 @@ class Glm4vVisionAttention(nn.Module):

232 # Per attention head and per partition values.

233 self.tp_size = (1 if use_data_parallel else

234 get_tensor_model_parallel_world_size())

237 - self.tp_rank = parallel_state.get_tensor_model_parallel_rank()

235 + self.tp_rank = (0 if use_data_parallel else

236 + parallel_state.get_tensor_model_parallel_rank())

237 self.hidden_size_per_attention_head = dist_utils.divide(

238 projection_size, num_heads)

239 self.num_attention_heads_per_partition = dist_utils.divide(

240 num_heads, self.tp_size)

241

243 - if use_data_parallel:

244 - self.qkv = ReplicatedLinear(

245 - input_size=embed_dim,

246 - output_size=3 * projection_size,

247 - bias=False,

248 - quant_config=quant_config,

249 - # Change qkv prefix to align with GLM-4.5V-FP8 quantization cfg

250 - prefix=f"{prefix}.qkv_proj"

251 - if quant_config else f"{prefix}.qkv",

252 - )

253 - self.proj = ReplicatedLinear(

254 - input_size=projection_size,

255 - output_size=embed_dim,

256 - quant_config=quant_config,

257 - prefix=f"{prefix}.proj",

258 - bias=False,

259 - )

260 - else:

261 - self.qkv = QKVParallelLinear(

262 - hidden_size=embed_dim,

263 - head_size=self.hidden_size_per_attention_head,

264 - total_num_heads=num_heads,

265 - total_num_kv_heads=num_heads,

266 - bias=False,

267 - quant_config=quant_config,

268 - # Change qkv prefix to align with GLM-4.5V-FP8 quantization cfg

269 - prefix=f"{prefix}.qkv_proj"

270 - if quant_config else f"{prefix}.qkv",

271 - )

272 - self.proj = RowParallelLinear(

273 - input_size=projection_size,

274 - output_size=embed_dim,

275 - quant_config=quant_config,

276 - prefix=f"{prefix}.proj",

277 - bias=False,

278 - )

242 + self.qkv = QKVParallelLinear(

243 + hidden_size=embed_dim,

244 + head_size=self.hidden_size_per_attention_head,

245 + total_num_heads=num_heads,

246 + total_num_kv_heads=num_heads,

247 + bias=False,

248 + quant_config=quant_config,

249 + # Change qkv prefix to align with GLM-4.5V-FP8 quantization cfg

250 + prefix=f"{prefix}.qkv_proj" if quant_config else f"{prefix}.qkv",

251 + disable_tp=use_data_parallel,

252 + )

253 + self.proj = RowParallelLinear(

254 + input_size=projection_size,

255 + output_size=embed_dim,

256 + quant_config=quant_config,

257 + prefix=f"{prefix}.proj",

258 + bias=False,

259 + disable_tp=use_data_parallel,

260 + )

261

262 # Detect attention implementation.

263 self.attn_backend: _Backend = get_vit_attn_backend(support_fa=True)

@@ -494,41 +476,31 @@ class Glm4vPatchMerger(nn.Module):

476 ) -> None:

477 super().__init__()

478 self.hidden_size = d_model

497 - if use_data_parallel:

498 - self.proj = ReplicatedLinear(

499 - input_size=self.hidden_size,

500 - output_size=self.hidden_size,

501 - bias=bias,

502 - quant_config=quant_config,

503 - prefix=f"{prefix}.proj",

504 - )

505 - else:

506 - self.proj = ColumnParallelLinear(

507 - self.hidden_size,

508 - self.hidden_size,

509 - bias=bias,

510 - gather_output=True,

511 - quant_config=quant_config,

512 - prefix=f"{prefix}.proj",

513 - )

479 + self.proj = ColumnParallelLinear(

480 + self.hidden_size,

481 + self.hidden_size,

482 + bias=bias,

483 + gather_output=True,

484 + quant_config=quant_config,

485 + prefix=f"{prefix}.proj",

486 + disable_tp=use_data_parallel,

487 + )

488 self.post_projection_norm = nn.LayerNorm(self.hidden_size)

515 - cls_gate_up = (MergedReplicatedLinear

516 - if use_data_parallel else MergedColumnParallelLinear)

517 - self.gate_up_proj = cls_gate_up(

489 + self.gate_up_proj = MergedColumnParallelLinear(

490 input_size=self.hidden_size,

491 output_sizes=[context_dim] * 2,

492 bias=bias,

493 quant_config=quant_config,

494 prefix=f"{prefix}.gate_up_proj",

495 + disable_tp=use_data_parallel,

496 )

524 - cls_down = (ReplicatedLinear

525 - if use_data_parallel else RowParallelLinear)

526 - self.down_proj = cls_down(

497 + self.down_proj = RowParallelLinear(

498 context_dim,

499 self.hidden_size,

500 bias=bias,

501 quant_config=quant_config,

502 prefix=f"{prefix}.down_proj",

503 + disable_tp=use_data_parallel,

504 )

505 self.act_fn = SiluAndMul()

506 self.extra_activation_func = nn.GELU()

vllm/model_executor/models/qwen2_5_vl.py

+25 -37

@@ -48,7 +48,6 @@ from vllm.model_executor.layers.layernorm import RMSNorm

48 # yapf: disable

49 from vllm.model_executor.layers.linear import (ColumnParallelLinear,

50 MergedColumnParallelLinear,

51 - MergedReplicatedLinear,

51 QKVParallelLinear,

52 ReplicatedLinear,

53 RowParallelLinear)

@@ -178,22 +177,20 @@ class Qwen2_5_VisionMLP(nn.Module):

177 prefix: str = "",

178 use_data_parallel: bool = False):

179 super().__init__()

181 - cls_gate_up_proj = (MergedReplicatedLinear if use_data_parallel else

182 - MergedColumnParallelLinear)

183 - self.gate_up_proj = cls_gate_up_proj(

180 + self.gate_up_proj = MergedColumnParallelLinear(

181 input_size=in_features,

182 output_sizes=[hidden_features] * 2, # [gate_proj, up_proj]

183 bias=bias,

184 quant_config=quant_config,

188 - prefix=f"{prefix}.gate_up_proj")

189 -

190 - cls_down_proj = (ReplicatedLinear

191 - if use_data_parallel else RowParallelLinear)

192 - self.down_proj = cls_down_proj(hidden_features,

193 - in_features,

194 - bias=bias,

195 - quant_config=quant_config,

196 - prefix=f"{prefix}.down_proj")

185 + prefix=f"{prefix}.gate_up_proj",

186 + disable_tp=use_data_parallel)

187 +

188 + self.down_proj = RowParallelLinear(hidden_features,

189 + in_features,

190 + bias=bias,

191 + quant_config=quant_config,

192 + prefix=f"{prefix}.down_proj",

193 + disable_tp=use_data_parallel)

194 self.act_fn = act_fn

195

196 def forward(self, x: torch.Tensor):

@@ -243,30 +240,21 @@ class Qwen2_5_VisionAttention(nn.Module):

240 self.num_attention_heads_per_partition = dist_utils.divide(

241 num_heads, self.tp_size)

242

246 - if use_data_parallel:

247 - self.qkv = ReplicatedLinear(embed_dim,

248 - self.hidden_size_per_attention_head *

249 - 3 * num_heads,

250 - bias=True,

251 - quant_config=quant_config,

252 - prefix=f"{prefix}.qkv")

253 -

254 - else:

255 - self.qkv = QKVParallelLinear(

256 - hidden_size=embed_dim,

257 - head_size=self.hidden_size_per_attention_head,

258 - total_num_heads=num_heads,

259 - total_num_kv_heads=num_heads,

260 - bias=True,

261 - quant_config=quant_config,

262 - prefix=f"{prefix}.qkv")

263 -

264 - cls_proj = (ReplicatedLinear

265 - if use_data_parallel else RowParallelLinear)

266 - self.proj = cls_proj(input_size=projection_size,

267 - output_size=embed_dim,

268 - quant_config=quant_config,

269 - prefix=f"{prefix}.proj")

243 + self.qkv = QKVParallelLinear(

244 + hidden_size=embed_dim,

245 + head_size=self.hidden_size_per_attention_head,

246 + total_num_heads=num_heads,

247 + total_num_kv_heads=num_heads,

248 + bias=True,

249 + quant_config=quant_config,

250 + prefix=f"{prefix}.qkv",

251 + disable_tp=use_data_parallel)

252 +

253 + self.proj = RowParallelLinear(input_size=projection_size,

254 + output_size=embed_dim,

255 + quant_config=quant_config,

256 + prefix=f"{prefix}.proj",

257 + disable_tp=use_data_parallel)

258

259 # Detect attention implementation.

260 self.attn_backend: _Backend = get_vit_attn_backend(support_fa=True)

vllm/model_executor/models/step3_vl.py

+27 -44

@@ -21,7 +21,6 @@ from vllm.distributed import get_tensor_model_parallel_world_size

21 from vllm.model_executor.layers.activation import get_act_fn

22 from vllm.model_executor.layers.linear import (ColumnParallelLinear,

23 QKVParallelLinear,

24 - ReplicatedLinear,

24 RowParallelLinear)

25 from vllm.model_executor.layers.quantization import QuantizationConfig

26 from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler

@@ -667,35 +666,21 @@ class Step3VisionAttention(nn.Module):

666

667 self.q_size = self.num_heads * self.head_dim

668

670 - if use_data_parallel:

671 - self.qkv_proj = ReplicatedLinear(

672 - self.embed_dim,

673 - 3 * self.q_size,

674 - bias=True,

675 - quant_config=quant_config,

676 - prefix=prefix,

677 - )

678 - self.out_proj = ReplicatedLinear(

679 - self.total_num_heads * self.head_dim,

680 - self.embed_dim,

681 - bias=True,

682 - quant_config=quant_config,

683 - prefix=prefix,

684 - )

685 - else:

686 - self.qkv_proj = QKVParallelLinear(

687 - self.embed_dim,

688 - self.head_dim,

689 - self.total_num_heads,

690 - bias=True,

691 - quant_config=quant_config,

692 - prefix=prefix,

693 - )

694 - self.out_proj = RowParallelLinear(self.embed_dim,

695 - self.embed_dim,

696 - bias=True,

697 - quant_config=quant_config,

698 - prefix=prefix)

669 + self.qkv_proj = QKVParallelLinear(

670 + self.embed_dim,

671 + self.head_dim,

672 + self.total_num_heads,

673 + bias=True,

674 + quant_config=quant_config,

675 + prefix=f"{prefix}.qkv_proj",

676 + disable_tp=use_data_parallel,

677 + )

678 + self.out_proj = RowParallelLinear(self.embed_dim,

679 + self.embed_dim,

680 + bias=True,

681 + quant_config=quant_config,

682 + prefix=f"{prefix}.out_proj",

683 + disable_tp=use_data_parallel)

684

685 def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):

686 return tensor.view(bsz, seq_len, self.num_heads,

@@ -740,20 +725,18 @@ class Step3VisionMLP(nn.Module):

725 super().__init__()

726 self.config = config

727 self.activation_fn = get_act_fn(config.hidden_act)

743 - cls_fc1 = (ReplicatedLinear

744 - if use_data_parallel else ColumnParallelLinear)

745 - self.fc1 = cls_fc1(config.hidden_size,

746 - config.intermediate_size,

747 - bias=True,

748 - quant_config=quant_config,

749 - prefix=prefix)

750 - cls_fc2 = (ReplicatedLinear

751 - if use_data_parallel else RowParallelLinear)

752 - self.fc2 = cls_fc2(config.intermediate_size,

753 - config.hidden_size,

754 - bias=True,

755 - quant_config=quant_config,

756 - prefix=prefix)

728 + self.fc1 = ColumnParallelLinear(config.hidden_size,

729 + config.intermediate_size,

730 + bias=True,

731 + quant_config=quant_config,

732 + prefix=f"{prefix}.fc1",

733 + disable_tp=use_data_parallel)

734 + self.fc2 = RowParallelLinear(config.intermediate_size,

735 + config.hidden_size,

736 + bias=True,

737 + quant_config=quant_config,

738 + prefix=f"{prefix}.fc2",

739 + disable_tp=use_data_parallel)

740

741 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:

742 hidden_states, _ = self.fc1(hidden_states)

vllm/model_executor/parameter.py

+11 -11

@@ -57,6 +57,8 @@ class BasevLLMParameter(Parameter):

57 weight_loader = _make_synced_weight_loader(weight_loader)

59 self._weight_loader = weight_loader

60 + self.tp_rank = get_tensor_model_parallel_rank()

61 + self.tp_size = get_tensor_model_parallel_world_size()

63 @property

64 def weight_loader(self):

@@ -116,10 +118,10 @@ class _ColumnvLLMParameter(BasevLLMParameter):

118 return self._output_dim

119

120 def load_column_parallel_weight(self, loaded_weight: torch.Tensor):

119 - tp_rank = get_tensor_model_parallel_rank()

121 shard_size = self.data.shape[self.output_dim]

122 loaded_weight = loaded_weight.narrow(self.output_dim,

122 - tp_rank * shard_size, shard_size)

123 + self.tp_rank * shard_size,

124 + shard_size)

125 assert self.data.shape == loaded_weight.shape

126 self.data.copy_(loaded_weight)

127

@@ -127,6 +129,7 @@ class _ColumnvLLMParameter(BasevLLMParameter):

129

130 shard_offset = kwargs.get("shard_offset")

131 shard_size = kwargs.get("shard_size")

132 +

133 # TODO: move these to PackedColumnParameter and PackedvLLMParameter

134 if isinstance(

135 self,

@@ -137,11 +140,11 @@ class _ColumnvLLMParameter(BasevLLMParameter):

140

141 param_data = self.data

142

140 - tp_rank = get_tensor_model_parallel_rank()

143 param_data = param_data.narrow(self.output_dim, shard_offset,

144 shard_size)

145 loaded_weight = loaded_weight.narrow(self.output_dim,

144 - tp_rank * shard_size, shard_size)

146 + self.tp_rank * shard_size,

147 + shard_size)

148 assert param_data.shape == loaded_weight.shape

149 param_data.copy_(loaded_weight)

150

@@ -161,8 +164,8 @@ class _ColumnvLLMParameter(BasevLLMParameter):

164 shard_offset=shard_offset, shard_size=shard_size)

165

166 param_data = self.data

164 - tp_rank = get_tensor_model_parallel_rank()

165 - shard_id = tp_rank if shard_id == "q" else tp_rank // num_heads

167 + shard_id = (self.tp_rank if shard_id == "q" else self.tp_rank //

168 + num_heads)

169 param_data = param_data.narrow(self.output_dim, shard_offset,

170 shard_size)

171 loaded_weight = loaded_weight.narrow(self.output_dim,

@@ -189,10 +192,10 @@ class RowvLLMParameter(BasevLLMParameter):

192 return self._input_dim

193

194 def load_row_parallel_weight(self, loaded_weight: torch.Tensor):

192 - tp_rank = get_tensor_model_parallel_rank()

195 shard_size = self.data.shape[self.input_dim]

196 loaded_weight = loaded_weight.narrow(self.input_dim,

195 - tp_rank * shard_size, shard_size)

197 + self.tp_rank * shard_size,

198 + shard_size)

199

200 if len(loaded_weight.shape) == 0:

201 loaded_weight = loaded_weight.reshape(1)

@@ -414,9 +417,6 @@ class SharedWeightParameter(BasevLLMParameter):

417 "weight_loader": self._fake_weight_loader

418 }

419

417 - self.tp_rank = get_tensor_model_parallel_rank()

418 - self.tp_size = get_tensor_model_parallel_world_size()

419 -

420 if self.tp_size > 1:

421 raise NotImplementedError(f"{self.__class__.__name__} does not "

422 "currently support tensor parallelism")