p2p: update speculative drafts from trained MTP weights

docs/advanced/p2p-weight-transfer.md

@@ -50,6 +50,7 @@ Both broadcast and P2P modes share the same bucketed weight-update pipeline in `
 | **Rollout engine rank assignment** | `protocols/utils/rollout_engine_rank_assignment.py` | Decides which rollout engine ranks each training rank sends weights to, for the rollout engines the rollout manager passes in (`engine_gpu_counts`). Senders are the training ranks that hold the same weights: those of one PP stage, or all ranks when PP is gathered (Megatron Bridge). Every (rollout engine, rollout engine rank) gets exactly one sender. First every sender gets one, so all of them send in parallel; the rest go to a sender that already sends the same rollout engine rank, so it prepares those weights once and sends them to several rollout engines. |
 | **Model replica** | `protocols/utils/model_replica.py` | An sglang model of the target rank's layout, built on `meta` with no parameter storage and brought, one module at a time on the GPU, to the state the engine loads updates into: its startup postprocess, then the restore every update session begins with. Its `weight_loader` functions re-shard all-gathered HF weights into the exact bytes the target rank expects, loaded into a transfer buffer. `ModelReplicas` keeps one replica per shard layout (the rank's parallelism and the sglang arguments that change its bytes and may differ between server groups, such as quantization and MoE and GEMM backends) for the whole trainer process, so a reconnect rebuilds only the connections to the new engines. Rollout engines whose expert placement a replica cannot reproduce (redundant experts, a non-trivial `init_expert_location`, EPLB, elastic EP, DWDP, KTransformers) are rejected at connect. |
 | **Transfer buffers** | `protocols/utils/transfer_buffers.py` | Two pinned host buffers per sender, sized to the larger of `--update-weight-buffer-size` and the largest complete parameter group, including alignment. Parameters sharing HF inputs stay in one group. Buffers are registered once per trainer process and reused only after every write reading them has returned; a write that failed or is still running after `--p2p-transfer-timeout` fails the update there. |
+| **Speculative draft** | `p2p.py`, `protocols/utils/model_replica.py` | With MTP training and same-model EAGLE, senders update both `target` and `draft` roles. Each runner receives every weight its own loader loads, as in the engine's own weight updates; a weight both loaders load into storage the draft shares with the target is written twice with the same bytes. Draft replicas use the engine's speculative MoE settings. Multi-layer EAGLE, separate draft checkpoints and a draft MoE runner or a2a backend other than the target's are rejected at connect, since sglang reloads a draft under the target's MoE backends. |
 | **Pipelined transfer** | `p2p.py`, `protocols/transports/mooncake.py` | While the writes reading one transfer buffer are in flight, the sender loads the next group of parameters into the other. Each rollout engine has its own write thread, so a stuck engine delays only its own writes. The end of the update waits for the remaining writes and fails it if any failed or is still running after `--p2p-transfer-timeout`. |
 
 The model replica runs the rollout engine's own sglang code in the trainer process: its model construction, startup postprocess and weight loaders. The trainer and the rollout engines therefore run the same sglang build, and their GPUs share an architecture: postprocess chooses kernels and layouts by architecture (TRT-LLM MoE layouts and UE8M0 scales on Blackwell, for example), so a replica on a Hopper trainer would lay out bytes for Hopper while Blackwell engines expect their own.

miles/backends/sglang_utils/sglang_api_client.py

@@ -136,21 +136,21 @@ async def update_weights_from_tensor(
             payload,
         )
 
-    async def get_remote_instance_transfer_engine_info(self, rank: int):
+    async def get_remote_instance_transfer_engine_info(self, rank: int, role: str):
         # TODO: will be changed to `remote_instance_transfer_engine_info` when the sglang side is ready.
         response = await GeneralHttpClientProvider.client().get(
             f"{self.server_url}/get_remote_instance_transfer_engine_info",
-            params={"rank": rank},
+            params={"rank": rank, "role": role},
             headers=self._headers,
             timeout=5.0,
         )
         response.raise_for_status()
         return response.json()["remote_instance_transfer_engine_info"]
 
-    async def get_parallelism_info(self, rank: int):
+    async def get_parallelism_info(self, rank: int, role: str):
         response = await GeneralHttpClientProvider.client().get(
             f"{self.server_url}/parallelism_config",
-            params={"rank": rank},
+            params={"rank": rank, "role": role},
             headers=self._headers,
             timeout=5.0,
         )