RDT weight sync: GPU->GPU zero copy transfer through SGLang Ray actor backend

miles/backends/megatron_utils/actor.py

@@ -250,7 +250,11 @@ def init(
         if self.args.vocab_size is None:
             self.args.vocab_size = self.tokenizer.vocab_size
 
-        if self.args.colocate:
+        if self.args.update_weight_transfer_mode == "rdt":
+            from .update_weight.update_weight_from_rdt import UpdateWeightFromRDT
+
+            update_weight_cls = UpdateWeightFromRDT
+        elif self.args.colocate:
             update_weight_cls = UpdateWeightFromTensor
         else:
             if self.args.update_weight_transfer_mode == "broadcast":

miles/backends/megatron_utils/update_weight/update_weight_from_distributed/p2p_transfer_utils.py

@@ -8,7 +8,6 @@
 import ray
 import torch
 import torch.distributed as dist
-from mooncake.engine import TransferEngine
 from ray.actor import ActorHandle
 from sglang.srt.server_args import ServerArgs
 from miles.backends.training_utils.parallel import get_parallel_state
@@ -206,6 +205,9 @@ def register_cpu_memory(params_dict: dict, transfer_engine) -> dict:
 
 
 def create_transfer_engine():
+    # Lazy: the RDT path imports this module but does not need mooncake.
+    from mooncake.engine import TransferEngine
+
     transfer_engine = TransferEngine()
     local_ip = ray._private.services.get_node_ip_address()
     transfer_engine.initialize(local_ip, "P2PHANDSHAKE", "rdma", "")

miles/backends/sglang_utils/sglang_engine.py

@@ -9,6 +9,9 @@
 import requests
 import sglang_router
 from packaging.version import parse
+from sglang.srt.environ import envs
+from sglang.srt.layers.dp_attention import compute_dp_attention_world_info
+from sglang.srt.ray import get_scheduler_actor_name
 from sglang.srt.server_args import ServerArgs
 from sglang.srt.utils import kill_process_tree
 from urllib3.exceptions import NewConnectionError
@@ -68,7 +71,12 @@ def _get_gpu_uuids(gpu_ids: list[int]) -> list[str | None]:
 
 
 def launch_server_process(server_args: ServerArgs) -> multiprocessing.Process:
-    from sglang.srt.entrypoints.http_server import launch_server
+    if server_args.use_ray:
+        # entrypoints.http_server ignores use_ray and starts mp.Process schedulers,
+        # which have no SchedulerActor for RDT to pull from.
+        from sglang.srt.ray.http_server import launch_server
+    else:
+        from sglang.srt.entrypoints.http_server import launch_server
 
     multiprocessing.set_start_method("spawn", force=True)
     server_args.host = server_args.host.strip("[]")
@@ -132,13 +140,16 @@ def __init__(
         base_gpu_id: int | None = None,
         sglang_overrides: dict | None = None,
         num_gpus_per_engine: int | None = None,
+        pg_bundles: list[int] | None = None,
     ):
         self.args = args
         self.rank = rank
         self.worker_type = worker_type
         self.base_gpu_id = base_gpu_id
         self.sglang_overrides = sglang_overrides or {}
         self.num_gpus_per_engine = num_gpus_per_engine
+        self.pg_bundles = pg_bundles
+        self._scheduler_actors = []
 
     def get_topology_info(self) -> dict:
         """Placement facts for the dashboard timeline. ``base_gpu_id`` is
@@ -244,8 +255,35 @@ def _sanity_check_server_args(actual_server_args, expect_server_args):
         _sanity_check_server_args(actual_server_args, expect_server_args)
 
     def _init_normal(self, server_args_dict):
-        logger.info(f"Launch HttpServerEngineAdapter at: {self.server_host}:{self.server_port}")
-        self.process = launch_server_process(ServerArgs(**server_args_dict))
+        use_rdt = self.args.update_weight_transfer_mode == "rdt"
+        if use_rdt:
+            if self.node_rank != 0:
+                # For a multi-node engine, the node-0 server's RayEngine spawns
+                # the SchedulerActors of ALL ranks (placed cross-node via the
+                # placement group), so non-zero node ranks launch nothing.
+                self.process = None
+                return
+            server_args_dict["use_ray"] = True
+            server_args_dict["enable_rdt_weight_sync"] = True
+            assert self.pg_bundles
+        logger.info(
+            f"Launch HttpServerEngineAdapter at: {self.server_host}:{self.server_port}"
+            f"{' (use_ray=True for RDT)' if use_rdt else ''}"
+        )
+        server_args = ServerArgs(**server_args_dict)
+        if use_rdt:
+            import ray
+
+            placement_group = ray.util.get_current_placement_group()
+            assert placement_group is not None
+            envs.SGLANG_RAY_BUNDLE_INDICES.set(",".join(str(bundle) for bundle in self.pg_bundles))
+            server_args.override(
+                "miles.rdt.ray_context",
+                placement_group=placement_group,
+                ray_runtime_env=dict(ray.get_runtime_context().runtime_env),
+                ray_namespace=ray.get_runtime_context().namespace,
+            )
+        self.process = launch_server_process(server_args)
 
         if self.node_rank == 0 and self.router_ip and self.router_port:
             if parse(sglang_router.__version__) <= parse("0.2.1") or self.args.use_miles_router:
@@ -465,6 +503,9 @@ def flush_cache(self):
     def shutdown(self):
         if self.args.rollout_external:
             return
+        if getattr(self, "process", None) is None:
+            # Non-zero node ranks of an RDT multi-node engine launch no server.
+            return
 
         logger.info(f"Shutdown engine {self.server_host}:{self.server_port}...")
         if self.node_rank == 0:
@@ -514,6 +555,50 @@ def unload_lora_adapter(self, lora_name: str):
             {"lora_name": lora_name},
         )
 
+    def get_scheduler_actors(self) -> list:
+        """Return this engine's SchedulerActor handles (RDT mode, use_ray=True).
+
+        Reconstructs each actor name with ``get_scheduler_actor_name`` from the
+        bundle list handed to RayEngine at launch.
+        """
+        if self._scheduler_actors:
+            return self._scheduler_actors
+
+        assert self.pg_bundles, "get_scheduler_actors requires the RDT bundle list"
+
+        import ray
+
+        tp_size = self.num_gpus_per_engine or self.args.rollout_num_gpus_per_engine
+        enable_dp_attention = bool(getattr(self.args, "sglang_enable_dp_attention", False))
+        dp_size = getattr(self.args, "sglang_dp_size", 1)
+        attn_cp_size = getattr(self.args, "sglang_attn_cp_size", 1)
+        rank0_node_ip = ray.util.get_node_ip_address()
+
+        actors = []
+        for rank, bundle_idx in enumerate(self.pg_bundles):
+            tp_rank = rank % tp_size
+            _, _, dp_rank, _ = compute_dp_attention_world_info(
+                enable_dp_attention, tp_rank, tp_size, dp_size, attn_cp_size
+            )
+            name = get_scheduler_actor_name(
+                rank0_node_ip=rank0_node_ip,
+                dp_rank=dp_rank,
+                pp_rank=rank // tp_size,
+                tp_rank=tp_rank,
+                port=self.server_port,
+                bundle_idx=bundle_idx,
+            )
+            try:
+                actors.append(ray.get_actor(name))
+            except ValueError as e:
+                raise RuntimeError(
+                    f"SchedulerActor {name!r} not found for engine "
+                    f"{self.server_host}:{self.server_port} rank={rank}"
+                ) from e
+
+        self._scheduler_actors = actors
+        return actors
+
     def release_memory_occupation(self, tags: list[str] = None):
         """Release memory occupation. Available tags: weights, kv_cache."""
         self.flush_cache()

miles/backends/training_utils/ci_utils.py

@@ -18,6 +18,10 @@ def check_kl(args: Namespace, log_dict: dict[str, float], step_id: int, accumula
             # LoRA weight conversion (Megatron → HF for SGLang) introduces
             # small floating-point differences, so use a relaxed threshold.
             assert abs(log_dict["train/ppo_kl"]) < 1e-8 and abs(log_dict["train/pg_clipfrac"]) < 1e-10, f"{log_dict=}"
+        elif getattr(args, "update_weight_transfer_mode", None) == "rdt":
+            # RDT's persistent staging/model allocations can perturb otherwise
+            # bitwise-identical forwards without affecting transferred weights.
+            assert abs(log_dict["train/ppo_kl"]) < 1e-8 and abs(log_dict["train/pg_clipfrac"]) < 1e-10, f"{log_dict=}"
         else:
             assert abs(log_dict["train/ppo_kl"]) < 1e-9 and abs(log_dict["train/pg_clipfrac"]) < 1e-10, f"{log_dict=}"
     if accumulated_step_id == 0 and "train/kl_loss" in log_dict and not args.use_rollout_routing_replay:

miles/ray/placement_group.py

@@ -12,6 +12,8 @@
 
 logger = logging.getLogger(__name__)
 
+MILES_RDT_PG_NAME = "miles_rdt_pg"
+
 
 def _select_train_group_class():
     if enable_experimental_ft_trainer():
@@ -48,13 +50,24 @@ def sort_key(x):
     return (node_ip_parts, gpu_id)
 
 
-def _create_placement_group(num_gpus):
+def _create_placement_group(num_gpus, is_rdt: bool = False):
     """Create a placement group with the specified number of GPUs."""
     if num_gpus == 0:
         return None, [], []
 
     bundles = [{"GPU": 1, "CPU": 1} for _ in range(num_gpus)]
-    pg = placement_group(bundles, strategy="PACK")
+    if is_rdt:
+        # Reusing this PG for the rollout SchedulerActors avoids double-booking the
+        # rollout GPUs, but sglang's engine is a separate Ray job, so the PG must be
+        # detached to be schedulable there. Other modes keep the job-scoped lifetime.
+        pg = placement_group(
+            bundles,
+            strategy="PACK",
+            name=MILES_RDT_PG_NAME,
+            lifetime="detached",
+        )
+    else:
+        pg = placement_group(bundles, strategy="PACK")
     num_bundles = len(bundles)
 
     ray.get(pg.ready())
@@ -74,7 +87,16 @@ def _create_placement_group(num_gpus):
         ray.kill(actor)
 
     bundle_infos = [(i, gpu_ids[i][0], gpu_ids[i][1]) for i in range(num_bundles)]
-    sorted_bundle_infos = sorted(bundle_infos, key=sort_key)
+    if is_rdt:
+        # Give the trainer the node PACK filled, so rollout bundles land where GPUs
+        # are still free: RayEngine STRICT_PACKs its SchedulerActors onto the engine
+        # actor's node and deadlocks if nothing there is unreserved.
+        node_bundle_counts: dict = {}
+        for _, node_identifier, _ in bundle_infos:
+            node_bundle_counts[node_identifier] = node_bundle_counts.get(node_identifier, 0) + 1
+        sorted_bundle_infos = sorted(bundle_infos, key=lambda info: (-node_bundle_counts[info[1]], *sort_key(info)))
+    else:
+        sorted_bundle_infos = sorted(bundle_infos, key=sort_key)
     pg_reordered_bundle_indices = [info[0] for info in sorted_bundle_infos]
     # Map from logical index -> physical GPU ID
     pg_reordered_gpu_ids = [gpu_ids[info[0]][1] for info in sorted_bundle_infos]
@@ -111,7 +133,9 @@ def create_placement_groups(args):
     num_gpus, rollout_offset = _get_placement_group_layout(args)
 
     logger.info(f"Creating placement group with {num_gpus} GPUs...")
-    pg, actor_pg_reordered_bundle_indices, actor_pg_reordered_gpu_ids = _create_placement_group(num_gpus)
+    pg, actor_pg_reordered_bundle_indices, actor_pg_reordered_gpu_ids = _create_placement_group(
+        num_gpus, is_rdt=args.update_weight_transfer_mode == "rdt"
+    )
 
     rollout_pg_reordered_bundle_indices = actor_pg_reordered_bundle_indices[rollout_offset:]
     rollout_pg_reordered_gpu_ids = actor_pg_reordered_gpu_ids[rollout_offset:]
@@ -126,13 +150,16 @@ def create_placement_groups(args):
 def allocate_train_group(
     args, num_nodes, num_gpus_per_node, pg, role: str, with_ref: bool, rollout_manager, with_opd_teacher: bool = False
 ):
+    # RDT pins one NIXL/NCCL rank per physical GPU, so it cannot time-share a device
+    # with a colocated rollout the way the fractional reservation allows.
+    num_gpus_per_actor = 1 if args.update_weight_transfer_mode == "rdt" else 0.4
     train_group_cls = _select_train_group_class()
     return train_group_cls(
         args=args,
         num_nodes=num_nodes,
         num_gpus_per_node=num_gpus_per_node,
         pg=pg,
-        num_gpus_per_actor=0.4,
+        num_gpus_per_actor=num_gpus_per_actor,
         role=role,
         with_ref=with_ref,
         rollout_manager=rollout_manager,

miles/ray/rollout/server_group.py

@@ -73,6 +73,9 @@ def start_engines(
 
         pg, reordered_bundle_indices, reordered_gpu_ids = self.pg
 
+        # RDT: reuse miles' PG so RayEngine does not auto-create a second one.
+        rdt_reuse_pg = self.args.update_weight_transfer_mode == "rdt"
+
         RolloutRayActor = ray.remote(SGLangEngine)
 
         new_engines = []
@@ -84,8 +87,11 @@ def start_engines(
                 continue
 
             global_rank = self.rank_offset + i
-            num_gpus = 0.2
-            num_cpus = num_gpus
+            # With PG reuse the SchedulerActors take the bundles' full GPUs
+            # (num_gpus=1 each), so the engine coordinator actor must not
+            # consume any GPU of its bundle.
+            num_gpus = 0 if rdt_reuse_pg else 0.2
+            num_cpus = 0.2
 
             gpu_index = self.gpu_offset + i * num_gpu_per_engine
             base_gpu_id = int(reordered_gpu_ids[gpu_index])
@@ -117,6 +123,15 @@ def start_engines(
             }
             env_vars.update(dumper_utils.get_sglang_env(self.args))
 
+            # The node-0 rank launches the sglang server (RayEngine), which
+            # spawns the SchedulerActors for ALL of the engine's ranks — so it
+            # gets the full engine's bundle list, spanning every node.
+            rdt_pg_kwargs = {}
+            if rdt_reuse_pg and i % self.nodes_per_engine == 0:
+                rdt_pg_kwargs = dict(
+                    pg_bundles=[reordered_bundle_indices[gpu_index + k] for k in range(self.num_gpus_per_engine)],
+                )
+
             rollout_engine = RolloutRayActor.options(
                 num_cpus=num_cpus,
                 num_gpus=num_gpus,
@@ -131,6 +146,7 @@ def start_engines(
                 base_gpu_id=base_gpu_id,
                 sglang_overrides=self.sglang_overrides,
                 num_gpus_per_engine=self.num_gpus_per_engine,
+                **rdt_pg_kwargs,
             )
 
             new_engines.append((global_rank, rollout_engine))

miles/ray/train/actor_factory.py

@@ -38,6 +38,17 @@ def allocate_gpus_for_actor(
         **args.train_env_vars,
     }
 
+    if args.update_weight_transfer_mode == "rdt":
+        # Keep Ray's mask: unmasking makes helper threads default to cuda:0 and NCCL
+        # re-init fails with "Duplicate GPU detected". NIXL uses driver APIs anyway.
+        env_vars["RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES"] = "0"
+        # Every SchedulerActor mapping this trainer's bucket costs a ~520 MiB CUDA
+        # context here, so tightly-packed trainers need the reclaimed fragmentation.
+        env_vars.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+        # Probe the NIXL backend at the real bucket size: EFA accepts small CUDA MRs
+        # through its host bounce pool even when GPUDirect is broken.
+        env_vars.setdefault("MILES_RDT_NIXL_VALIDATE_BYTES", str(args.update_weight_buffer_size))
+
     if source_patcher_config := args.dumper_source_patcher_config_train:
         env_vars["DUMPER_SOURCE_PATCHER_CONFIG"] = source_patcher_config
 
@@ -71,11 +82,16 @@ def allocate_gpus_for_actor(
         actor_impl = FSDPTrainRayActor
 
     ft = args.use_fault_tolerance
-    TrainRayActor = ray.remote(
-        num_gpus=1,
-        runtime_env={"env_vars": env_vars},
-        **(dict(concurrency_groups={"heartbeat_status": 1, "default": 1, "fault_injector": 1}) if ft else {}),
-    )(_with_ft_concurrency_groups(actor_impl) if ft else actor_impl)
+    remote_kwargs = {"num_gpus": 1, "runtime_env": {"env_vars": env_vars}}
+    if ft:
+        remote_kwargs["concurrency_groups"] = {"heartbeat_status": 1, "default": 1, "fault_injector": 1}
+    elif args.update_weight_transfer_mode == "rdt":
+        # update_weights() blocks this actor in ray.get() while Ray's transport
+        # threads serve the NIXL reads of the objects it owns -- one per engine rank
+        # it feeds. At concurrency 1 the blocking call starves them.
+        rdt_tp_size = getattr(args, "rollout_num_gpus_per_engine", 1)
+        remote_kwargs["max_concurrency"] = 1 + rdt_tp_size
+    TrainRayActor = ray.remote(**remote_kwargs)(_with_ft_concurrency_groups(actor_impl) if ft else actor_impl)
 
     # Create worker actors
     actor_handles = []

miles/utils/arguments.py

@@ -819,10 +819,12 @@ def add_rollout_arguments(parser):
             )
             parser.add_argument(
                 "--update-weight-transfer-mode",
-                choices=["broadcast", "p2p", "disk-delta"],
+                choices=["broadcast", "p2p", "disk-delta", "rdt"],
                 default="broadcast",
                 help=(
                     "The method to transfer weights to remote rollout engines during update weight. "
+                    "'broadcast' = NCCL broadcast; 'p2p' = mooncake RDMA write; "
+                    "'rdt' = Ray Direct Transport (NIXL RDMA pull, requires sglang use_ray=True). "
                     "'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."
@@ -3212,14 +3214,14 @@ def miles_validate_args(args):
         args.check_weight_update_equal = True
 
     # always true on offload for colocate at the moment.
-    if args.update_weight_transfer_mode == "p2p":
+    if args.update_weight_transfer_mode in ("p2p", "rdt"):
         assert not args.colocate, (
-            "P2P weight transfer mode is not compatible with --colocate. "
-            "Please use broadcast mode or disable colocate."
+            f"{args.update_weight_transfer_mode} weight transfer mode is not compatible with "
+            "--colocate. Please use broadcast mode or disable colocate."
         )
         assert (
             getattr(args, "prefill_num_servers", None) is None
-        ), "P2P weight transfer mode has not been tested when PD is enabled."
+        ), f"{args.update_weight_transfer_mode} weight transfer mode has not been tested when PD is enabled."
         assert args.lora_rank <= 0, "LoRA weight sync is not supported for p2p (RDMA) weight transfer."
 
     if args.update_weight_transfer_mode == "disk-delta":

requirements.txt

@@ -19,7 +19,7 @@ pylatexenc
 pytest-asyncio
 pyyaml
 qwen_vl_utils # for VLM
-ray[default]
+ray[default]>=2.56.0
 ring_flash_attn; platform_system == "Linux"
 safetensors>=0.8.0  # samples-reply wire codec; malformed-payload exception contract validated on 0.8.0
 sglang-router>=0.2.3