Skip to content
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.

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

[core][compiled graphs]Support reduce scatter and all gather collective for gpu communicator in compiled graph #50037

Open
wants to merge 11 commits into
base: master
Choose a base branch
from

Conversation

anyadontfly
Copy link
Contributor

Why are these changes needed?

Currently we do not have other collective operations except allreduce in Ray Compiled Graphs, we plan to add the other collective operations required in FSDP in the future.

Proposed API:
For reduce scatter:

import ray.experimental.collective as collective

with InputNode() as inp:
    dag = [worker.return_tensor.bind(inp) for worker in workers]
    dag = collective.reducescatter.bind(dag, ReduceOp.SUM)
    dag = MultiOutputNode(dag)

For all agther:

import ray.experimental.collective as collective

with InputNode() as inp:
    dag = [worker.return_tensor.bind(inp) for worker in workers]
    dag = collective.allgather.bind(dag)
    dag = MultiOutputNode(dag)

Related issue number

Meta-issue: #47983

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 :(

@jcotant1 jcotant1 added the core Issues that should be addressed in Ray Core label Jan 23, 2025
@anyadontfly anyadontfly changed the title Support reduce scatter and all gather collective for gpu communicator in compiled graph [core][compiled graphs]Support reduce scatter and all gather collective for gpu communicator in compiled graph Jan 23, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
core Issues that should be addressed in Ray Core
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants