RDT weight sync: GPU->GPU zero copy transfer through SGLang Ray actor backend - #1313
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
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.
launch_engine runs inside a child actor of the miles job, so the SchedulerActor handles come back directly from start() instead of being rediscovered by name, and the rollout PG no longer needs the detached lifetime or the miles_rdt_pg name.
The facade set it in its own process, but launch_engine runs in the SGLangServerActor child, where a parent's os.environ mutation never arrives. Pass the bundle list through start() and set it there.
sglang's launch_engine now takes placement_group directly, so SGLangServerActor resolves its own current placement group and hands it over, replacing the miles.rdt.ray_context override on ServerArgs.
update_weight_from_rdt lives in the megatron backend, so fail fast in arg validation instead of at weight-sync time.
RDT reserves a full GPU per trainer rank, while shared PPO places actor and critic ranks in the same bundles and leaves the critic unschedulable. Reject that combination during argument validation and cover the valid and invalid RDT configurations.
| 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": |
There was a problem hiding this comment.
Update: I traced this through Ray 2.56 and retract this P1. This elif is intentional: defining any named concurrency group makes Ray create a separate executor for the built-in default group even when its concurrency remains 1 (manager, condition). The unannotated update_weights() in the FT actor therefore already runs off the core-worker main thread; the non-FT path needs max_concurrency > 1 only to force that same off-main-thread executor.
Also, "default" in concurrency_groups is only a named group and has no bound method here, so increasing it would not affect update_weights(). NIXL is one-sided, and the required SGLang PR #27723 gives the destination SchedulerActor max_concurrency=2 so pull_weights can run beside run_event_loop. Setting actor-level max_concurrency=1+tp for FT would instead allow ordinary trainer RPCs to overlap and potentially race. No behavior change is justified without an FT+RDT E2E hang.
There was a problem hiding this comment.
kept the concurrency setup as-is.
| self, | ||
| rollout_engines: Sequence[ActorHandle], | ||
| rollout_engine_lock: ActorHandle | None = None, | ||
| engine_gpu_counts: Sequence[int] | None = None, |
There was a problem hiding this comment.
[P2] Build RDT targets from the per-engine GPU counts
--sglang-config allows each server group to override num_gpus_per_engine, and RolloutManager passes the resolved engine_gpu_counts here, but this method ignores them and uses a RemoteTransferPlan built from the global args.rollout_num_gpus_per_engine. For example, with four rollout GPUs, a TP=2 CLI default, and groups producing counts [1, 1, 2], the plan creates (0,1) and (1,1) (out-of-range scheduler ranks) and never targets engine 2, so connection fails at engine_actors[t.engine_rank] or leaves that engine stale. Build or update the plan from engine_gpu_counts, or explicitly reject heterogeneous RDT layouts.
There was a problem hiding this comment.
Added validation: connect_rollout_engines() now raises a ValueError if engine_gpu_counts is heterogeneous (any count != --rollout-num-gpus-per-engine). This fails safely before touching any connection state.
Allowed cases: Uniform, None, and empty counts still pass unchanged.
Deferred work: The ideal fix for heterogeneous layouts is to build the transfer plan per-engine. Since there is no RDT use case for this yet, I'm deferring it until one actually shows up.
| shape, dtype, nbytes = self.param_specs[name] | ||
| itemsize = torch.empty((), dtype=dtype).element_size() | ||
| offset = ((offset + itemsize - 1) // itemsize) * itemsize | ||
| assert offset + nbytes <= capacity, ( |
There was a problem hiding this comment.
[P2] Validate or grow the RDT bucket before the first sync
The shared bucketizer deliberately lets an oversized first update unit through when the current bucket is empty, but stage() requires every complete destination parameter to fit in this fixed allocation. --update-weight-buffer-size still defaults to 512 MiB, while a legal TP=1 Qwen3-30B-A3B rollout has a bf16 embed_tokens parameter of 622,329,856 bytes (~0.58 GiB); the new E2E also raises the buffer to 1 GiB because of this parameter. Selecting rdt with otherwise valid/default arguments therefore reaches this assertion on the first weight update. Once param_specs are available in connect_rollout_engines(), allocate at least max(configured_size, max_param_nbytes) or reject the configuration there with a clear error.
There was a problem hiding this comment.
Dynamic bucket sizing: Each engine-rank bucket is now allocated at max(--update-weight-buffer-size, largest destination param nbytes) based on that rank's param_specs. Added a warning when the buffer needs to grow.
Better CI coverage: Removed the hand-tuned 1 GiB override from the 30B e2e test. It now runs at the 512 MiB default, which forces the new growth path to trigger on the 0.58 GiB embed_tokens in every CI run.
| **args.train_env_vars, | ||
| } | ||
|
|
||
| if args.update_weight_transfer_mode == "rdt": |
There was a problem hiding this comment.
[P2] Add the new transfer-mode field to fast-test defaults
This unconditional access already makes pytest -q tests/fast/ray/test_actor_factory.py fail at this line because its SimpleNamespace has no update_weight_transfer_mode. The shared tests/fast/ray/rollout/conftest.py::make_args() also lacks the field, while ServerGroup.start_engines() now reads it directly and raises the same AttributeError before actor creation. Add update_weight_transfer_mode="broadcast" to both fixtures, matching the parser default, so existing non-RDT fast tests still exercise their intended paths.
There was a problem hiding this comment.
Updated test configs to match the parser default. Set update_weight_transfer_mode="broadcast" in:
-
tests/fast/ray/rollout/conftest.py::make_args()
-
The SimpleNamespace in test_actor_factory.py
sglang_engine imports this module unconditionally, but only the RDT path has (or needs) sglang.srt.ray.http_server and the server entrypoints it drags in — resolve them when start() runs instead.
connect_rollout_engines now rejects heterogeneous engine GPU counts instead of silently planning wrong transfers, and the fixed GPU bucket grows (with a warning) when the largest destination parameter exceeds --update-weight-buffer-size, so the 30B e2e test no longer needs to hand-tune the buffer size. Fast-test arg fixtures gain the update_weight_transfer_mode field the code now reads.
|
Two failed CI not related to the PR |
The mkdir and hf download don't need a GPU slot; run them on the CPU path like test_external_rollout does.
The bucket auto-grows past --update-weight-buffer-size when a param needs it, but that path warns; size it explicitly so the test runs clean.
The GPU bucket is lifetime-registered with NIXL, so growing it per flush is off the table; instead pack the ready params into rounds that fit the fixed bucket and pull each round before staging the next. _get_transfer_ready_params returns name->shards so a round carries exactly its own HF tensors, and _staging_span keeps the packer and stage() computing alignment identically. The 30B e2e test drops the explicit --update-weight-buffer-size again: capacity no longer limits a flush.
|
Need to pin sglang version to pass the final CI |
|
kicked off ci |
5 participants
ci-sglang-pr: #27723
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.