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

python/sglang/srt/models/nemotron_h.py

@@ -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.")

test/registered/unit/models/test_nemotron_h_weight_loading.py

@@ -32,6 +32,17 @@ class _FakeParam:
         self.loaded = (param, loaded_weight, name, shard_id, expert_id)
 
 
+class _RecordingParam:
+    """Param whose weight_loader matches the plain (non-expert) load path:
+    weight_loader(param, loaded_weight)."""
+
+    def __init__(self):
+        self.loaded_weight = None
+
+    def weight_loader(self, param, loaded_weight):
+        self.loaded_weight = loaded_weight
+
+
 class TestNemotronHWeightLoading(unittest.TestCase):
     def _make_minimal_model(self, named_parameters=()):
         model = object.__new__(NemotronHForCausalLM)
@@ -82,6 +93,47 @@ class TestNemotronHWeightLoading(unittest.TestCase):
             ),
         )
 
+    def test_mtp_keeps_shared_embed_tokens_and_lm_head(self):
+        """MTP draft load must keep the shared embed_tokens + lm_head, not only
+        mtp.layers.*. Regression for issue #21138: dropping the shared embedding
+        / head makes the draft accept zero tokens (accept rate 0.00). The remap
+        rewrites "embeddings" -> "embed_tokens" before the MTP filter, so the
+        whitelist must match the post-remap names."""
+        embed = _RecordingParam()
+        head = _RecordingParam()
+        mtp_layer = _RecordingParam()
+        skipped = _RecordingParam()
+        model = self._make_minimal_model(
+            [
+                ("model.embed_tokens.weight", embed),
+                ("lm_head.weight", head),
+                ("model.layers.0.norm.weight", mtp_layer),
+                ("model.layers.5.norm.weight", skipped),
+            ]
+        )
+        # Production remap: backbone -> model, embeddings -> embed_tokens.
+        model.remap_prefix = {"backbone": "model"}
+        model.remap_substr = {"embeddings": "embed_tokens"}
+
+        w_embed, w_head, w_mtp, w_skip = (torch.ones(1) for _ in range(4))
+        weights = [
+            ("backbone.embeddings.weight", w_embed),  # -> model.embed_tokens.weight
+            ("lm_head.weight", w_head),  # -> lm_head.weight
+            ("mtp.layers.0.norm.weight", w_mtp),  # -> model.layers.0.norm.weight
+            ("backbone.layers.5.norm.weight", w_skip),  # non-MTP target -> skipped
+        ]
+
+        model.load_weights(weights, is_mtp=True)
+
+        self.assertIs(embed.loaded_weight, w_embed, "shared embed_tokens dropped")
+        self.assertIs(head.loaded_weight, w_head, "shared lm_head dropped")
+        self.assertIs(
+            mtp_layer.loaded_weight, w_mtp, "mtp.layers.* not remapped/loaded"
+        )
+        self.assertIsNone(
+            skipped.loaded_weight, "non-MTP target weight should be skipped"
+        )
+
 
 if __name__ == "__main__":
     unittest.main()