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. |
+| **Speculative draft** | `p2p.py`, `protocols/utils/model_replica.py` | With MTP training and same-model EAGLE, senders update both `target` and `draft` roles. The target receives shared embedding and LM-head weights once; the draft receives its MTP weights. Draft replicas use the engine's speculative MoE settings. Multi-layer EAGLE and separate draft checkpoints are rejected at connect. |
 | **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 every write and fails it if any write failed or is still running after `--p2p-transfer-timeout`. |
 
 ## Supported Model Architectures

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

miles/backends/training_utils/weight_update/protocols/p2p.py

@@ -1,6 +1,7 @@
 import logging
 from argparse import Namespace
 from collections.abc import Callable, Iterator, Sequence
+from dataclasses import dataclass
 from typing import NamedTuple
 
 import torch
@@ -16,10 +17,12 @@
 from miles.backends.training_utils.weight_update.protocols.utils.model_replica import (
     ModelReplica,
     ModelReplicas,
+    RolloutEngineRankConfig,
     assert_replica_matches_shard,
     compute_param_group_nbytes,
     pack_into_buffers,
     query_rollout_engine_rank_configs,
+    query_runner_roles,
 )
 from miles.backends.training_utils.weight_update.protocols.utils.rollout_engine_rank_assignment import (
     assign_rollout_engine_ranks,
@@ -38,28 +41,31 @@ class _ReplicaTarget(NamedTuple):
     remote_shards: list[RemoteShard]
 
 
-class UpdateWeightP2P(WeightTransferProtocol):
-    """Writes weight updates straight into the rollout engines' GPU memory over Mooncake.
+@dataclass
+class _RunnerSender:
+    """Per-connection write targets and per-sync staging for one runner role, backed by cached replicas."""
+
+    model_replicas: ModelReplicas
+    replica_targets: list[_ReplicaTarget]
+    model_param_stager: ModelParamStager | None = None
+
+    def loads_hf_tensor(self, hf_name: str) -> bool:
+        return hf_name in self.model_replicas.hf_name_mapping.param_names_by_hf_name
 
-    Each sender loads the HF tensors of the updater's buckets with a model replica of each rollout engine rank it
-    sends to, into transfer buffers in the bytes that rank's loader would write, and writes those bytes to the
-    rank's published addresses. A write that failed or is still running after --p2p-transfer-timeout fails the
-    update when its transfer buffer is next loaded, or at the end of the base weights.
-    """
+
+class UpdateWeightP2P(WeightTransferProtocol):
+    """Write target and supported MTP draft weights to published GPU addresses over Mooncake."""
 
     def __init__(self, args: Namespace) -> None:
         super().__init__(args)
         if args.sglang_pp_size != 1:
             raise NotImplementedError("Rollout pipeline parallelism is not tested yet.")
         self.global_rank = dist.get_rank(group=get_gloo_group())
         self._hf_tensor_specs: dict[str, HfTensorSpec] | None = None
-        self._model_param_stager: ModelParamStager | None = None
-        self._model_replicas = ModelReplicas(
-            model_path=args.hf_checkpoint, transfer_buffer_device=MooncakeTransport.transfer_buffer_device
-        )
+        self._model_replicas_by_runner_role: dict[str, ModelReplicas] = {}
         self._transport: MooncakeTransport | None = None
         self._transfer_buffers: TransferBuffers | None = None
-        self._replica_targets: list[_ReplicaTarget] = []
+        self._runner_senders: list[_RunnerSender] = []
 
     def begin_sync(
         self, weight_version: int, iter_buckets: Callable[..., Iterator[list[tuple[str, torch.Tensor]]]]
@@ -72,10 +78,11 @@ def begin_sync(
                 for hf_name, tensor in bucket
             }
         if self.is_sender:
-            self._model_replicas.map_hf_names(self._hf_tensor_specs)
-            self._model_param_stager = ModelParamStager(self._model_replicas.hf_name_mapping)
+            for runner_sender in self._runner_senders:
+                runner_sender.model_replicas.map_hf_names(self._hf_tensor_specs)
+                runner_sender.model_param_stager = ModelParamStager(runner_sender.model_replicas.hf_name_mapping)
             if self._transfer_buffers is None:
-                self._log_hf_names_no_replica_loads()
+                self._log_hf_names_no_runner_loads()
                 self._transfer_buffers = self._create_transfer_buffers()
         return True
 
@@ -84,37 +91,51 @@ def after_base_weights(self) -> None:
         if not self.is_sender:
             return
         self._transfer_buffers.wait_for_all_writes()
-        self._model_param_stager.assert_all_done()
+        for runner_sender in self._runner_senders:
+            runner_sender.model_param_stager.assert_all_done()
 
     def send_bucket(self, converted_named_tensors: list[tuple[str, torch.Tensor]]) -> None:
-        """Write complete parameter groups; skip HF tensors the engine's loader ignores."""
+        """Write complete parameter groups, routing shared embed/head weights through the target only."""
         if not self.is_sender or not converted_named_tensors:
             return
         unknown_hf_names = [hf_name for hf_name, _ in converted_named_tensors if hf_name not in self._hf_tensor_specs]
         assert not unknown_hf_names, (
             f"{unknown_hf_names[:5]} ({len(unknown_hf_names)} in all) were not in the trainer's first pass, so no "
             "replica was mapped for them"
         )
-        ready_hf_tensors_by_param_group = self._model_param_stager.stage(
-            (hf_name, tensor) for hf_name, tensor in converted_named_tensors if self._replica_loads(hf_name)
-        )
+        unclaimed_hf_tensors = converted_named_tensors
+        for runner_sender in self._runner_senders:
+            claimed_hf_tensors = [
+                (hf_name, tensor) for hf_name, tensor in unclaimed_hf_tensors if runner_sender.loads_hf_tensor(hf_name)
+            ]
+            unclaimed_hf_tensors = [
+                (hf_name, tensor)
+                for hf_name, tensor in unclaimed_hf_tensors
+                if not runner_sender.loads_hf_tensor(hf_name)
+            ]
+            self._send_runner_bucket(runner_sender, claimed_hf_tensors)
+
+        converted_named_tensors.clear()
+
+    def _send_runner_bucket(
+        self, runner_sender: _RunnerSender, converted_named_tensors: list[tuple[str, torch.Tensor]]
+    ) -> None:
+        ready_hf_tensors_by_param_group = runner_sender.model_param_stager.stage(converted_named_tensors)
         for param_groups in pack_into_buffers(
             ready_hf_tensors_by_param_group,
-            self._model_replicas.transfer_buffer_param_layouts,
+            runner_sender.model_replicas.transfer_buffer_param_layouts,
             self._transfer_buffers.buffer_nbytes,
         ):
             param_names = [param_name for param_group in param_groups for param_name in param_group]
             hf_tensors = [
                 hf_tensor for param_group in param_groups for hf_tensor in ready_hf_tensors_by_param_group[param_group]
             ]
-            for target in self._replica_targets:
+            for target in runner_sender.replica_targets:
                 buffer = self._transfer_buffers.acquire()
                 param_bytes_by_name = target.model_replica.load_into(buffer, param_names, hf_tensors)
                 writes = self._transport.write(target.remote_shards, param_bytes_by_name)
                 self._transfer_buffers.release_after(buffer, list(zip(target.remote_shards, writes, strict=True)))
 
-        converted_named_tensors.clear()
-
     def connect(
         self,
         rollout_engines: Sequence[SGLangApiClient],
@@ -130,39 +151,64 @@ def connect(
         self.is_sender = bool(assignments)
 
         if self.is_sender:
-            configs_by_rollout_engine_rank = query_rollout_engine_rank_configs(rollout_engines, assignments)
+            runner_roles = query_runner_roles(rollout_engines, assignments, selector)
+            configs_by_runner_role = {
+                runner_role: query_rollout_engine_rank_configs(rollout_engines, assignments, runner_role)
+                for runner_role in runner_roles
+            }
             if self._transport is None:
                 self._transport = MooncakeTransport()
-            remote_shards_by_rollout_engine_rank = self._transport.connect(rollout_engines, assignments)
-            self._replica_targets = []
-            for rollout_engine_rank, remote_shards in remote_shards_by_rollout_engine_rank.items():
-                config = configs_by_rollout_engine_rank[rollout_engine_rank]
-                model_replica = self._model_replicas.get_or_build(config)
-                for remote_shard in remote_shards:
-                    assert_replica_matches_shard(
-                        model_replica,
-                        remote_shard.published_nbytes_by_name,
-                        published_by=f"rollout engine {remote_shard.rollout_engine_ind} rank {rollout_engine_rank}",
-                    )
-                self._replica_targets.append(_ReplicaTarget(model_replica, remote_shards))
-
-    def _replica_loads(self, hf_name: str) -> bool:
-        return hf_name in self._model_replicas.hf_name_mapping.param_names_by_hf_name
-
-    def _log_hf_names_no_replica_loads(self) -> None:
-        hf_names_no_replica_loads = sorted(
-            hf_name for hf_name in self._hf_tensor_specs if not self._replica_loads(hf_name)
+            remote_shards_by_runner_role = self._transport.connect(rollout_engines, assignments, runner_roles)
+            self._runner_senders = [
+                self._connect_runner(
+                    runner_role, configs_by_runner_role[runner_role], remote_shards_by_runner_role[runner_role]
+                )
+                for runner_role in runner_roles
+            ]
+
+    def _connect_runner(
+        self,
+        runner_role: str,
+        configs_by_rollout_engine_rank: dict[int, RolloutEngineRankConfig],
+        remote_shards_by_rollout_engine_rank: dict[int, list[RemoteShard]],
+    ) -> _RunnerSender:
+        if runner_role not in self._model_replicas_by_runner_role:
+            self._model_replicas_by_runner_role[runner_role] = ModelReplicas(
+                model_path=self.args.hf_checkpoint, transfer_buffer_device=self._transport.transfer_buffer_device
+            )
+        model_replicas = self._model_replicas_by_runner_role[runner_role]
+        replica_targets = []
+        for rollout_engine_rank, remote_shards in remote_shards_by_rollout_engine_rank.items():
+            model_replica = model_replicas.get_or_build(configs_by_rollout_engine_rank[rollout_engine_rank])
+            for remote_shard in remote_shards:
+                assert_replica_matches_shard(
+                    model_replica,
+                    remote_shard.published_nbytes_by_name,
+                    published_by=(
+                        f"the {runner_role} of rollout engine {remote_shard.rollout_engine_ind} "
+                        f"rank {rollout_engine_rank}"
+                    ),
+                )
+            replica_targets.append(_ReplicaTarget(model_replica, remote_shards))
+        return _RunnerSender(model_replicas, replica_targets)
+
+    def _log_hf_names_no_runner_loads(self) -> None:
+        hf_names_no_runner_loads = sorted(
+            hf_name
+            for hf_name in self._hf_tensor_specs
+            if not any(runner_sender.loads_hf_tensor(hf_name) for runner_sender in self._runner_senders)
         )
-        if hf_names_no_replica_loads:
+        if hf_names_no_runner_loads:
             logger.info(
-                f"p2p skips {len(hf_names_no_replica_loads)} HF tensors the rollout engines' loaders ignore, e.g. "
-                f"{hf_names_no_replica_loads[:5]}"
+                f"p2p skips {len(hf_names_no_runner_loads)} HF tensors the rollout engines' loaders ignore, e.g. "
+                f"{hf_names_no_runner_loads[:5]}"
             )
 
     def _create_transfer_buffers(self) -> TransferBuffers:
         largest_param_group_nbytes = max(
-            compute_param_group_nbytes(param_group, self._model_replicas.transfer_buffer_param_layouts)
-            for param_group in self._model_param_stager.param_groups
+            compute_param_group_nbytes(param_group, runner_sender.model_replicas.transfer_buffer_param_layouts)
+            for runner_sender in self._runner_senders
+            for param_group in runner_sender.model_param_stager.param_groups
         )
         return TransferBuffers(
             _NUM_TRANSFER_BUFFERS,

miles/backends/training_utils/weight_update/protocols/transports/mooncake.py

@@ -21,10 +21,11 @@ class RemoteWeightLocation(NamedTuple):
 
 @dataclass(frozen=True)
 class RemoteShard:
-    """One rank of one rollout engine as a write target: its Mooncake session and where each weight lives."""
+    """Mooncake session and published weight addresses for one engine rank's target or draft runner."""
 
     rollout_engine_ind: int
     rollout_engine_rank: int
+    runner_role: str
     session_id: str
     weight_locations_by_name: dict[str, RemoteWeightLocation]
 
@@ -51,9 +52,12 @@ def __init__(self) -> None:
         self._write_executors_by_rollout_engine_ind: dict[int, ThreadPoolExecutor] = {}
 
     def connect(
-        self, rollout_engines: Sequence[SGLangApiClient], assignments: Sequence[RolloutEngineRankAssignment]
-    ) -> dict[int, list[RemoteShard]]:
-        """Returns the shards of the rollout engine ranks in `assignments`, by rollout engine rank.
+        self,
+        rollout_engines: Sequence[SGLangApiClient],
+        assignments: Sequence[RolloutEngineRankAssignment],
+        runner_roles: Sequence[str],
+    ) -> dict[str, dict[int, list[RemoteShard]]]:
+        """Return assigned shards by runner role and rollout rank.
 
         Starts new write threads: a rollout engine index now names a new engine, which must not queue behind the
         writes of the one it replaced.
@@ -62,13 +66,19 @@ def connect(
             write_executor.shutdown(wait=False)
         self._write_executors_by_rollout_engine_ind = {}
         return {
-            assignment.rollout_engine_rank: [
-                _query_remote_shard(
-                    rollout_engines[rollout_engine_ind], rollout_engine_ind, assignment.rollout_engine_rank
-                )
-                for rollout_engine_ind in assignment.rollout_engine_indices
-            ]
-            for assignment in assignments
+            runner_role: {
+                assignment.rollout_engine_rank: [
+                    _query_remote_shard(
+                        rollout_engines[rollout_engine_ind],
+                        rollout_engine_ind,
+                        assignment.rollout_engine_rank,
+                        runner_role,
+                    )
+                    for rollout_engine_ind in assignment.rollout_engine_indices
+                ]
+                for assignment in assignments
+            }
+            for runner_role in runner_roles
         }
 
     def register_memory(self, tensor: torch.Tensor) -> None:
@@ -113,17 +123,18 @@ def _write_shard(self, remote_shard: RemoteShard, tensors_by_name: dict[str, tor
 
 
 def _query_remote_shard(
-    rollout_engine: SGLangApiClient, rollout_engine_ind: int, rollout_engine_rank: int
+    rollout_engine: SGLangApiClient, rollout_engine_ind: int, rollout_engine_rank: int, runner_role: str
 ) -> RemoteShard:
     session_id, weights_info = async_utils.run(
-        rollout_engine.get_remote_instance_transfer_engine_info(rank=rollout_engine_rank)
+        rollout_engine.get_remote_instance_transfer_engine_info(rank=rollout_engine_rank, role=runner_role)
     )
     assert (
         session_id is not None
-    ), f"rollout engine {rollout_engine_ind} rank {rollout_engine_rank} has no Mooncake session"
+    ), f"the {runner_role} of rollout engine {rollout_engine_ind} rank {rollout_engine_rank} has no Mooncake session"
     return RemoteShard(
         rollout_engine_ind=rollout_engine_ind,
         rollout_engine_rank=rollout_engine_rank,
+        runner_role=runner_role,
         session_id=session_id,
         weight_locations_by_name={name: RemoteWeightLocation(*location) for name, location in weights_info.items()},
     )
@@ -139,7 +150,10 @@ def _create_transfer_engine() -> Any:
 
 
 def _assert_tensors_fit(remote_shard: RemoteShard, tensors_by_name: dict[str, torch.Tensor]) -> None:
-    target = f"rollout engine {remote_shard.rollout_engine_ind} rank {remote_shard.rollout_engine_rank}"
+    target = (
+        f"the {remote_shard.runner_role} of rollout engine {remote_shard.rollout_engine_ind} "
+        f"rank {remote_shard.rollout_engine_rank}"
+    )
     for name, tensor in tensors_by_name.items():
         assert name in remote_shard.weight_locations_by_name, f"{target} publishes no {name}"
         location = remote_shard.weight_locations_by_name[name]