[Bugfix] Add missing packed_modules_mapping to DeepseekV2ForCausalLM (#22352)

Signed-off-by: Felix Marty <Felix.Marty@amd.com>
This commit is contained in:
fxmarty-amd 2025-08-07 15:30:48 +02:00 committed by GitHub
parent 766bc8162c
commit 7e0b121812
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -726,6 +726,9 @@ class DeepseekV2Model(nn.Module):
class DeepseekV2ForCausalLM(nn.Module, SupportsPP, MixtureOfExperts):
packed_modules_mapping = {
"gate_up_proj": ["gate_proj", "up_proj"],
}
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
super().__init__()
@ -733,6 +736,19 @@ class DeepseekV2ForCausalLM(nn.Module, SupportsPP, MixtureOfExperts):
quant_config = vllm_config.quant_config
self.config = config
self.quant_config = quant_config
# `packed_modules_mapping` needs to be modified before
# initializing DeepseekV2Model, as it is passed inplace to
# quantization config init and may be used to select the
# quant_method for relevant layers during initialization.
self.fuse_qkv_a_proj = hasattr(
config, "q_lora_rank") and config.q_lora_rank is not None
if self.fuse_qkv_a_proj:
self.packed_modules_mapping["fused_qkv_a_proj"] = [
"q_a_proj",
"kv_a_proj_with_mqa",
]
self.model = DeepseekV2Model(vllm_config=vllm_config,
prefix=maybe_prefix(prefix, "model"))
if get_pp_group().is_last_rank: