diff --git a/python/sglang/srt/models/qwen3_5_mtp.py b/python/sglang/srt/models/qwen3_5_mtp.py index ae1d0b48e..2012068db 100644 --- a/python/sglang/srt/models/qwen3_5_mtp.py +++ b/python/sglang/srt/models/qwen3_5_mtp.py @@ -323,6 +323,23 @@ class Qwen3_5ForCausalLMMTP(nn.Module): loaded_params: set[str] = set() for name, loaded_weight in weights: + # The last-stage MTP draft cannot share the target embedding on PP0. + # Load the checkpoint embedding into its retained local copy instead + # of leaving the torch.empty() allocation uninitialized. + if name in ( + "model.embed_tokens.weight", + "model.language_model.embed_tokens.weight", + ): + param_name = "model.embed_tokens.weight" + if param_name in params_dict: + param = params_dict[param_name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + loaded_params.add(param_name) + continue + if "rotary_emb.inv_freq" in name: continue diff --git a/test/registered/unit/models/test_qwen3_5_pipeline_parallel.py b/test/registered/unit/models/test_qwen3_5_pipeline_parallel.py index f239893a3..9a9116f08 100644 --- a/test/registered/unit/models/test_qwen3_5_pipeline_parallel.py +++ b/test/registered/unit/models/test_qwen3_5_pipeline_parallel.py @@ -1,8 +1,11 @@ import unittest from types import SimpleNamespace +import torch + from sglang.srt.layers.utils import PPMissingLayer from sglang.srt.models.qwen3_5 import Qwen3_5MoeForConditionalGeneration +from sglang.srt.models.qwen3_5_mtp import Qwen3_5ForCausalLMMTP from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -10,6 +13,18 @@ register_cpu_ci(est_time=2, suite="base-a-test-cpu") class TestQwen3_5PipelineParallel(CustomTestCase): + @staticmethod + def _make_mtp_weight_loader_stub(): + model = Qwen3_5ForCausalLMMTP.__new__(Qwen3_5ForCausalLMMTP) + torch.nn.Module.__init__(model) + model.model = torch.nn.Module() + model.model.embed_tokens = torch.nn.Embedding(4, 3) + model.config = SimpleNamespace(num_experts=None) + model.quant_config = None + with torch.no_grad(): + model.model.embed_tokens.weight.fill_(torch.nan) + return model + @staticmethod def _get_num_fused_shared_experts(layers, start_layer, end_layer): model = SimpleNamespace( @@ -64,6 +79,26 @@ class TestQwen3_5PipelineParallel(CustomTestCase): self.assertEqual(num_fused_shared_experts, 0) + def test_mtp_loads_vl_target_embedding_for_last_pp_stage(self): + model = self._make_mtp_weight_loader_stub() + expected = torch.arange(12, dtype=torch.float32).reshape(4, 3) + + loaded = model.load_weights( + [("model.language_model.embed_tokens.weight", expected)] + ) + + self.assertEqual(loaded, {"model.embed_tokens.weight"}) + torch.testing.assert_close(model.model.embed_tokens.weight, expected) + + def test_mtp_loads_text_target_embedding_for_last_pp_stage(self): + model = self._make_mtp_weight_loader_stub() + expected = torch.arange(12, dtype=torch.float32).reshape(4, 3) + + loaded = model.load_weights([("model.embed_tokens.weight", expected)]) + + self.assertEqual(loaded, {"model.embed_tokens.weight"}) + torch.testing.assert_close(model.model.embed_tokens.weight, expected) + if __name__ == "__main__": unittest.main()