[core][compiled graphs] Support reduce scatter collective in compiled graph #49404

anyadontfly · 2024-12-22T23:30:02Z

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:

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)

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

Signed-off-by: Puyuan Yao <[email protected]> rebase to updated main branch

Signed-off-by: Puyuan Yao <[email protected]>

anyadontfly · 2024-12-23T09:31:32Z

python/ray/dag/collective_node.py

+            communicator.allreduce(send_buf, recv_buf, self._op)
+        elif self.comm_op == self.REDUCESCATTER:
+            world_size = len(self._actor_handles)
+            assert (


I'm planning on adding a test on throwing an error when input tensors' first dimension is not divisible by the number of reduce scatter participants. This is a RayTaskError, I'm not sure what is the expected behavior of compiled graph when this error is generated.

Sounds good. Raise ValueError for now.

dengwxn

Overall looks good! Left some comments to polish.

dengwxn · 2024-12-23T18:23:16Z

python/ray/experimental/collective/reducescatter.py

+    PARENT_CLASS_NODE_KEY,
+)
+from ray.experimental.channel.torch_tensor_type import Communicator, TorchTensorType
+from ray.experimental.util.types import ReduceOp


Create two new enums AllReduceOp and ReduceScatterOp inherited from ReduceOp. They do not need to introduce new enum values, just use pass. Such that we don't need a str to tell whether this is a allreduce or reducescatter, simply just check if their types are AllReduceOp and ReduceScatterOp.

dengwxn · 2024-12-23T18:24:38Z

python/ray/experimental/collective/reducescatter.py

+    def bind(
+        self,
+        input_nodes: List["ray.dag.DAGNode"],
+        op: ReduceOp = ReduceOp.SUM,


After introducing the two new enums, this would be ReduceScatterOp.SUM. Also update for allreduce to use AllReduceOp.SUM. Make sure update all the places including allreduce.py and tests.

dengwxn · 2024-12-23T18:26:02Z

python/ray/dag/tests/experimental/test_torch_tensor_dag.py

+
+
+@pytest.mark.parametrize("ray_start_regular", [{"num_cpus": 4}], indirect=True)
+def test_torch_tensor_nccl_reduce_scatter_different_shapes_among_participants(


Is this test same as all_reduce_wrong_shape? If so, we should use the same names for both. wrong_shape is fine to keep.

dengwxn · 2024-12-23T18:26:42Z

python/ray/dag/tests/experimental/test_torch_tensor_dag.py

+            self,
+            send_buf: "torch.Tensor",
+            recv_buf: "torch.Tensor",
+            op: ReduceOp = ReduceOp.SUM,


After introducing two new enums, we need to update all of these.

dengwxn · 2024-12-23T18:27:02Z

python/ray/dag/collective_node.py

@@ -31,9 +31,13 @@ class _CollectiveOperation:
    3. Actor handles match the custom NCCL group if specified.
    """

+    ALLREDUCE = "ar"


After introducing the two new enums, we don't need these.

dengwxn · 2024-12-23T18:27:55Z

python/ray/dag/collective_node.py

+            world_size = len(self._actor_handles)
+            assert (
+                send_buf.shape[0] % world_size == 0
+            ), "Input tensor's first dimension should be divisible by "


Raise a ValueError here. For all the user-input errors, we would raise a ValueError instead of asserting. Assert is used for checking internal states.

dengwxn · 2024-12-23T18:28:13Z

python/ray/dag/collective_node.py

+            communicator.allreduce(send_buf, recv_buf, self._op)
+        elif self.comm_op == self.REDUCESCATTER:
+            world_size = len(self._actor_handles)
+            assert (


Sounds good. Raise ValueError for now.

anyadontfly · 2024-12-24T20:18:07Z

python/ray/experimental/util/types.py

@@ -9,6 +9,23 @@ class _CollectiveOp(Enum):

 @PublicAPI
 class ReduceOp(_CollectiveOp):


I changed the types to this because python has strict requirements on subclassing enum.

anyadontfly added 6 commits December 20, 2024 00:56

first try

a9c2692

Signed-off-by: Puyuan Yao <[email protected]> rebase to updated main branch

adding reducescatter wrapper

7ee67b4

Signed-off-by: Puyuan Yao <[email protected]>

sync changes

8b53da0

passed basic test on default nccl group

50658c4

added reduce scatter for cpu communicator, added tests

be4cb1d

lint

aee414c

anyadontfly commented Dec 23, 2024

View reviewed changes

jcotant1 added core Issues that should be addressed in Ray Core compiled-graphs labels Dec 23, 2024

dengwxn reviewed Dec 23, 2024

View reviewed changes

anyadontfly changed the title ~~[compiled graphs] Support reduce scatter collective in compiled graph~~ [core][compiled graphs] Support reduce scatter collective in compiled graph Dec 23, 2024

anyadontfly added 2 commits December 24, 2024 01:38

fix first try

3fc93b2

small fix

90ee5fd

anyadontfly commented Dec 24, 2024

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[core][compiled graphs] Support reduce scatter collective in compiled graph #49404

[core][compiled graphs] Support reduce scatter collective in compiled graph #49404

anyadontfly commented Dec 22, 2024 •

edited

Loading

anyadontfly Dec 23, 2024

dengwxn Dec 23, 2024

dengwxn left a comment

dengwxn Dec 23, 2024

dengwxn Dec 23, 2024

dengwxn Dec 23, 2024

dengwxn Dec 23, 2024

dengwxn Dec 23, 2024

dengwxn Dec 23, 2024

dengwxn Dec 23, 2024

anyadontfly Dec 24, 2024



		@pytest.mark.parametrize("ray_start_regular", [{"num_cpus": 4}], indirect=True)
		def test_torch_tensor_nccl_reduce_scatter_different_shapes_among_participants(

		@@ -9,6 +9,23 @@ class _CollectiveOp(Enum):

		@PublicAPI
		class ReduceOp(_CollectiveOp):

[core][compiled graphs] Support reduce scatter collective in compiled graph #49404

Are you sure you want to change the base?

[core][compiled graphs] Support reduce scatter collective in compiled graph #49404

Conversation

anyadontfly commented Dec 22, 2024 • edited Loading

Why are these changes needed?

Related issue number

Checks

Choose a reason for hiding this comment

Choose a reason for hiding this comment

dengwxn left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

anyadontfly commented Dec 22, 2024 •

edited

Loading