From 68c5ef98a8983f70bf95b0695eb0e9f689af8422 Mon Sep 17 00:00:00 2001 From: Aaron Kang Date: Mon, 13 Jul 2026 14:13:34 -0400 Subject: [PATCH] [Bugfix] Pass quant_config to DiffusionGemma's ParallelLMHead DiffusionGemmaForConditionalGeneration builds its ParallelLMHead without quant_config or prefix, so the head is always constructed as an unquantized layer. Checkpoints that quantize lm_head fail to load with: ValueError: There is no module or parameter named 'lm_head.weight_scale' in DiffusionGemmaForConditionalGeneration Every other quantization-aware model passes both arguments (see llama.py), and this file already threads quant_config and maybe_prefix() into the vision tower and language backbone a few lines above. Read vllm_config.quant_config directly rather than the local quant_config name, which is only bound inside the 'if vision_config is not None' branch and would raise NameError on a text-only config. No behaviour change for existing checkpoints: quant_config is None for unquantized models, and the published quantized checkpoints all list lm_head in their ignore list. Signed-off-by: Aaron Kang --- vllm/model_executor/models/diffusion_gemma.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py index 11a10131df17..8467f89f775e 100644 --- a/vllm/model_executor/models/diffusion_gemma.py +++ b/vllm/model_executor/models/diffusion_gemma.py @@ -239,6 +239,8 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.lm_head = ParallelLMHead( num_embeddings=text_config.vocab_size, embedding_dim=text_config.hidden_size, + quant_config=vllm_config.quant_config, + prefix=maybe_prefix(prefix, "lm_head"), ) if text_config.tie_word_embeddings: