diff --git a/vllm/models/deepseek_v4/nvidia/dspark.py b/vllm/models/deepseek_v4/nvidia/dspark.py index e4e258372a2..bde088edae1 100644 --- a/vllm/models/deepseek_v4/nvidia/dspark.py +++ b/vllm/models/deepseek_v4/nvidia/dspark.py @@ -43,6 +43,7 @@ from vllm.model_executor.models.utils import maybe_prefix from .model import ( DeepseekV4DecoderLayer, + DeepseekV4Model, make_deepseek_v4_expert_params_mapping, ) @@ -277,6 +278,11 @@ class DSparkDeepseekV4ForCausalLM(nn.Module): assert vllm_config.speculative_config is not None self.draft_model_config = vllm_config.speculative_config.draft_model_config self.config = self.draft_model_config.hf_config + self.quant_config = vllm_config.quant_config + self.pad_shared_expert = ( + getattr(self.quant_config, "weight_block_size", None) is not None + and not vllm_config.parallel_config.use_sequence_parallel_moe + ) self.model = DSparkDeepseekV4Model( vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") ) @@ -396,6 +402,12 @@ class DSparkDeepseekV4ForCausalLM(nn.Module): else ".weight_scale_inv" ) name = name.removesuffix(".scale") + suffix + if ".shared_experts.w2" in name: + name = name.replace(".shared_experts.w2", ".shared_experts.down_proj") + if self.pad_shared_expert and ".shared_experts." in name: + loaded_weight = DeepseekV4Model._pad_shared_expert_weight( + self.quant_config, name, loaded_weight + ) # E8M0 expert scales: keep raw exponent bytes. if ".experts." in name: @@ -440,10 +452,6 @@ class DSparkDeepseekV4ForCausalLM(nn.Module): params_dict[name][: narrow.shape[0]].copy_(narrow) loaded_params.add(name) continue - if ".shared_experts.w2" in name: - name = name.replace( - ".shared_experts.w2", ".shared_experts.down_proj" - ) if name.endswith(".ffn.gate.bias"): name = name.replace( ".ffn.gate.bias", ".ffn.gate.e_score_correction_bias" diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index ddc2fe0f4bc..96a3c77074a 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -1181,7 +1181,9 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): for name, loaded_weight in weights: if pad_shared_expert and ".shared_experts." in name: - loaded_weight = self._pad_shared_expert_weight(name, loaded_weight) + loaded_weight = self._pad_shared_expert_weight( + self.quant_config, name, loaded_weight + ) for param_name, weight_name, shard_id in stacked_params_mapping: # Skip non-stacked layers and experts (experts handled below). if ".experts." in name: @@ -1256,15 +1258,18 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): return loaded_params + @staticmethod def _pad_shared_expert_weight( - self, name: str, loaded_weight: torch.Tensor + quant_config: QuantizationConfig | None, + name: str, + loaded_weight: torch.Tensor, ) -> torch.Tensor: """Zero-pad a block-FP8 shared-expert weight/scale on its intermediate axis so the standard TP loaders split it into even, block-aligned shards (trailing ranks get the zero pad). gate (w1)/up (w3) [I, H] pad dim 0; down (w2 -> down_proj) [H, I] pads dim 1. """ - block_size = getattr(self.quant_config, "weight_block_size", None) + block_size = getattr(quant_config, "weight_block_size", None) assert block_size is not None # Round the intermediate axis up to a whole number of TP shards. The axis # is in elements for weights (step = block) and in blocks for scales. diff --git a/vllm/models/deepseek_v4/nvidia/mtp.py b/vllm/models/deepseek_v4/nvidia/mtp.py index 8aa3d2c9a29..ee036967a7e 100644 --- a/vllm/models/deepseek_v4/nvidia/mtp.py +++ b/vllm/models/deepseek_v4/nvidia/mtp.py @@ -52,6 +52,7 @@ from vllm.sequence import IntermediateTensors from .model import ( DeepseekV4DecoderLayer, + DeepseekV4Model, make_deepseek_v4_expert_params_mapping, ) @@ -265,6 +266,10 @@ class DeepSeekV4MTP(nn.Module): super().__init__() self.config = vllm_config.model_config.hf_config self.quant_config = vllm_config.quant_config + self.pad_shared_expert = ( + getattr(self.quant_config, "weight_block_size", None) is not None + and not vllm_config.parallel_config.use_sequence_parallel_moe + ) self.model = DeepSeekV4MultiTokenPredictor( vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") ) @@ -387,6 +392,12 @@ class DeepSeekV4MTP(nn.Module): else ".weight_scale_inv" ) name = name.removesuffix(".scale") + suffix + if ".shared_experts.w2" in name: + name = name.replace(".shared_experts.w2", ".shared_experts.down_proj") + if self.pad_shared_expert and ".shared_experts." in name: + loaded_weight = DeepseekV4Model._pad_shared_expert_weight( + self.quant_config, name, loaded_weight + ) for param_name, weight_name, shard_id in stacked_params_mapping: # Skip non-stacked layers and experts (experts handled below). if ".experts." in name: @@ -442,10 +453,6 @@ class DeepSeekV4MTP(nn.Module): loaded_params.add(name) continue else: - if ".shared_experts.w2" in name: - name = name.replace( - ".shared_experts.w2", ".shared_experts.down_proj" - ) if name.endswith(".ffn.gate.bias"): # ``e_score_correction_bias`` lives on the gate # under a different attribute name.