From 29f8b8c1744df4b5907f64ce8565a363611f501b Mon Sep 17 00:00:00 2001 From: biondizzle Date: Fri, 15 May 2026 02:17:16 +0000 Subject: [PATCH] fix: load lm_head.weight in outer model before forwarding to inner lm_head lives on DeepseekV4ForCausalLM, not DeepseekV4Model. The inner load_weights silently drops it (not in params_dict). Extract it in the outer loader, load it directly, then forward the rest to the inner model. --- vllm/patches/deepseek_v4.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/vllm/patches/deepseek_v4.py b/vllm/patches/deepseek_v4.py index a527c357..2fbc4fc5 100644 --- a/vllm/patches/deepseek_v4.py +++ b/vllm/patches/deepseek_v4.py @@ -2205,11 +2205,24 @@ class DeepseekV4ForCausalLM(nn.Module): return getattr(self.model, "_mtp_hidden_buffer", None) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + # lm_head lives on this outer model, not on the inner DeepseekV4Model. + # The inner load_weights silently drops lm_head.weight via + # "if name not in params_dict: continue". Extract it here. + rest = [] + for name, loaded_weight in weights: + if name == "lm_head.weight" or name.endswith(".lm_head.weight"): + param = self.lm_head.weight + weight_loader = getattr(param, "weight_loader", + default_weight_loader) + weight_loader(param, loaded_weight) + else: + rest.append((name, loaded_weight)) # Use the model-level loader which handles NVFP4 expert mapping, # uint8→bf16 unpacking for MergedColumnParallelLinear, and # bf16→NVFP4 quantization for unquantized layers. # AutoWeightsLoader bypasses this logic and would break NVFP4 loading. - loaded_params = self.model.load_weights(weights) + loaded_params = self.model.load_weights(rest) + loaded_params.add("lm_head.weight") self.model.finalize_mega_moe_weights() self.model._convert_nvfp4_post_load() if int(os.environ.get('NVFP4_DEBUG', '0')):