[NemotronH] Load shared embed_tokens/lm_head in MTP draft weights (#30456)

2026-07-16 17:20:38 +00:00
parent 3bc7c54111
commit 1f34911de7
2 changed files with 63 additions and 1 deletions
@@ -1125,7 +1125,17 @@ class NemotronHForCausalLM(nn.Module):
            name = replace_prefix(name, self.remap_prefix)
            name = replace_substrings(name, self.remap_substr)
            if is_mtp:
-                if "mtp" not in name:
+                # Keep the MTP draft layers (mtp.layers.*) and the shared
+                # embed_tokens / lm_head weights. The remap above already
+                # rewrote "embeddings" -> "embed_tokens", so the shared draft
+                # weights must be whitelisted by their post-remap names; without
+                # this the draft loads no embedding / head and accepts zero
+                # draft tokens (accept rate 0.00), see issue #21138.
+                if (
+                    "mtp" not in name
+                    and "embed_tokens" not in name
+                    and "lm_head" not in name
+                ):
                    continue

                name = name.replace("mtp.layers.", "model.layers.")