[Bugfix] Tie lm_head.weight for Nemotron Parse when checkpoint omits it (#53020)
Signed-off-by: Anish Kumar <aniskumar@nvidia.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
aniskumar-nv committed
Oct 1, 2026 at 09:26 UTC
83cadd65d9ec9e7c2839339a8f70f201246b540b
2 files changed
+19
-1
tests/models/multimodal/generation/test_nemotron_parse.py
+10
-1
@@ -104,7 +104,16 @@ def run_test(
104
)
105
106
107
-@pytest.mark.parametrize("model", ["nvidia/NVIDIA-Nemotron-Parse-v1.2"])
107
+@pytest.mark.parametrize(
108
+ "model",
109
+ [
110
+ "nvidia/NVIDIA-Nemotron-Parse-v1.2",
111
+ # v2.0's checkpoint ties lm_head to the decoder's input embeddings
112
+ # instead of shipping a separate lm_head.weight tensor; this
113
+ # exercises that path (see tie_weights() in nemotron_parse.py).
114
+ "nvidia/NVIDIA-Nemotron-Parse-2.0",
115
+ ],
116
+)
117
@pytest.mark.parametrize("dtype", ["bfloat16"])
118
@pytest.mark.parametrize("num_logprobs", [5])
119
def test_models(
vllm/model_executor/models/nemotron_parse.py
+9
@@ -568,6 +568,15 @@ class NemotronParseForConditionalGeneration(nn.Module, SupportsMultiModal):
568
self.lm_head = ParallelLMHead(
569
config.decoder.vocab_size, config.decoder.d_model, quant_config=quant_config
570
)
571
+ # Some checkpoints (e.g. compact exports) tie the output head to the
572
+ # decoder's input embeddings instead of materializing a separate
573
+ # lm_head.weight tensor.
574
+ tie_word_embeddings = bool(
575
+ getattr(config, "tie_word_embeddings", False)
576
+ or getattr(config.decoder, "tie_word_embeddings", False)
577
+ )
578
+ if tie_word_embeddings:
579
+ self.lm_head = self.lm_head.tie_weights(self.decoder.embed_tokens)
580
self.logits_processor = LogitsProcessor(
581
self.vocab_size, config.decoder.vocab_size
582
)