[core][compiled graphs] Unify scheduling for NCCL operation nodes#53111
[core][compiled graphs] Unify scheduling for NCCL operation nodes#53111stephanie-wang merged 40 commits intoray-project:masterfrom
Conversation
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>
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>
|
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? |
|
When one operation node has zero in-degree, it will update the I think there's no ambiguity when one operation node is ready. There might be ambiguity for What about we keep |
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>
python/ray/dag/dag_node_operation.py
Outdated
| # The downstream NCCL read nodes are already updated. | ||
| if node.is_nccl_write: | ||
| continue |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
Ack. Move the check to below, it's equivalent to check if out_node in visited_nodes.
stephanie-wang
left a comment
There was a problem hiding this comment.
Thanks for this contribution, looks much simpler than before! Left a couple suggestions.
Signed-off-by: Weixin Deng <weixin@cs.washington.edu>
Signed-off-by: Weixin Deng <weixin@cs.washington.edu>
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_idxsandpending_sync_idxs. The synchronous nodes denote the P2P send/recv nodes or the collective nodes. The NCCL P2P/collective operation is ready whensync_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
git commit -s) in this PR.scripts/format.shto lint the changes in this PR.method in Tune, I've added it in
doc/source/tune/api/under thecorresponding
.rstfile.