- Notifications
You must be signed in to change notification settings - Fork6.2k
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 ourterms of service andprivacy statement. We’ll occasionally send you account related emails.
Already on GitHub?Sign in to your account
(WIP) [ADAG] Support dag.experimental_compile(_custom_nccl_group= nccl_group) in aDAG#47987
base:master
Are you sure you want to change the base?
Conversation
Signed-off-by: Yuhan Ruan <andyubryh@gmail.com>
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others.Learn more.
Initial pass. Left some comments.
| @@ -322,8 +322,6 @@ def test_torch_tensor_custom_comm(ray_start_regular): | |||
| sender = actor_cls.remote() | |||
| receiver = actor_cls.remote() | |||
| from cupy.cuda import nccl | |||
| class TestNcclGroup(GPUCommunicator): | |||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others.Learn more.
@stephanie-wang Rui movesTestNcclGroup back inside each test in the patch PR. This is not necessary though.
Signed-off-by: Yuhan Ruan <andyubryh@gmail.com>
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others.Learn more.
@stephanie-wang Do we still need to create a default NCCL group if the specified custom NCCL group does not contain all the P2P actors?
Signed-off-by: Yuhan Ruan <andyubryh@gmail.com>
I think it's better to throw an error. Better to do as little automagically as possible. |
python/ray/dag/compiled_dag_node.py Outdated
| # The NCCL group ID for P2P send/recv operations. | ||
| self._nccl_group_id_p2p: Optional[str] = None | ||
| # The `custom_nccl_group` specified in `experimental_compile`. | ||
| self._custom_nccl_group_p2p: Optional[GPUCommunicator] = custom_nccl_group |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others.Learn more.
| self._custom_nccl_group_p2p:Optional[GPUCommunicator]=custom_nccl_group | |
| self._default_nccl_group_p2p:Optional[GPUCommunicator]=custom_nccl_group |
Is this ever different from the group ID in the next line?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others.Learn more.
There were 2 cases where they may be different.
- If a default NCCL group is specified in `experimental_compile`, but it is never used in the DAG (e.g., there is no P2P NCCL send/recvs in the DAG), the `default_nccl_group_id_p2p` will be None.
- If no default NCCL group is specified in `experimental_compile`, but there are some NCCL send/recv operations in the DAG whose `transport="nccl"`, either a NCCL group specified somewhere else in the DAG is reused for these P2P operations, or a new `_NcclGroup` is created. In this case, `default_nccl_group_id_p2p` is set to the id of that (reused/created) NCCL group.
To make it simpler, I've set it to thedefault_nccl_group_p2p's ID, if one was provided and used, and None otherwise.
In the case where a default NCCL group is provided but never used (and thus never initialized) in the DAG, the ID is still set to None.
python/ray/dag/compiled_dag_node.py Outdated
| @@ -765,7 +765,18 @@ def _preprocess(self) -> None: | |||
| self.input_task_idx, self.output_task_idx = None, None | |||
| self.actor_task_count.clear() | |||
| nccl_actors_p2p: Set["ray.actor.ActorHandle"] = set() | |||
| # Actors and DAG nodes involved in NCCL P2P send/recv with `transport=nccl`, | |||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others.Learn more.
The logic in this block of code is getting quite complicated. Let's try to clean it up a bit:
- Separate it out into a helper function.
- Ideally, we do it in these steps and none of them should be interleaved:
- Collect mapping of set(actors) <> NCCL group specified (None for "default").
- Validate all groups. I.e. if NCCL group was specified, check that actors are actually in that group
- Initialize groups if needed.
- Set type hints on DAG nodes.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others.Learn more.
Factored steps 2-4 out into a helper function. Kept step 1 in_preprocess.
For step 1, I created:
- A mapping of each NCCL group specified (None for "default") -> sender DAG nodes that specify it as their
transport. - A mapping of each NCCL group specified (None for "default") -> sender and receiver DAG nodes that use it as their
transport, where the receivers are downstream DAG nodes of the senders.
The first mapping is used to set the NCCL group ids in the type hints of the senders. The second mapping corresponds to the set of all actors that use the NCCL group and is used for validation. Maybe there is a simpler way to do this?
Signed-off-by: Yuhan Ruan <andyubryh@gmail.com>
Signed-off-by: Yuhan Ruan <andyubryh@gmail.com>
Signed-off-by: Yuhan Ruan <andyubryh@gmail.com>
Signed-off-by: Yuhan Ruan <andyubryh@gmail.com>
Signed-off-by: Yuhan Ruan <andyubryh@gmail.com>
@anyscalesam This PR is close to a final review. Could you add a go badge for running all tests? Thanks! |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others.Learn more.
Looks great! Did a cursory pass for now (a question and a comment you can ignore), will do a more in-depth one when we meet later today
Signed-off-by: Yuhan Ruan <andyubryh@gmail.com>
This pull request has been automatically marked as stale because it has not had recent activity. It will be closed in 14 days if no further activity occurs. Thank you for your contributions.
|
Why are these changes needed?
This supports the API to specify a custom NCCL group for P2P communication when compiling an accelerated DAG. Additionally, this PR allows the user to specify more than 1 custom NCCL group for the same compiled DAG. Specifically, the following code is supported.
The custom NCCL group provided for
dag.experimental_compilewill be used as the default P2P communicator forcompiled_dag. If a dag node's type hint specifiestransport="nccl", the default communicator will be used for its P2P NCCL communication. If a dag node's type hint specifiestransport=another_custom_nccl_group, that custom NCCL group will be used instead.Requirements:
transport=custom_nccl_group, the sender and receiver involved must be members in that custom group, otherwise an error is raised.transport="nccl"and a default custom NCCL group is specified indag.experimental_compile, the default group must contain the sender and receiver, otherwise an error is raised.In the future, instead of raising an error for Requirement 2, we can consider constructing a
_NcclGroupfor the P2P actors that are not members of the group specified inexperimental_compile.Related issue number
Closes#47540.
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.