Add disk and shared-memory weight sync for engines without NCCL or CUDA IPC

An engine on another accelerator stack (sglang-torchtpu) shares neither NCCL
nor CUDA IPC with the trainer, and serves no /pull_weights, so none of the
existing transfer modes reach it.

--update-weight-transfer-mode disk writes the gathered HF buckets as a fresh
safetensors checkpoint on a shared filesystem and has the engines reload it
through /update_weights_from_disk. --update-weight-transfer-mode shm copies each
flattened bucket into a shared-memory tensor and hands its handle to engines on
the trainer's hosts through /update_weights_from_tensor, so nothing is written
or parsed. In both, the ranks holding identical gathered buckets split them
round-robin and copy only after the collective gathers, so the host copies run
in parallel.

shm runs inside the begin/end weight-update session; disk pauses the engines
around its one reload call, as disk-delta does. Engine calls go through the
rollout-cell updaters, so a failed engine is reported per cell.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Author: HJSang

docs/advanced/disaggregated-rollout.md

@@ -107,6 +107,8 @@ miles-managed path:
 | `broadcast_packed` | Pack each weight bucket into one byte buffer and issue one NCCL broadcast | Non-colocated Megatron with SGLang mixed-dtype flattened-bucket support; trades contiguous bucket storage for fewer collectives |
 | `p2p` | Convert and re-shard weights, then write them directly to rollout-rank memory over RDMA | miles-managed, in-cluster jobs with direct rank-to-rank connectivity; see [P2P Weight Transfer](/advanced/p2p-weight-transfer) |
 | `disk-delta` | Publish changed canonical checkpoint bytes to shared storage, let rollout hosts materialize them locally, then reload | Trainer and rollout cannot share an NCCL fabric, or model-sized full-weight transfer dominates the update |
+| `disk` | Write the full HF checkpoint to shared storage on every sync, then reload it via `/update_weights_from_disk` | Engines on another accelerator stack (e.g. TPU) with no NCCL, CUDA IPC or `/pull_weights` |
+| `shm` | Copy each weight bucket into a host shared-memory tensor and hand the engines its handle via `/update_weights_from_tensor` | Like `disk`, with a single-node trainer and the engines on the same host |
 
 These are weight synchronization choices, not different rollout APIs. In the
 broadcast and P2P modes, miles transfers tensors into known engine ranks directly.

miles/backends/training_utils/weight_update/protocol.py

@@ -77,7 +77,7 @@ def pop_metrics(self) -> dict[str, float]:
 
 def get_weight_transfer_protocol(args: Namespace) -> WeightTransferProtocol:
     mode = getattr(args, "update_weight_transfer_mode", "broadcast")
-    if mode not in ("broadcast", "broadcast_packed", "p2p", "disk-delta"):
+    if mode not in ("broadcast", "broadcast_packed", "p2p", "disk-delta", "disk", "shm"):
         raise ValueError(f"Unknown --update-weight-transfer-mode {mode!r}")
     if mode == "broadcast_packed" and (getattr(args, "train_backend", None) != "megatron" or args.colocate):
         raise ValueError("broadcast_packed requires Megatron non-colocated weight transfer")
@@ -93,6 +93,14 @@ def get_weight_transfer_protocol(args: Namespace) -> WeightTransferProtocol:
         from miles.backends.training_utils.weight_update.protocols.delta import UpdateWeightFromDiskDelta
 
         return UpdateWeightFromDiskDelta(args)
+    if mode == "disk":
+        from miles.backends.training_utils.weight_update.protocols.disk import UpdateWeightFromDisk
+
+        return UpdateWeightFromDisk(args)
+    if mode == "shm":
+        from miles.backends.training_utils.weight_update.protocols.shm import UpdateWeightFromSharedMemory
+
+        return UpdateWeightFromSharedMemory(args)
     if mode == "p2p":
         from miles.backends.training_utils.weight_update.protocols.p2p import UpdateWeightP2P
 

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

