fix(step3p5-mtp): honor exclude_modules for the MTP head via prefix (#48883)
Signed-off-by: Chan Nguyen <channguyen@nvidia.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Chanh Nguyen committed
Jul 29, 2026 at 09:58 UTC
f98061ce6cb0326978363fb6c0489fe75a9aac5d
1 file changed
+11
-2
vllm/model_executor/models/step3p5_mtp.py
+11
-2
@@ -32,11 +32,18 @@ class SharedHead(nn.Module):
32
self,
33
config: PretrainedConfig,
34
quant_config: QuantizationConfig | None = None,
35
+ prefix: str = "",
36
) -> None:
37
super().__init__()
38
self.norm = GemmaRMSNorm(config.hidden_size, config.rms_norm_eps)
39
+ # Give the head its prefix so the quant config's exclude_modules matcher
40
+ # can skip it; without one it defaults to "" and never matches, so a
41
+ # checkpoint-excluded (BF16) MTP head gets quantized -> load crash.
42
self.head = ParallelLMHead(
39
- config.vocab_size, config.hidden_size, quant_config=quant_config
43
+ config.vocab_size,
44
+ config.hidden_size,
45
+ quant_config=quant_config,
46
+ prefix=f"{prefix}.head",
47
)
48
49
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
@@ -55,7 +62,9 @@ class Step3p5AMultiTokenPredictorLayer(nn.Module):
62
self.enorm = GemmaRMSNorm(config.hidden_size, config.rms_norm_eps)
63
self.hnorm = GemmaRMSNorm(config.hidden_size, config.rms_norm_eps)
64
self.eh_proj = nn.Linear(config.hidden_size * 2, config.hidden_size, bias=False)
58
- self.shared_head = SharedHead(config=config, quant_config=quant_config)
65
+ self.shared_head = SharedHead(
66
+ config=config, quant_config=quant_config, prefix=f"{prefix}.shared_head"
67
+ )
68
self.mtp_block = Step3p5DecoderLayer(
69
vllm_config,
70
prefix=f"{prefix}.mtp_block",