alibaba / alibaba/ROLL

[RFC] Integrate Ray Core RDT for Weight Synchronization

Open
#431 3 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
3.4k
Forks
312
Avg merge
1h 2m
Merged PRs (30d)
2

Description

## Background
ROLL supports weight synchronization approaches between training engines and inference engines:
- Collocated (gloo + ccl broadcast)
- Separated (ccl Broadcast)

## Limitations of Legacy ccl Broadcast
1. Single-GPU source constraint: Only trainer rank 0 participates in broadcast. Model weights must fit entirely in this single GPU. For large models (e.g., 70B+ with high TP), this creates memory pressure—rank 0 must hold complete weights before broadcasting.

2. Transfer-trim synchronous cycle: Inference workers receive full weight tensors then trim redundant partitions (e.g., keep only their TP slice). Each batch follows:

```bash
# e.g. Separated
[rank0] broadcast → [all infer workers] recv full tensor → trim → next batch
```

Trimming and receiving cannot pipeline; all workers wait for full tensor before trimming.

## Ray Direct Transport
We can use RDT(Ray Direct Transport) to optimize the above issues.

[Ray Direct Transport (RDT)](https://docs.ray.io/en/latest/ray-core/api/direct-transport.html) enables pull-based P2P without collective group (for one-sided backends):

- `ray.put(tensor, _tensor_transport="nixl"|"yr")` - store tensor reference (NIXL for NVIDIA GPU, YR for Ascend NPU)
- `ray.get(ref)` - pull tensor (zero-copy for NIXL/YR)

## Advantages
- No NCCL group: Avoid init_custom_process_group complexity (master_addr/port allocation, rank assignment, timeout handling)

- RDMA transport + NPU support: RDT integration provides one-sided RDMA for efficient GPU transfers, and through YR backend (via [ray-ascend](https://github.com/Ascend/ray-ascend)) extends shard transfer support to Ascend NPU clusters.

Based on RDT, we can make two-stage improvements for weight synchronization:
- [ ] Firstly, replace the broadcast with `ray.put` and `ray.get` so that workers can asynchronously obtain the latest parameters.
- [ ] Secondly, further optimize the communication link for parameter synchronization, changing the `gather -> broadcast` to P2P data transmission between the **train worker TP** and **infer worker TP**.

Contributor guide

No contributing guide indexed for this repository

Research direction

No file, test, or entry point is named. Start by locating the existing collocated and separated ccl broadcast implementations, then read the linked Ray Direct Transport documentation and compare the proposed ray.put/ray.get flow with the current synchronization path. Done requires a decided scope and an implemented, tested RDT-based synchronization approach.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
distributed-systems
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Needs clarification
Newbie friendliness
32/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.