@@ -0,0 +1,127 @@
+import json
+import logging
+import os
+import shutil
+from argparse import Namespace
+from collections.abc import Sequence
+
+import safetensors.torch
+import torch
+import torch.distributed as dist
+
+from miles.backends.sglang_utils.sglang_api_client import SGLangApiClient
+from miles.backends.training_utils.parallel import ParallelState
+from miles.backends.training_utils.weight_update.hf_weight_iterator import WeightUpdatePlacement
+from miles.backends.training_utils.weight_update.protocol import WeightTransferProtocol
+from miles.backends.training_utils.weight_update.protocols.utils.data_replica import get_data_replica_rank_and_size
+from miles.backends.training_utils.weight_update.rollout_cell_updater import create_rollout_cell_updaters
+from miles.backends.training_utils.weight_update.session import (
+    check_weight_sync_results,
+    maybe_pause_engines,
+    maybe_resume_engines,
+)
+from miles.utils import async_utils
+from miles.utils.distributed_utils import get_gloo_group
+
+logger = logging.getLogger(__name__)
+
+# everything else in the HF checkpoint (config, tokenizer, remote code) is copied into each version
+_WEIGHT_FILE_SUFFIXES = (".safetensors", ".bin", ".pt", ".pth", ".index.json")
+
+
+class UpdateWeightFromDisk(WeightTransferProtocol):
+    """
+    Full-checkpoint weight sync over a shared filesystem. The ranks holding the gathered HF buckets
+    split them round-robin and write each as a safetensors shard of a fresh version dir, rank 0 adds
+    the index and HF assets, and the engines reload it via /update_weights_from_disk. Needs no NCCL,
+    CUDA IPC or /pull_weights, so it works with engines on another accelerator stack.
+    """
+
+    # /update_weights_from_disk replaces every weight in one call, so no begin/end session
+    use_weight_update_session = False
+
+    def __init__(self, args: Namespace) -> None:
+        super().__init__(args)
+        self.root_dir = args.update_weight_disk_dir
+        self._version_dir: str | None = None
+        self._previous_version_dir: str | None = None
+        self._weight_map: dict[str, str] = {}
+        self._bucket_index = 0
+        self._owned_buckets: list[tuple[int, list[tuple[str, torch.Tensor]]]] = []
+
+    def connect(
+        self,
+        rollout_engines: Sequence[SGLangApiClient],
+        engine_gpu_counts: Sequence[int] | None,
+        engine_gpu_offsets: Sequence[int] | None,
+        engine_cell_ids: Sequence[str],
+        parallel_state: ParallelState,
+        placement: WeightUpdatePlacement,
+        selector: str,
+    ) -> None:
+        self.rollout_engines = rollout_engines
+        self.cell_updaters_of_cell_id = create_rollout_cell_updaters(self.args, self.rollout_engines, engine_cell_ids)
+        self.group_name = "miles-disk"
+        self._replica_rank, self._replica_size = get_data_replica_rank_and_size(parallel_state, placement)
+        self.is_sender = True
+
+    def begin_sync(self, weight_version: int, iter_buckets) -> bool:
+        self._version_dir = os.path.join(self.root_dir, f"weight_v{weight_version:06d}")
+        self._weight_map = {}
+        self._bucket_index = 0
+        self._owned_buckets = []
+        os.makedirs(self._version_dir, exist_ok=True)
+        return True
+
+    def send_bucket(self, bucket: list[tuple[str, torch.Tensor]]) -> None:
+        if self._bucket_index % self._replica_size == self._replica_rank:
+            self._owned_buckets.append((self._bucket_index, bucket))
+        self._bucket_index += 1
+
+    def after_base_weights(self) -> None:
+        # written only after the collective gathers finish, so the ranks' host copies run in parallel
+        for bucket_index, bucket in self._owned_buckets:
+            tensors = {name: tensor.detach().to("cpu").contiguous() for name, tensor in bucket}
+            shard = f"rank{dist.get_rank():03d}-shard{bucket_index:05d}.safetensors"
+            safetensors.torch.save_file(tensors, os.path.join(self._version_dir, shard))
+            self._weight_map.update({name: shard for name in tensors})
+        self._owned_buckets = []
+
+    def finalize(self, weight_version: int) -> None:
+        self._write_index()
+        self._reload_engines(weight_version)
+        self._drop_previous_version()
+
+    def _write_index(self) -> None:
+        weight_maps: list = [None] * dist.get_world_size()
+        dist.all_gather_object(weight_maps, self._weight_map, group=get_gloo_group())
+        if dist.get_rank() == 0:
+            weight_map = {name: shard for rank_map in weight_maps for name, shard in rank_map.items()}
+            index = {"metadata": {}, "weight_map": weight_map}
+            with open(os.path.join(self._version_dir, "model.safetensors.index.json"), "w") as f:
+                json.dump(index, f)
+            for fname in os.listdir(self.args.hf_checkpoint):
+                path = os.path.join(self.args.hf_checkpoint, fname)
+                if os.path.isfile(path) and not fname.endswith(_WEIGHT_FILE_SUFFIXES):
+                    shutil.copy(path, self._version_dir)
+            logger.info("[disk] wrote %d tensors to %s", len(weight_map), self._version_dir)
+
+    def _reload_engines(self, weight_version: int) -> None:
+        if dist.get_rank() == 0:
+            cell_updaters = list(self.cell_updaters_of_cell_id.values())
+            maybe_pause_engines(self.args, cell_updaters)
+            results = async_utils.wait_futures(
+                [
+                    updater.submit_client_call(
+                        "update_weights_from_disk", model_path=self._version_dir, weight_version=str(weight_version)
+                    )
+                    for updater in cell_updaters
+                ]
+            )
+            check_weight_sync_results(results, is_lora=False)
+            maybe_resume_engines(self.args, cell_updaters)
+
+    def _drop_previous_version(self) -> None:
+        if dist.get_rank() == 0 and self._previous_version_dir not in (None, self._version_dir):
+            shutil.rmtree(self._previous_version_dir, ignore_errors=True)
+        self._previous_version_dir = self._version_dir

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

