-
Notifications
You must be signed in to change notification settings - Fork 5.5k
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][experimental] Accelerated DAG NCCL-based p2p channels for torch.Tensors #45092
Merged
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Signed-off-by: Stephanie Wang <swang@cs.berkeley.edu>
Signed-off-by: Stephanie Wang <swang@cs.berkeley.edu>
Signed-off-by: Your Name <you@example.com>
Signed-off-by: Stephanie Wang <swang@cs.berkeley.edu>
Signed-off-by: Stephanie Wang <swang@cs.berkeley.edu>
Signed-off-by: Your Name <you@example.com>
Signed-off-by: Your Name <you@example.com>
Signed-off-by: Your Name <you@example.com>
… dag-gpu-channels
Signed-off-by: Your Name <you@example.com>
Signed-off-by: Your Name <you@example.com>
Signed-off-by: Stephanie Wang <swang@cs.berkeley.edu>
Is this ready for another pass? |
Signed-off-by: Stephanie Wang <swang@cs.berkeley.edu>
Yes, it's ready! I removed the multi-GPU CI tests, will merge these in separately after #45253. |
rkooo567
approved these changes
May 10, 2024
stephanie-wang
added a commit
that referenced
this pull request
May 13, 2024
stephanie-wang
added a commit
that referenced
this pull request
May 14, 2024
Fix another broken test from #45092. Signed-off-by: Stephanie Wang <swang@cs.berkeley.edu>
8 tasks
ryanaoleary
pushed a commit
to ryanaoleary/ray
that referenced
this pull request
Jun 6, 2024
…h.Tensors (ray-project#45092) ## Why are these changes needed? This adds a NCCL-based transport option for torch.tensors. Here is an example of the API: ```python with InputNode() as inp: dag = sender.send.bind(inp) dag = dag.with_type_hint(TorchTensorType(SHAPE, DTYPE, transport="nccl")) dag = receiver.recv.bind(dag) compiled_dag = dag.experimental_compile() ``` When `transport="nccl"` is specified, upon compile(), Ray will initialize a NCCL group with the actors involved. The reading actor(s) will `recv` on the NCCL communicator instead of reading from the default shared-memory channel. This PR also refactors channel types so that we now create `ChannelInterfaces` based on the type hints that appear in the DAG, either a `TorchTensorType` or the default `SharedMemoryType`. Current limitations: - p2p only, no collectives - Synchronizes CUDA stream after receiving data. This is because kernels following the NCCL op have no guarantee that the op succeeded, so it is not safe to read the received buffer unless we know that the op succeeded. - Shape and dtype of the tensor must be declared at compile time. --------- Signed-off-by: Stephanie Wang <swang@cs.berkeley.edu> Co-authored-by: SangBin Cho <rkooo567@gmail.com> Signed-off-by: Ryan O'Leary <ryanaoleary@google.com>
ryanaoleary
pushed a commit
to ryanaoleary/ray
that referenced
this pull request
Jun 6, 2024
Fixes test that was broken due a hidden merge conflict between ray-project#45092 and ray-project#45119. Related issue number Closes ray-project#45264. Signed-off-by: Stephanie Wang <swang@cs.berkeley.edu> Signed-off-by: Ryan O'Leary <ryanaoleary@google.com>
ryanaoleary
pushed a commit
to ryanaoleary/ray
that referenced
this pull request
Jun 6, 2024
Fix another broken test from ray-project#45092. Signed-off-by: Stephanie Wang <swang@cs.berkeley.edu> Signed-off-by: Ryan O'Leary <ryanaoleary@google.com>
ryanaoleary
pushed a commit
to ryanaoleary/ray
that referenced
this pull request
Jun 6, 2024
…h.Tensors (ray-project#45092) ## Why are these changes needed? This adds a NCCL-based transport option for torch.tensors. Here is an example of the API: ```python with InputNode() as inp: dag = sender.send.bind(inp) dag = dag.with_type_hint(TorchTensorType(SHAPE, DTYPE, transport="nccl")) dag = receiver.recv.bind(dag) compiled_dag = dag.experimental_compile() ``` When `transport="nccl"` is specified, upon compile(), Ray will initialize a NCCL group with the actors involved. The reading actor(s) will `recv` on the NCCL communicator instead of reading from the default shared-memory channel. This PR also refactors channel types so that we now create `ChannelInterfaces` based on the type hints that appear in the DAG, either a `TorchTensorType` or the default `SharedMemoryType`. Current limitations: - p2p only, no collectives - Synchronizes CUDA stream after receiving data. This is because kernels following the NCCL op have no guarantee that the op succeeded, so it is not safe to read the received buffer unless we know that the op succeeded. - Shape and dtype of the tensor must be declared at compile time. --------- Signed-off-by: Stephanie Wang <swang@cs.berkeley.edu> Co-authored-by: SangBin Cho <rkooo567@gmail.com> Signed-off-by: Ryan O'Leary <ryanaoleary@google.com>
ryanaoleary
pushed a commit
to ryanaoleary/ray
that referenced
this pull request
Jun 6, 2024
Fixes test that was broken due a hidden merge conflict between ray-project#45092 and ray-project#45119. Related issue number Closes ray-project#45264. Signed-off-by: Stephanie Wang <swang@cs.berkeley.edu> Signed-off-by: Ryan O'Leary <ryanaoleary@google.com>
ryanaoleary
pushed a commit
to ryanaoleary/ray
that referenced
this pull request
Jun 6, 2024
Fix another broken test from ray-project#45092. Signed-off-by: Stephanie Wang <swang@cs.berkeley.edu> Signed-off-by: Ryan O'Leary <ryanaoleary@google.com>
ryanaoleary
pushed a commit
to ryanaoleary/ray
that referenced
this pull request
Jun 7, 2024
…h.Tensors (ray-project#45092) ## Why are these changes needed? This adds a NCCL-based transport option for torch.tensors. Here is an example of the API: ```python with InputNode() as inp: dag = sender.send.bind(inp) dag = dag.with_type_hint(TorchTensorType(SHAPE, DTYPE, transport="nccl")) dag = receiver.recv.bind(dag) compiled_dag = dag.experimental_compile() ``` When `transport="nccl"` is specified, upon compile(), Ray will initialize a NCCL group with the actors involved. The reading actor(s) will `recv` on the NCCL communicator instead of reading from the default shared-memory channel. This PR also refactors channel types so that we now create `ChannelInterfaces` based on the type hints that appear in the DAG, either a `TorchTensorType` or the default `SharedMemoryType`. Current limitations: - p2p only, no collectives - Synchronizes CUDA stream after receiving data. This is because kernels following the NCCL op have no guarantee that the op succeeded, so it is not safe to read the received buffer unless we know that the op succeeded. - Shape and dtype of the tensor must be declared at compile time. --------- Signed-off-by: Stephanie Wang <swang@cs.berkeley.edu> Co-authored-by: SangBin Cho <rkooo567@gmail.com>
ryanaoleary
pushed a commit
to ryanaoleary/ray
that referenced
this pull request
Jun 7, 2024
Fixes test that was broken due a hidden merge conflict between ray-project#45092 and ray-project#45119. Related issue number Closes ray-project#45264. Signed-off-by: Stephanie Wang <swang@cs.berkeley.edu>
ryanaoleary
pushed a commit
to ryanaoleary/ray
that referenced
this pull request
Jun 7, 2024
Fix another broken test from ray-project#45092. Signed-off-by: Stephanie Wang <swang@cs.berkeley.edu>
GabeChurch
pushed a commit
to GabeChurch/ray
that referenced
this pull request
Jun 11, 2024
…h.Tensors (ray-project#45092) ## Why are these changes needed? This adds a NCCL-based transport option for torch.tensors. Here is an example of the API: ```python with InputNode() as inp: dag = sender.send.bind(inp) dag = dag.with_type_hint(TorchTensorType(SHAPE, DTYPE, transport="nccl")) dag = receiver.recv.bind(dag) compiled_dag = dag.experimental_compile() ``` When `transport="nccl"` is specified, upon compile(), Ray will initialize a NCCL group with the actors involved. The reading actor(s) will `recv` on the NCCL communicator instead of reading from the default shared-memory channel. This PR also refactors channel types so that we now create `ChannelInterfaces` based on the type hints that appear in the DAG, either a `TorchTensorType` or the default `SharedMemoryType`. Current limitations: - p2p only, no collectives - Synchronizes CUDA stream after receiving data. This is because kernels following the NCCL op have no guarantee that the op succeeded, so it is not safe to read the received buffer unless we know that the op succeeded. - Shape and dtype of the tensor must be declared at compile time. --------- Signed-off-by: Stephanie Wang <swang@cs.berkeley.edu> Co-authored-by: SangBin Cho <rkooo567@gmail.com> Signed-off-by: gchurch <gabe1church@gmail.com>
GabeChurch
pushed a commit
to GabeChurch/ray
that referenced
this pull request
Jun 11, 2024
Fixes test that was broken due a hidden merge conflict between ray-project#45092 and ray-project#45119. Related issue number Closes ray-project#45264. Signed-off-by: Stephanie Wang <swang@cs.berkeley.edu> Signed-off-by: gchurch <gabe1church@gmail.com>
GabeChurch
pushed a commit
to GabeChurch/ray
that referenced
this pull request
Jun 11, 2024
Fix another broken test from ray-project#45092. Signed-off-by: Stephanie Wang <swang@cs.berkeley.edu> Signed-off-by: gchurch <gabe1church@gmail.com>
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Why are these changes needed?
This adds a NCCL-based transport option for torch.tensors. Here is an example of the API:
When
transport="nccl"
is specified, upon compile(), Ray will initialize a NCCL group with the actors involved. The reading actor(s) willrecv
on the NCCL communicator instead of reading from the default shared-memory channel.This PR also refactors channel types so that we now create
ChannelInterfaces
based on the type hints that appear in the DAG, either aTorchTensorType
or the defaultSharedMemoryType
.Current limitations:
TODO: Add unit and release tests to PR,
Checks
git commit -s
) in this PR.scripts/format.sh
to lint the changes in this PR.method in Tune, I've added it in
doc/source/tune/api/
under thecorresponding
.rst
file.