[core][compiled graphs] Meta-issue: Support collective communication ops

[stephanie-wang]

Description

This is a meta-issue to track progress for tasks related to collective communication. See RFC for more details.

Roadmap:

• Initial support for allreduce: [aDAG] Support all reduce collective in aDAG #47621
• [WIP] Unify p2p and collective op codepaths: (WIP) [core][compiled graphs] Unify code paths for NCCL P2P and collectives scheduling #48649
• [WIP] Support for other all-to-all patterns, e.g., allgather, etc.
• Support all-to-one patterns, e.g., reduce
• [core][compiled graphs] Execute multiple p2p NCCL transfers with collective ops #47938
• Error detection for users passing different shapes on different workers. Possibilities:
• Support direct_return=False, i.e. cases where user returns a mix of CPU and GPU data
• Support batching multiple tensors into one communication op
• Support for non-compiled graphs
• [core][compiled graphs] Add CPU-based NCCL communicator for development #47936

Use case

No response

[jeffreyjeffreywang]

Hey @stephanie-wang, just wanted to make sure this meta-task is still on the roadmap before I start the implementation. I'm particularly interested in #47938, but I think it's better to start with supporting all-to-all, all-to-one, and one-to-all collectives. I took a look at the RFC, and here is my understanding:

• All-to-one: Pass the reader worker handle to the collective call

workers = [Worker.options(num_gpus=1).remote() for _ in range(3)]
nccl_group_handle = ray.collective.NcclGroup(workers)
with InputNode() as inp:
  results = [worker.fwd.bind(inp) for worker in workers]
  # Pass the worker handle to the collective call.
  dag = ray.collective.gather.bind(
    results, workers[0],
    transport=nccl_group_handle)
  dag = workers[0].sync.bind(dag)

# Errors if `gather` reader is not part of the group.
dag = dag.experimental_compile()

• One-to-all: Pass the sender worker handle to the collective call

workers = [Worker.options(num_gpus=1).remote() for _ in range(3)]
nccl_group_handle = ray.collective.NcclGroup(workers)
# One-to-all pattern.
with InputNode() as inp:
  result = workers[0].fwd.bind(inp)
  results = ray.collective.broadcast.bind(
    result, workers,
    transport=nccl_group_handle)
  dag = MultiOutputNode(results)

# Errors if `broadcast` sender is not part of the group.
dag = dag.experimental_compile()

Is this the design that we agreed upon? If this looks good to you, I'll create issues to track the support for all-to-all, all-to-one, and one-to-all patterns.

[stephanie-wang]

Hey @stephanie-wang, just wanted to make sure this meta-task is still on the roadmap before I start the implementation. I'm particularly interested in #47938, but I think it's better to start with supporting all-to-all, all-to-one, and one-to-all collectives. I took a look at the RFC, and here is my understanding:

* All-to-one: Pass the reader worker handle to the collective call

workers = [Worker.options(num_gpus=1).remote() for _ in range(3)]
nccl_group_handle = ray.collective.NcclGroup(workers)
with InputNode() as inp:
  results = [worker.fwd.bind(inp) for worker in workers]
  # Pass the worker handle to the collective call.
  dag = ray.collective.gather.bind(
    results, workers[0],
    transport=nccl_group_handle)
  dag = workers[0].sync.bind(dag)

# Errors if `gather` reader is not part of the group.
dag = dag.experimental_compile()

* One-to-all: Pass the sender worker handle to the collective call

workers = [Worker.options(num_gpus=1).remote() for _ in range(3)]
nccl_group_handle = ray.collective.NcclGroup(workers)
# One-to-all pattern.
with InputNode() as inp:
  result = workers[0].fwd.bind(inp)
  results = ray.collective.broadcast.bind(
    result, workers,
    transport=nccl_group_handle)
  dag = MultiOutputNode(results)

# Errors if `broadcast` sender is not part of the group.
dag = dag.experimental_compile()

Is this the design that we agreed upon? If this looks good to you, I'll create issues to track the support for all-to-all, all-to-one, and one-to-all patterns.

Yes, everything here is still on the roadmap and open to contributions!

[jeffreyjeffreywang]

Thanks Stephanie, here are the links to the issues for all-to-one and one-to-all.

• All-to-one: [core][compiled graph] Support all-to-one collective ops (e.g. gather) #49324
• One-to-all: [core][compiled graph] Support one-to-all collective ops (e.g. broadcast) #49325