@@ -0,0 +1,103 @@
+from argparse import Namespace
+from collections import deque
+from collections.abc import Sequence
+from concurrent.futures import Future
+
+import torch
+import torch.multiprocessing
+from sglang.srt.utils import MultiprocessingSerializer
+from sglang.srt.weight_sync.tensor_bucket import FlattenedTensorBucket
+
+from miles.backends.sglang_utils.sglang_api_client import SGLangApiClient
+from miles.backends.training_utils.parallel import ParallelState
+from miles.backends.training_utils.weight_update.hf_weight_iterator import WeightUpdatePlacement
+from miles.backends.training_utils.weight_update.protocol import WeightTransferProtocol
+from miles.backends.training_utils.weight_update.protocols.utils.data_replica import get_data_replica_rank_and_size
+from miles.backends.training_utils.weight_update.rollout_cell_updater import create_rollout_cell_updaters
+from miles.backends.training_utils.weight_update.session import check_weight_sync_results
+from miles.utils import async_utils
+
+
+# enough buckets in flight to keep the engines loading while the next one is copied out; arbitrary otherwise
+_MAX_BUCKETS_IN_FLIGHT = 4
+
+
+class UpdateWeightFromSharedMemory(WeightTransferProtocol):
+    """
+    Host weight sync through shared memory, for engines on the trainer's hosts. The ranks holding the
+    gathered buckets split them round-robin, copy each flattened bucket into a shared-memory tensor
+    once the collective gathers are done, and hand the engines its handle through
+    /update_weights_from_tensor; each engine maps it in place, with no file written or parsed.
+    A bucket is published as soon as it is on the host, so the copies overlap the engine loads, and at most
+    `_MAX_BUCKETS_IN_FLIGHT` host copies are alive at once, so host memory stays a few buckets at any model size.
+    """
+
+    def __init__(self, args: Namespace) -> None:
+        super().__init__(args)
+        # shared-memory tensors travel by /dev/shm file name, which an unrelated process can open
+        torch.multiprocessing.set_sharing_strategy("file_system")
+        self._bucket_index = 0
+        self._owned_buckets: list[list[tuple[str, torch.Tensor]]] = []
+        # each published host copy with its engine calls; a copy lives until every engine has loaded it
+        self._in_flight: deque[tuple[torch.Tensor, list[Future]]] = deque()
+
+    def connect(
+        self,
+        rollout_engines: Sequence[SGLangApiClient],
+        engine_gpu_counts: Sequence[int] | None,
+        engine_gpu_offsets: Sequence[int] | None,
+        engine_cell_ids: Sequence[str],
+        parallel_state: ParallelState,
+        placement: WeightUpdatePlacement,
+        selector: str,
+    ) -> None:
+        self.rollout_engines = rollout_engines
+        self.cell_updaters_of_cell_id = create_rollout_cell_updaters(self.args, self.rollout_engines, engine_cell_ids)
+        self.group_name = "miles-shm"
+        self.is_sender = True
+        self._selector = selector
+        self._engine_tp_sizes = engine_gpu_counts
+        self._replica_rank, self._replica_size = get_data_replica_rank_and_size(parallel_state, placement)
+
+    def begin_sync(self, weight_version: int, iter_buckets) -> bool:
+        self._bucket_index = 0
+        self._owned_buckets = []
+        return True
+
+    def send_bucket(self, bucket: list[tuple[str, torch.Tensor]]) -> None:
+        if self._bucket_index % self._replica_size == self._replica_rank:
+            self._owned_buckets.append(bucket)
+        self._bucket_index += 1
+
+    def after_base_weights(self) -> None:
+        # copied out only after the collective gathers finish, so the ranks' host copies run in parallel
+        cell_updaters = list(self.cell_updaters_of_cell_id.values())
+        for bucket in self._owned_buckets:
+            if len(self._in_flight) == _MAX_BUCKETS_IN_FLIGHT:
+                self._wait_for_oldest_bucket()
+            flattened = FlattenedTensorBucket(named_tensors=bucket)
+            device_tensor = flattened.get_flattened_tensor()
+            host_tensor = torch.empty_like(device_tensor, device="cpu").share_memory_()
+            host_tensor.copy_(device_tensor)
+            payload = MultiprocessingSerializer.serialize(
+                {"flattened_tensor": host_tensor, "metadata": flattened.get_metadata()}, output_str=True
+            )
+            futures = [
+                updater.submit_client_call(
+                    "update_weights_from_tensor",
+                    serialized_named_tensors=[payload] * tp_size,
+                    load_format="flattened_bucket",
+                    selector=self._selector,
+                )
+                for updater, tp_size in zip(cell_updaters, self._engine_tp_sizes, strict=True)
+            ]
+            self._in_flight.append((host_tensor, futures))
+        self._owned_buckets = []
+
+    def finalize(self, weight_version: int) -> None:
+        while self._in_flight:
+            self._wait_for_oldest_bucket()
+
+    def _wait_for_oldest_bucket(self) -> None:
+        _, futures = self._in_flight.popleft()
+        check_weight_sync_results(async_utils.wait_futures(futures), is_lora=False)

