8000 [core] Support broadcast and reduce collective for compiled graphs by jeffreyjeffreywang · Pull Request #53625 · ray-project/ray · GitHub
[go: up one dir, main page]
More Web Proxy on the site http://driver.im/
Skip to content

[core] Support broadcast and reduce collective for compiled graphs #53625

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Open
wants to merge 5 commits into
base: master
Choose a base branch
from
Open

[core] Support broadcast and reduce collective for compiled graphs #53625

wants to merge 5 commits into from

Conversation

jeffreyjeffreywang
Copy link
Contributor
@jeffreyjeffreywang jeffreyjeffreywang commented Jun 6, 2025 ̶ 8000 6;
edited
Loading

Why are these changes needed?

allreduce, reducescatter, allgather collectives are supported in compiled graphs. As part of the effort (#47983) of supporting collective operations in compiled graphs, we aim to also support broadcast and reduce.

Sample usage:

computes = [
    worker.send_tensor.bind(tensor) for worker, tensor in zip(workers, inp)
]

# Broadcast (providing root node to broadcast tensors from)
collectives = collective.broadcast.bind(computes[0], computes)

# Reduce (providing root node to reduce tensors to and the reduction operation)
collectives = collective.reduce.bind(computes[0], computes, op=reduce_op)

recvs = [
    worker.recv_tensor.bind(collective)
    for worker, collective in zip(workers, collectives)
]
dag = MultiOutputNode(recvs)

Related issue number

Resolves #49325 and #49324

Checks

  • I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
  • I've run scripts/format.sh to lint the changes in this PR.
  • I've included any doc changes needed for https://docs.ray.io/en/master/.
    • I've added any new APIs to the API Reference. For example, if I added a
      method in Tune, I've added it in doc/source/tune/api/ under the
      corresponding .rst file.
  • I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
  • Testing Strategy
    • Unit tests
    • Release tests
    • This PR is not tested :(

@codope codope requested a review from stephanie-wang June 7, 2025 00:42
@jeffreyjeffreywang jeffreyjeffreywang requested review from a team as code owners June 8, 2025 20:52
@stephanie-wang stephanie-wang self-assigned this Jun 10, 2025
jeffreyjeffreywang
jeffreyjeffreywang commented Jun 16, 2025
View reviewed changes
jeffreyjeffreywang added 4 commits June 17, 2025 22:33
Support broadcast collective in compiled graph
5ac3b91 
Signed-off-by: jeffreyjeffreywang <jeffjeffreywang@gmail.com>
Fix CI test errors
bd81159 
Signed-off-by: jeffreyjeffreywang <jeffjeffreywang@gmail.com>
Fix error
170c79b 
Signed-off-by: jeffreyjeffreywang <jeffjeffreywang@gmail.com>
Support reduce collective for compiled graph
dcef0c2 
Signed-off-by: jeffreyjeffreywang <jeffjeffreywang@gmail.com>
Remove merge duplicates
6aa3318 
Signed-off-by: jeffreyjeffreywang <jeffjeffreywang@gmail.com>
@jeffreyjeffreywang
Copy link
Contributor Author

Resolving messed-up commits...

jeffreyjeffreywang
jeffreyjeffreywang commented Jun 18, 2025
View reviewed changes
python/ray/dag/collective_node.py
@@ -16,6 +16,8 @@
AllGatherOp,
AllReduceOp,
ReduceScatterOp,
BroadcastOp,
ReduceOperation,
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not thrilled about this naming. Using ReduceOp for the collective could conflict with actual reduction ops like min, max, and sum (code). Open to better naming ideas.

@jeffreyjeffreywang jeffreyjeffreywang changed the title [core] Support broadcast collective for compiled graphs [core] Support broadcast and reduce collective for compiled graphs Jun 18, 2025
@aslonnie aslonnie removed request for a team June 24, 2025 07:04
@aslonnie aslonnie removed request for a team and aslonnie June 24, 2025 07:04
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@stephanie-wang stephanie-wang Awaiting requested review from stephanie-wang

@pcmoritz pcmoritz Awaiting requested review from pcmoritz

6CF1

@kevin85421 kevin85421 Awaiting requested review from kevin85421

@edoakes edoakes Awaiting requested review from edoakes

@richardliaw richardliaw Awaiting requested review from richardliaw

@raulchen raulchen Awaiting requested review from raulchen

At least 1 approving review is required to merge this pull request.

Assignees

@stephanie-wang stephanie-wang

Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.

[core][compiled graph] Support one-to-all collective ops (e.g. broadcast)
2 participants
@jeffreyjeffreywang @stephanie-wang
0