fix(p2p): publish parameter dtypes for weight transfers

python/sglang/srt/entrypoints/engine_info_bootstrap_server.py

@@ -43,6 +43,7 @@ def __init__(self, host: str, port: int):
 
         # Storage: {(role, tp_rank): (session_id, weights_info_dict)}
         self.transfer_engine_info: Dict[Tuple[str, int], Tuple] = {}
+        self.weight_dtypes: Dict[Tuple[str, int], dict] = {}
         # Storage: {(role, tp_rank): parallelism_config_dict}
         self.parallelism_config: Dict[Tuple[str, int], dict] = {}
         self.lock = threading.Lock()
@@ -67,6 +68,7 @@ def register_transfer_engine_info(data: dict):
                         session_id,
                         weights_info_dict,
                     )
+                    self.weight_dtypes[(role, tp_rank)] = info.get("weight_dtypes", {})
 
                 logger.info(
                     f"Registered transfer engine info for role={role} tp_rank={tp_rank}, "
@@ -84,14 +86,19 @@ def get_transfer_engine_info(rank: int, role: str = "target"):
 
             with self.lock:
                 info = self.transfer_engine_info.get((role, rank))
+                weight_dtypes = self.weight_dtypes.get((role, rank), {})
 
             if info is None:
                 raise HTTPException(
                     status_code=404,
                     detail=f"No transfer engine info for role {role} rank {rank}",
                 )
 
-            return {"rank": rank, "remote_instance_transfer_engine_info": list(info)}
+            return {
+                "rank": rank,
+                "remote_instance_transfer_engine_info": list(info),
+                "weight_dtypes": weight_dtypes,
+            }
 
         config = uvicorn.Config(app, host=host, port=port, log_level="warning")
         self._server = uvicorn.Server(config)

python/sglang/srt/model_executor/model_runner_components/remote_instance_weight_transporter.py

@@ -158,6 +158,10 @@ def _register_to_engine_info_bootstrap(
             "transfer_engine_info": {
                 "session_id": self.session_id,
                 "weights_info_dict": self.weight_info,
+                "weight_dtypes": {
+                    name: str(weight.dtype).removeprefix("torch.")
+                    for name, weight in self.model.named_parameters()
+                },
             },
         }
 

test/registered/unit/model_executor/model_runner_components/test_remote_instance_weight_roles.py

@@ -44,6 +44,8 @@ def _start_bootstrap_server():
 
 
 def _publish(role, session_id, port):
+    model = torch.nn.Module()
+    model.add_module(role, torch.nn.Linear(4, 1, bias=False, dtype=torch.bfloat16))
     model_config = SimpleNamespace(
         remote_instance_weight_loader_backend=None, engine_info_bootstrap_port=port
     )
@@ -65,7 +67,7 @@ def _publish(role, session_id, port):
     ):
         transporter = RemoteInstanceWeightTransporter(
             server_args=SimpleNamespace(registers_parallelism_config=lambda: True),
-            get_model=lambda: torch.nn.Module(),
+            get_model=lambda: model,
             gpu_id=0,
         )
         transporter.engine = object()
@@ -98,6 +100,14 @@ def get(path, **params):
             ["draft-session", {"draft.weight": [1, 4, 2]}],
         )
         self.assertEqual(get("get_parallelism_config"), {"published_by": "target"})
+        self.assertEqual(
+            get("get_transfer_engine_info")["weight_dtypes"],
+            {"target.weight": "bfloat16"},
+        )
+        self.assertEqual(
+            get("get_transfer_engine_info", role="draft")["weight_dtypes"],
+            {"draft.weight": "bfloat16"},
+        )
         self.assertEqual(
             get("get_parallelism_config", role="draft"), {"published_by": "draft"}
         )