RDT weight sync: GPU->GPU zero copy transfer through SGLang Ray actor backend - #1313
RDT weight sync: GPU->GPU zero copy transfer through SGLang Ray actor backend#1313xyuzh wants to merge 11 commits into
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for RDT/NIXL weight synchronization (point-to-point RDMA) between trainer and rollout engines on Anyscale H100 clusters, adding new entrypoints, job configurations, and READMEs for Qwen3.5-35B-A3B and Qwen3-8B. The code review highlights several important issues: potential bucket overflow crashes and AttributeError exceptions in the new RDT weight sync implementation, incorrect environment variables (PYTHONBUFFERED instead of PYTHONUNBUFFERED) in multiple entrypoint scripts, a duplicate method definition in rollout.py, and potential name-based discovery failures for local loopback addresses in sglang_engine.py. Addressing these issues will improve the robustness, compatibility, and logging behavior of the new synchronization mechanism.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| if not self._is_source or not converted_named_tensors: | ||
| return | ||
|
|
||
| transfer_ready_params, ready_hf_tensors = self._get_transfer_ready_params(converted_named_tensors) |
There was a problem hiding this comment.
If self._shared_param_mapper is None (e.g., when no rollout engines are connected or the transfer plan has no targets for this rank), calling _get_transfer_ready_params will raise an AttributeError when trying to map parameters. Additionally, we should clear converted_named_tensors to prevent memory accumulation and potential leaks.
if not self._is_source or not converted_named_tensors:
return
if self._shared_param_mapper is None:
converted_named_tensors.clear()
return
transfer_ready_params, ready_hf_tensors = self._get_transfer_ready_params(converted_named_tensors)| tp_size = getattr(self.args, "rollout_num_gpus_per_engine", 1) | ||
| host = self.server_host.strip("[]") | ||
| name_prefix = f"sglang_scheduler_node{host}" | ||
|
|
There was a problem hiding this comment.
In single-node or local development environments, self.server_host might be set to 127.0.0.1, 0.0.0.0, or localhost. However, sglang registers named actors using the node's private IP address (resolved via Ray). If we use 127.0.0.1 directly, the name-based discovery will fail to match any actors. We should resolve these local loopback addresses to the actual Ray node IP address.
tp_size = getattr(self.args, "rollout_num_gpus_per_engine", 1)
host = self.server_host.strip("[]")
if host in ("127.0.0.1", "0.0.0.0", "localhost"):
host = ray.util.get_node_ip_address()
name_prefix = f"sglang_scheduler_node{host}"There was a problem hiding this comment.
we use the sglang-miles branch from sglang
11c9524 to
f84e615
Compare
xyuzh
changed the title
b76f454 to
00d331d
Compare
xyuzh
changed the title
3426879 to
ae7f0b7
Compare
| weights_ref = ray.put(tensor_views, _tensor_transport="nixl") | ||
| weight_refs.append(weights_ref) | ||
| for actor in meta.actors: | ||
| futures.append(actor.pull_weights.remote([weights_ref], transfer_ready_params)) |
There was a problem hiding this comment.
sgl-project/sglang#27723
There is a pr at sglang side
|
Also do we have tests? |
26026aa to
cc786c1
Compare
xyuzh
changed the title
I have done testing using the Anyscale job config, would follow up with MILES team for guidance to add test |
cc786c1 to
c67e3cd
Compare
Squash of the 12-commit xinyu/rdt-weight-sync branch onto upstream/main, reconciled with two upstream changes that landed in parallel: - Actor allocation moved out of the now-frozen v1 RayTrainGroup into miles/ray/train/actor_factory.py:allocate_gpus_for_actor. The RDT env-var setup (RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES, PYTORCH_CUDA_ALLOC_CONF, MILES_RDT_NIXL_VALIDATE_BYTES) and the max_concurrency bump for the tensor-transport serve threads now live in the factory (non-FT path). - Upstream's 'disk-delta' transfer mode coexists with 'rdt': merged the --update-weight-transfer-mode choices list and help text. Also: whole-GPU allocation when not colocated (nixl needs a whole GPU), detached/named PG opt-in for no-double-booking, multi-node bundle-count sort, and the new update_weight_from_rdt.py transport.
c67e3cd to
c414b94
Compare
pg_id/pg_bundles were never passed by SGLangEngine's caller, and sglang's RayEngine never reads MILES_RDT_PG_ID/BUNDLES, so the reuse path was dead in both directions. What it did leave behind was a fixed PG name that collides between concurrent jobs and, under MILES_RDT_REUSE_PG=1, a detached PG that survives a crashed job and blocks every later run. The bundle reorder stays: it is what actually keeps sglang's auto-created STRICT_PACK group off the trainer's fully-booked node.
|
Can you support weight transfer for the multi-node SGLang engine? |
This reverts commit 4c5a351. Restoring pending the author's input -- the plumbing is a partially-landed feature, not an accident, and the call on whether to finish it here or split it out is theirs. For that discussion, what the removed commit was based on: - server_group.py never passes pg_id/pg_bundles to SGLangEngine, so the env vars are never written. - sglang RayEngine resolves its PG from server_args.placement_group or get_current_placement_group(); it never reads MILES_RDT_PG_ID/BUNDLES. - What does work today is the bundle reorder, which keeps sglang's auto-created STRICT_PACK group off the trainer's fully-booked node. - Live cost while incomplete: MILES_RDT_PG_NAME is a fixed name that collides between concurrent jobs, and under MILES_RDT_REUSE_PG=1 the PG is detached with no cleanup path, so a crashed job blocks every later run.
|
Could you add an e2e CI test, in a similar format to the p2p weight update test (tests/e2e/megatron/test_qwen3_30B_A3B_p2p.py), for the RDT weight update? thanks! |
The pg_id / pg_bundles parameters on SGLangEngine had no caller, so the rollout SchedulerActors always auto-created a second placement group and double-booked the rollout GPUs. Pass them from ServerGroup.start_engines for the engine's node-0 rank, which is the rank whose RayEngine spawns the SchedulerActors for every rank of the engine, so it gets the whole engine's bundle list rather than just its own node's. The reserved bundles now hold the actors, so the engine coordinator actor no longer takes a fraction of a GPU. A multi-node engine also needs its non-zero node ranks to launch nothing: RayEngine places all the SchedulerActors itself through the placement group, so a second server per extra node would duplicate them. Return early there and skip the matching shutdown, which has no process to wait on. Drop the MILES_RDT_REUSE_PG gate. Reuse is not optional once the engine coordinator holds no GPU, and an auto-created PG cannot co-locate its rank-0 bundle with the engine after miles has reserved every GPU on that node.
yes multi-node SGLang engine supported |
test added, tested CI pass locally |
Hand the rollout PlacementGroup and bundle indices straight to sglang via server_args.override(), plus the actor's runtime env and SGLANG_RAY_BUNDLE_INDICES, rather than smuggling a hex PG id through env vars for the child job to look up. Also switch the replica bootstrap to set_global_server_args_for_scheduler() and force nnodes=1, since the replica is node-local even when the rollout deployment spans nodes. Scope SchedulerActor discovery to the local node IP: ports are unique per host, so (node IP, port) is the real engine key on multi-node rollouts. Track engine-connection freshness in UpdateWeightFromRDT so the actor's reconnect path matches the other weight updaters, and re-register destination buffers when scheduler actors are first resolved so registrations survive repeated pulls. Relax the CI kl assertion for rdt to the same 1e-8 threshold used for LoRA: RDT's persistent staging allocations perturb the forward without changing transferred weights. Require ray>=2.56.0 for the PG/runtime-env plumbing above.
| # 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 |
There was a problem hiding this comment.
This breaks Miles' contract of using the process to control the lifecycle of server engines, which will conflict with the process-based control logics, including cleanup, fault tolerance, etc...
There was a problem hiding this comment.
In non-rdt path, we also return None for the process for non-zero node ranks
https://github.com/xyuzh/miles/blob/xinyu/rdt-weight-sync/miles/backends/sglang_utils/sglang_engine.py#L69-L80
The engine only needs the process on rank 0 node
| {"lora_name": lora_name}, | ||
| ) | ||
|
|
||
| def get_scheduler_actors(self) -> list: |
There was a problem hiding this comment.
This function looks a bit of hardcoded
There was a problem hiding this comment.
addressed hardcoding by adding get_scheduler_actor_name as a helper function
| base_gpu_id: int | None = None, | ||
| sglang_overrides: dict | None = None, | ||
| num_gpus_per_engine: int | None = None, | ||
| placement_group: PlacementGroup | None = None, |
There was a problem hiding this comment.
importing ray here and using placement group does not match the files' scope, these are usually controlled under /ray
There was a problem hiding this comment.
SGLangEngine subclasses miles.ray.ray_actor.RayActor
placement_group removed, ray only imported at method level
SGLang exposes get_scheduler_actor_name, so rebuild each SchedulerActor name from the bundle list handed to RayEngine (accounting for dp-attention rank mapping) and look it up directly, instead of matching node IP / port / tp-rank substrings across all namespaces. Also drop the placement_group constructor arg: the engine actor already runs inside miles' rollout PG, so ray.util.get_current_placement_group() gives the same handle without threading it through.
The mp.Process child calls ray.init(address="auto") without a namespace, so it lands in a fresh anonymous namespace and registers its SchedulerActors there. Discovery from the training side then has to scan every namespace to find them. Forward the launching actor's namespace alongside the placement group and runtime env so the schedulers register where the rest of the deployment already lives.
4 participants
Summary
RDT weight sync shares P2P and nccl distributed weight sync bucketed all-gather + HF conversion pipeline, moving the bucket payload over NIXL (
ray.put(_tensor_transport="nixl")+ RDMA pull) instead of holding a full GPU replica of the rollout model. We have measured faster weight sync speed across models than P2P RDMA and NCCL.CORRECTNESS: Weight equality verified with
--check-weight-update-equal.Performance
RDT / NIXL is the fastest of the three weight-sync transports — a per-rank, zero-copy GPU→GPU RDMA pull with no CPU model replica.
The trainer all-gathers params bucket-by-bucket into a small, reusable fixed-size GPU staging bucket (no full replica); the first
rollout_num_gpusranks are the transfer sources. Each flush, every rollout rank issues a concurrent NIXL RDMA pull (set_target_for_ref→param.data) of its shard over NVLink (same-node) or EFA/LIBFABRIC (cross-node) — no NCCL broadcast, no host bounce.NCCL broadcast vs RDMA P2P vs RDT / NIXL
Three transports, all moving the same trainer → rollout weights:
Per-sync
update_weightson the same model set (seconds, lower is better; bold = fastest per row):NCCL and RDMA P2P (Mooncake TransferEngine) measured on the same harness; RDT / NIXL from miles validation on Anyscale H100 / EFA, byte-equal every sync (
--check-weight-update-equal; RDT is fastest on 4 of 5 models — up to ~2.2× over NCCL — with no CPU replica; the Mooncake RDMA path edges it only on Qwen3-30B, and is itself slower than plain NCCL on two models (GLM-Z1-9B, GLM-4.7-Flash).param.data; the Mooncake RDMA path stages a full model copy in CPU per trainer rank.Design
UpdateWeightFromRDTinheritsDistBucketedWeightUpdateMixin(same bucketed TP/EP all-gather + HF conversion as P2P). Each engine rank is backed by a small fixed-size GPU staging bucket on its source trainer rank; per flush,load_weightswrites the TP-rank-correct sglang shard into bucket views, which are shipped viaray.put(views, _tensor_transport="nixl")and pulled intoparam.dataon eachSchedulerActor.cuMemCreate) segment — VMM memory cannot export legacy CUDA-IPC handles, so UCX'scuda_ipclane silently drops and NIXL falls back to software-emulated RMA over eager TCP fragments (~0.3 GB/s vs ~150 GB/s over NVLink). Bisected toPYTORCH_CUDA_ALLOC_CONF=expandable_segments:Trueon the source; the bucket is now allocated withexpandable_segmentstemporarily forced off.ray.experimental.register_nixl_memory) — Ray otherwise ties registration toObjectReflifetime, so every flush re-pinned GiBs of GPU memory and invalidated the remote-agent cache, forcing a re-handshake.prefer_libfabric_nixl_backend): validates a realistic-size CUDA registration before committing, because EFA accepts small MRs through its host bounce pool even when GPUDirect is broken.Engine discovery is Ray-native:
SchedulerActors register stable names carrying the engine's HTTP port, andget_scheduler_actorsmatches on_port{port}_/_tp{rank}_tokens across namespaces, raising on ambiguous or partial matches. A per-sync[RDT] sync phase breakdownlog splits stage/load/put/submit/pull_wait.Status / requirements
MILES_RDT_REUSE_PG=1): sglang reuses miles' already-reserved rollout bundles instead of auto-creating a second placement group. Validated on GLM-4.5-Air-106B TP8 (40 GPU / 5 nodes).SchedulerActor.pull_weights, RayEngine named actors,enable_engine_info_bootstrap.