[Bugfix] Load Qwen3.5 MTP embedding under PP (#37471)

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
 

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()