miles/utils/arguments.py

@@ -1114,7 +1114,7 @@ def add_rollout_arguments(parser):
             )
             parser.add_argument(
                 "--update-weight-transfer-mode",
-                choices=["broadcast", "broadcast_packed", "p2p", "disk-delta"],
+                choices=["broadcast", "broadcast_packed", "p2p", "disk-delta", "disk", "shm"],
                 default="broadcast",
                 help=(
                     "The method to transfer weights to remote rollout engines during update weight. "
@@ -1125,17 +1125,22 @@ def add_rollout_arguments(parser):
                     "exceed --update-weight-buffer-size. "
                     "'disk-delta' diffs each sync against a CPU snapshot of the previous one and publishes "
                     "only the changed bytes to --update-weight-disk-dir; each engine's /pull_weights applies "
-                    "them into a host-local checkpoint that the engine reloads from."
+                    "them into a host-local checkpoint that the engine reloads from. "
+                    "'disk' writes the full HF checkpoint to --update-weight-disk-dir on every sync and has "
+                    "the engines reload it via /update_weights_from_disk; it needs no NCCL, CUDA IPC or "
+                    "/pull_weights, so it works with engines on another accelerator stack (e.g. TPU). "
+                    "'shm' hands engines on the trainer's hosts shared-memory tensors through "
+                    "/update_weights_from_tensor, skipping the files 'disk' writes and parses."
                 ),
             )
             parser.add_argument(
                 "--update-weight-disk-dir",
                 type=str,
                 default=None,
                 help=(
-                    "Filesystem directory disk-delta weight sync publishes to: one delta directory "
-                    "(changed tensors only) per sync, written by the trainer and read by every "
-                    "rollout host. Required for --update-weight-transfer-mode=disk-delta."
+                    "Filesystem directory disk and disk-delta weight sync publish to, written by the "
+                    "trainer and read by every rollout host: one full checkpoint (disk) or one delta "
+                    "directory with the changed tensors only (disk-delta) per sync. Required for both modes."
                 ),
             )
             parser.add_argument(
@@ -3439,7 +3444,7 @@ def miles_validate_args(args):
             setattr(args, k, v)
 
     mode = args.update_weight_transfer_mode
-    if mode not in ("broadcast", "broadcast_packed", "p2p", "disk-delta"):
+    if mode not in ("broadcast", "broadcast_packed", "p2p", "disk-delta", "disk", "shm"):
         raise ValueError(f"Unknown --update-weight-transfer-mode {mode!r}")
     if mode == "broadcast_packed" and (args.train_backend != "megatron" or args.colocate):
         raise ValueError("broadcast_packed requires Megatron non-colocated weight transfer")
@@ -3883,6 +3888,32 @@ def miles_validate_args(args):
                 f"weights restored from --load={args.requested_load}."
             )
 
