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 15, 14:45 UTC
[Bugfix] Handle `process_weights_after_loading` for `QKVCrossParallelLinear` (#15328)
Signed-off-by: Isotr0py <2037008807@qq.com>
Isotr0py committed
Apr 9, 2025 at 01:02 UTC
40b4284fe3fb70722aeb5d187a03a6b7a3a98565
3 files changed
+33
-6
vllm/model_executor/layers/linear.py
+21
-6
@@ -1353,6 +1353,7 @@ class QKVCrossParallelLinear(LinearBase):
1353
prefix=f"{prefix}.kv_proj_encoder")
1354
1355
# `kv_proj_encoder.num_kv_heads` accounts for sharding with tp>1.
1356
+ self.q_size = self.q_proj_decoder.output_size_per_partition
1357
self.kv_size = self.kv_proj_encoder.num_kv_heads * head_size
1358
1359
if bias:
@@ -1364,20 +1365,31 @@ class QKVCrossParallelLinear(LinearBase):
1365
else:
1366
self.bias = None
1367
1368
+ def process_weights_after_loading(self):
1369
+ for layer in self.proj.values():
1370
+ if self.quant_method is not None:
1371
+ self.quant_method.process_weights_after_loading(layer)
1372
+
1373
@property
1374
def q_proj_decoder(self) -> ColumnParallelLinear:
1375
layer = self.proj["q_proj_decoder"]
1376
for name, param in self.named_parameters():
1371
- target_param = getattr(layer, name)
1372
- self.sync_weight_attrs(param, target_param, mode="q_proj_decoder")
1377
+ target_param = getattr(layer, name, None)
1378
+ if target_param is not None:
1379
+ self.sync_weight_attrs(param,
1380
+ target_param,
1381
+ mode="q_proj_decoder")
1382
return layer
1383
1384
@property
1385
def kv_proj_encoder(self) -> QKVParallelLinear:
1386
layer = self.proj["kv_proj_encoder"]
1387
for name, param in self.named_parameters():
1379
- target_param = getattr(layer, name)
1380
- self.sync_weight_attrs(param, target_param, mode="kv_proj_encoder")
1388
+ target_param = getattr(layer, name, None)
1389
+ if target_param is not None:
1390
+ self.sync_weight_attrs(param,
1391
+ target_param,
1392
+ mode="kv_proj_encoder")
1393
return layer
1394
1395
def sync_weight_attrs(
@@ -1466,11 +1478,14 @@ class QKVCrossParallelLinear(LinearBase):
1478
if loaded_shard_id == "q" else self.kv_proj_encoder)
1479
target_param = self.select_proj_params(layer, param)
1480
shard_id_args = (loaded_shard_id, ) if loaded_shard_id != "q" else ()
1469
- layer.weight_loader(target_param, loaded_weight, *shard_id_args)
1481
+ if self.quant_method.__class__.__name__ in WEIGHT_LOADER_V2_SUPPORTED:
1482
+ layer.weight_loader_v2(target_param, loaded_weight, *shard_id_args)
1483
+ else:
1484
+ layer.weight_loader(target_param, loaded_weight, *shard_id_args)
1485
1486
def extra_repr(self) -> str:
1487
s = f"in_features={self.input_size}"
1473
- s += f", q_size={self.q_proj_decoder.output_size_per_partition}"
1488
+ s += f", q_size={self.q_size}"
1489
s += f", kv_size={self.kv_size}"
1490
s += f", bias={self.bias is not None}"
1491
s += f", tp_size={get_tensor_model_parallel_world_size()}"
vllm/model_executor/layers/quantization/fp8.py
+3
@@ -254,6 +254,7 @@ class Fp8LinearMethod(LinearMethodBase):
254
weight_loader=weight_loader,
255
)
256
scale[:] = torch.finfo(torch.float32).min
257
+ set_weight_attrs(scale, {"scale_type": "weight_scale"})
258
layer.register_parameter("weight_scale", scale)
259
else:
260
assert self.quant_config.activation_scheme == "dynamic"
@@ -268,6 +269,7 @@ class Fp8LinearMethod(LinearMethodBase):
269
weight_loader=weight_loader,
270
)
271
scale[:] = torch.finfo(torch.float32).min
272
+ set_weight_attrs(scale, {"scale_type": "weight_scale"})
273
# The weight_scale_inv name is intentional for deepseekv3
274
layer.register_parameter("weight_scale_inv", scale)
275
@@ -278,6 +280,7 @@ class Fp8LinearMethod(LinearMethodBase):
280
weight_loader=weight_loader)
281
282
scale[:] = torch.finfo(torch.float32).min
283
+ set_weight_attrs(scale, {"scale_type": "input_scale"})
284
layer.register_parameter("input_scale", scale)
285
else:
286
layer.register_parameter("input_scale", None)
vllm/model_executor/model_loader/loader.py
+9
@@ -33,11 +33,15 @@ from vllm.distributed import (get_tensor_model_parallel_rank,
33
get_tensor_model_parallel_world_size)
34
from vllm.envs import VLLM_USE_MODELSCOPE
35
from vllm.logger import init_logger
36
+# yapf conflicts with isort for this block
37
+# yapf: disable
38
from vllm.model_executor.layers.linear import (LinearBase,
39
MergedColumnParallelLinear,
40
+ QKVCrossParallelLinear,
41
QKVParallelLinear,
42
ReplicatedLinear,
43
RowParallelLinear)
44
+# yapf: enable
45
from vllm.model_executor.layers.quantization.base_config import (
46
QuantizeMethodBase)
47
from vllm.model_executor.model_loader.tensorizer import (
@@ -160,6 +164,11 @@ def _initialize_model(
164
def _process_weights_after_loading(model: nn.Module, model_config: ModelConfig,
165
target_device: torch.device) -> None:
166
for _, module in model.named_modules():
167
+ if isinstance(module, QKVCrossParallelLinear):
168
+ # NOTE(Isotr0py): special case for cross QKV layer because
169
+ # q and kv proj aren't registered as submodules intentionally
170
+ module.process_weights_after_loading()
171
+ continue
172
quant_method = getattr(module, "quant_method", None)
173
if isinstance(quant_method, QuantizeMethodBase):
174
# When quant methods need to process weights after loading