feat(atom): direct trainer-to-rollout weight transfer over RDMA - #33
Draft
i-chaochen wants to merge 2 commits into
Draft
i-chaochen wants to merge 2 commits into
i-chaochen wants to merge 2 commits into
Conversation
`ATOMRayServer` could only reach workers through the fixed utility commands ATOM happens to register, so anything else -- capability queries, and the RDMA weight-group methods still to come -- had no route at all. - `collective_rpc()` and `get_capabilities()` passthroughs on the server, plus replica-wide versions on the manager. Both hop through an executor: the engine call blocks on a queue, and running it on the actor's event loop would stall `generate` for a whole weight sync, the way `_generate_blocking` already avoids. - `wait_for_requests_to_drain` was `return True`, making every caller's drain barrier a no-op -- a weight sync could begin while responses were still in flight, producing rollouts from a half-updated model. ATOM exposes `is_finished()`, so the barrier is now real, with a timeout and a warning when it expires. The two weight-sync call sites deliberately keep using `broadcast_utility_command_sync`. `EngineUtilityHandler.process_queue` clears `_is_rl_weights_offloaded` by matching the command name against `update_weights_shm`/`update_weights_ipc` with `is_last=True`; routing them through `collective_rpc` would rename the command, leave the engine flagged as offloaded, and make the busy loop keep skipping model execution -- a silent rollout failure. Those paths already work and have their own handlers; the generic route is for methods that have none. Requires the matching ATOM change (generic collective_rpc + capability discovery). Co-authored-by: Cursor <cursoragent@cursor.com>
Completes the ATOM RDMA path. The receiver and the trainer's sender both existed; nothing connected them, so `weight_sync.backend: rdma` with `generation_backend: atom` failed at "RDMA weight group is not initialized". ATOM rollout fell back to safetensors on shared storage. `_sync_weights_rdma` in the trainer is already backend-agnostic, so this is mostly a matter of `ATOMReplicaManager` implementing the same interface `VLLMReplicaManager` does -- a test asserts the three signatures match, since a drift shows up as an AttributeError mid-sync. On the server: `rdma_preflight`, `init_rdma_weight_group`, `receive_weights_rdma`, `destroy_rdma_weight_group`, all over `collective_rpc` so every rank is reached and every rank's reply is collected. On the manager: group construction, arming the receivers, teardown, and the rank layout. Rank 0 is the trainer; replica r takes the block at `1 + r * ranks_per_replica`. That stride is `tp * dp`, not `tp`: a replica running DP internally is several engines behind one actor handle and each of their TP ranks joins separately. Striding by `tp` alone would start replica 1 inside replica 0's block and the rendezvous would never complete -- all ranks sit until the timeout with nothing naming the cause. The rank arithmetic is tested as pure arithmetic, including that every rank is claimed exactly once. Two checks run before the rendezvous rather than after, because afterwards a failure is a hang rather than a message: `rdma_preflight` on every participant, and a capability check that every rank reports `rdma_weight_receive`. The latter uses ATOM's general capability negotiation instead of a bespoke handshake. Trainer-side FP8 (`weight_sync.fp8_quantization_location=trainer`) is rejected with NotImplementedError rather than ignored. ATOM's receive path is BF16-only, and silently dropping the flag would leave the trainer quantising while ATOM expected BF16 -- garbage output, no error. It defaults off, so the BF16 path is unaffected. Receivers are armed with un-awaited refs: all of them and the sender must be inside the broadcast at once, and joining them in turn would park the first receiver waiting for a sender waiting for this call to return. Requires the two ATOM-side changes (generic collective_rpc, and the transactional RDMA receiver). 14 new tests. LumenRL unit suite 609 -> 623 passed, errors and skips unchanged. Not yet run on hardware: the 9-rank sync needs two GPU nodes at once. Co-authored-by: Cursor <cursoragent@cursor.com>
i-chaochen
marked this pull request as draft
September 19, 2026 18:00
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
Completes the ATOM RDMA path. The receiver and the trainer's sender both
existed; nothing connected them, so
weight_sync.backend: rdmawithgeneration_backend: atomfailed at "RDMA weight group is not initialized" andATOM rollout fell back to safetensors on shared storage.
Requires ATOM PRs ROCm/ATOM#2297 and ROCm/ATOM#2298, and stacks on #32
Technical Details
_sync_weights_rdmain the trainer is already backend-agnostic, so this ismostly
ATOMReplicaManagerimplementing the interfaceVLLMReplicaManageralready has. A test asserts the three signatures match, since a drift shows up
as an AttributeError mid-sync rather than at import.
Server:
rdma_preflight,init_rdma_weight_group,receive_weights_rdma,destroy_rdma_weight_group, all overcollective_rpcso every rank is reachedand every reply collected.
Manager: group construction, arming receivers, teardown, and the rank layout.
Rank 0 is the trainer; replica
rtakes the block at1 + r * ranks_per_replica. That stride istp * dp, nottp-- a replicarunning DP internally is several engines behind one actor handle, and each of
their TP ranks joins separately. Striding by
tpalone would start replica 1inside replica 0's block and the rendezvous would never complete, leaving all
ranks parked until the timeout with nothing naming the cause. The arithmetic is
tested directly, including that every rank is claimed exactly once.
Two checks run before the rendezvous rather than after, because afterwards a
failure is a hang and not a message:
rdma_preflighton every participant, anda capability check that every rank reports
rdma_weight_receive-- using ATOM'sgeneral capability negotiation rather than a bespoke handshake.
Trainer-side FP8 is rejected with NotImplementedError rather than ignored.
ATOM's receive path is BF16-only, and silently dropping the flag would leave the
trainer quantising while ATOM expected BF16 -- garbage output, no error. It
defaults off, so the BF16 path is unaffected.
Receivers are armed with un-awaited refs: all of them and the sender must be
inside the broadcast at once, and joining them in turn would park the first
receiver waiting for a sender waiting for the call to return.
Test Plan
14 new tests covering the rank layout across six topologies, that every rank is
claimed exactly once, that a
tp-only stride is wrong under DP, theinit/receive/teardown guards, signature parity with the vLLM manager, and the
FP8 rejection. Rank arithmetic is tested as arithmetic, away from Ray.
Hardware validation is the real gate and has not run: a 9-rank sync on two GPU
nodes, confirming
NCCL INFO Using network IBandNET/IB/.../GDRDMAin thelogs, since a config gap degrades silently to TCP and throughput alone will not
show it.
Test Result
black --checkclean;ruffno findings in the added regionsNot yet run on hardware -- two GPU nodes have not been available
simultaneously. This should not merge as "RDMA works" until that run is green.