+    if args.update_weight_transfer_mode in ("disk", "shm"):
+        mode = args.update_weight_transfer_mode
+        # colocate always syncs through CUDA IPC (get_weight_transfer_protocol)
+        assert not args.colocate, f"--update-weight-transfer-mode {mode} needs disaggregated engines; drop --colocate"
+        assert args.lora_rank <= 0, f"LoRA weight sync is not supported for --update-weight-transfer-mode {mode}"
+        assert (
+            args.prefill_num_servers is None
+        ), f"--update-weight-transfer-mode {mode} has not been tested when PD is enabled"
+    if args.update_weight_transfer_mode == "shm":
+        # the engines open the shared-memory tensors by /dev/shm file name, so they must share the trainer's host
+        assert args.actor_num_nodes == 1, "--update-weight-transfer-mode shm needs the trainer on one node"
+    if args.update_weight_transfer_mode == "disk":
+        # /update_weights_from_disk reloads the draft model from the same directory, which has no draft weights
+        assert args.sglang_speculative_algorithm is None or args.enable_mtp_training, (
+            "--update-weight-transfer-mode disk reloads the speculative draft model from the trainer's "
+            "checkpoint; set --enable-mtp-training or drop --sglang-speculative-algorithm"
+        )
+        assert args.update_weight_disk_dir, (
+            "--update-weight-transfer-mode disk needs --update-weight-disk-dir on a filesystem "
+            "both the trainer and the rollout engines can read"
+        )
+        assert os.path.isdir(args.hf_checkpoint), (
+            "--update-weight-transfer-mode disk needs --hf-checkpoint to be a local directory: "
+            "each published version copies its config and tokenizer files"
+        )
+
     if args.colocate:
         if args.offload_train is None:
             args.offload_train = True