diff --git a/python/sglang/srt/models/nemotron_h.py b/python/sglang/srt/models/nemotron_h.py index 6a4752630..58253830b 100644 --- a/python/sglang/srt/models/nemotron_h.py +++ b/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.") diff --git a/test/registered/unit/models/test_nemotron_h_weight_loading.py b/test/registered/unit/models/test_nemotron_h_weight_loading.py index 6a1221ff3..ed4c21d23 100644 --- a/test/registered/unit/models/test_nemotron_h_weight_loading.py +++ b/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()