Skip to content

[core][compiled graphs] Unify scheduling for NCCL operation nodes#53111

Merged
stephanie-wang merged 40 commits intoray-project:masterfrom
dengwxn:sync-op-pr
May 29, 2025
Merged

[core][compiled graphs] Unify scheduling for NCCL operation nodes#53111
stephanie-wang merged 40 commits intoray-project:masterfrom
dengwxn:sync-op-pr

Conversation

@dengwxn
Copy link
Copy Markdown

@dengwxn dengwxn commented May 17, 2025

Why are these changes needed?

This PR unifies the scheduling implementation for the NCCL P2P and collective operation nodes. The logic remains the same: (1) P2P case: When a NCCL send node is selected, its downstream NCCL recv nodes are also selected; (2) Collective case: When a NCCL collective node is selected, its corresponding NCCL collective nodes are also selected. Previously, the NCCL P2P case was implemented by selecting the recv nodes if a send node is detected, and the NCCL collective case was implemented by maintaining a set of pending collective nodes.

We unify the implementation for both cases. Concretely, they both maintain a set of (pending) synchronous nodes named sync_idxs and pending_sync_idxs. The synchronous nodes denote the P2P send/recv nodes or the collective nodes. The NCCL P2P/collective operation is ready when sync_idxs == pending_sync_idxs.

Test cases are updated to reflect the use of synchronous nodes for both NCCL P2P and collective nodes.

Related issue number

This PR is a follow-up of #53007. They are parts of #48649 planning to be merged incrementally.

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

Weixin Deng added 24 commits May 12, 2025 23:25
Signed-off-by: Weixin Deng <weixin@cs.washington.edu>
Signed-off-by: Weixin Deng <weixin@cs.washington.edu>
Signed-off-by: Weixin Deng <weixin@cs.washington.edu>
Signed-off-by: Weixin Deng <weixin@cs.washington.edu>
Signed-off-by: Weixin Deng <weixin@cs.washington.edu>
Signed-off-by: Weixin Deng <weixin@cs.washington.edu>
Signed-off-by: Weixin Deng <weixin@cs.washington.edu>
Signed-off-by: Weixin Deng <weixin@cs.washington.edu>
Signed-off-by: Weixin Deng <weixin@cs.washington.edu>
Signed-off-by: Weixin Deng <weixin@cs.washington.edu>
Signed-off-by: Weixin Deng <weixin@cs.washington.edu>
Signed-off-by: Weixin Deng <weixin@cs.washington.edu>
Signed-off-by: Weixin Deng <weixin@cs.washington.edu>
Signed-off-by: Weixin Deng <weixin@cs.washington.edu>
Signed-off-by: Weixin Deng <weixin@cs.washington.edu>
Signed-off-by: Weixin Deng <weixin@cs.washington.edu>
Signed-off-by: Weixin Deng <weixin@cs.washington.edu>
Signed-off-by: Weixin Deng <weixin@cs.washington.edu>
Signed-off-by: Weixin Deng <weixin@cs.washington.edu>
Signed-off-by: Weixin Deng <weixin@cs.washington.edu>
Signed-off-by: Weixin Deng <weixin@cs.washington.edu>
Signed-off-by: Weixin Deng <weixin@cs.washington.edu>
Signed-off-by: Weixin Deng <weixin@cs.washington.edu>
@dengwxn
Copy link
Copy Markdown
Author

dengwxn commented May 17, 2025

cc @stephanie-wang

Weixin Deng added 5 commits May 16, 2025 18:31
Signed-off-by: Weixin Deng <weixin@cs.washington.edu>
Signed-off-by: Weixin Deng <weixin@cs.washington.edu>
Signed-off-by: Weixin Deng <weixin@cs.washington.edu>
Signed-off-by: Weixin Deng <weixin@cs.washington.edu>
Weixin Deng added 5 commits May 20, 2025 16:21
Signed-off-by: Weixin Deng <weixin@cs.washington.edu>
Signed-off-by: Weixin Deng <weixin@cs.washington.edu>
Signed-off-by: Weixin Deng <weixin@cs.washington.edu>
Signed-off-by: Weixin Deng <weixin@cs.washington.edu>
Signed-off-by: Weixin Deng <weixin@cs.washington.edu>
@stephanie-wang
Copy link
Copy Markdown
Contributor

I'm a bit confused by the logic for when to mark the synchronous nodes as ready. Right now the logic seems to be that if one of the nodes becomes ready, then they should all become ready. But isn't it that none of them are ready until all of them become ready?

@dengwxn
Copy link
Copy Markdown
Author

dengwxn commented May 21, 2025

When one operation node has zero in-degree, it will update the ready_sync_idxs for itself and other operation nodes in the same NCCL operation.
Then one operation node is ready (self.is_ready() returns true) when sync_idxs == ready_sync_idxs.

I think there's no ambiguity when one operation node is ready. There might be ambiguity for ready_sync_idxs. What do you think is a better name for the ready_sync_idxs?

What about we keep self.is_ready() and rename ready_sync_idxs into something like pending_sync_idxs?

Weixin Deng added 3 commits May 27, 2025 16:24
Signed-off-by: Weixin Deng <weixin@cs.washington.edu>
Signed-off-by: Weixin Deng <weixin@cs.washington.edu>
Signed-off-by: Weixin Deng <weixin@cs.washington.edu>
Comment on lines +743 to +745
# The downstream NCCL read nodes are already updated.
if node.is_nccl_write:
continue
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Is it possible to remove this special case for is_nccl_write? Seems it could work to just treat nccl write nodes as any other node.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

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

Ack. Move the check to below, it's equivalent to check if out_node in visited_nodes.

Copy link
Copy Markdown
Contributor

@stephanie-wang stephanie-wang left a comment

Choose a reason for hiding this comment

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

Thanks for this contribution, looks much simpler than before! Left a couple suggestions.

Weixin Deng added 3 commits May 27, 2025 17:44
Signed-off-by: Weixin Deng <weixin@cs.washington.edu>
Signed-off-by: Weixin Deng <weixin@cs.washington.edu>
@stephanie-wang stephanie-wang added the go add ONLY when ready to merge, run all tests label May 28, 2025
@stephanie-wang stephanie-wang merged commit aec2791 into ray-project:master May 29, 2025
6 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

go add ONLY when ready to merge, run all tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants