chore: import upstream snapshot with attribution

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

@@ -0,0 +1,469 @@
+# coding: utf-8
+import os
+import sys
+from typing import Optional
+
+import pytest
+import torch
+
+import ray
+import ray.cluster_utils
+from ray.dag import InputNode, MultiOutputNode
+from ray.dag.compiled_dag_node import CompiledDAG
+from ray.dag.dag_node_operation import _DAGNodeOperationType
+from ray.tests.conftest import *  # noqa
+
+if sys.platform != "linux" and sys.platform != "darwin":
+    pytest.skip("Skipping, requires Linux or Mac.", allow_module_level=True)
+
+USE_GPU = os.environ.get("RAY_PYTEST_USE_GPU") == "1"
+
+if not USE_GPU:
+    pytest.skip("Skipping, these tests require GPUs.", allow_module_level=True)
+
+
+@ray.remote(num_cpus=0, num_gpus=1)
+class Worker:
+    def __init__(self, rank: Optional[int] = None):
+        self.rank = rank
+        self.trace = []
+
+    def fwd(self, value):
+        self.trace.append(("FWD", self.rank))
+        return value
+
+    def bwd(self, value):
+        self.trace.append(("BWD", self.rank))
+        return value
+
+    def pop_trace(self):
+        trace = self.trace
+        self.trace = []
+        return trace
+
+    def read_input(self, input):
+        return input
+
+    def send(self, shape, dtype, value: int, send_tensor=True):
+        if not send_tensor:
+            return 1
+        return torch.ones(shape, dtype=dtype, device=self.device) * value
+
+    def recv(self, tensor):
+        # Check that tensor got loaded to the correct device.
+        assert tensor.device == self.device
+        return (tensor[0].item(), tensor.shape, tensor.dtype)
+
+    def no_op(self, value):
+        return value
+
+    def no_op_two(self, value1, value2):
+        return value1, value2
+
+
+def generate_1f1b_dag(
+    num_workers: int, num_microbatches: int, num_lead_microbatches: int
+) -> CompiledDAG:
+    workers = [Worker.remote(rank) for rank in range(num_workers)]
+
+    with ray.dag.InputNode() as inp:
+        fwd_queues = [[] for _ in range(num_workers)]
+        bwd_queues = [[] for _ in range(num_workers)]
+        # Once a worker's counter reaches 0, it cannot execute another fwd until it
+        # executes a bwd first.
+        fwd_counter = [num_lead_microbatches - i for i in range(num_workers)]
+        # All of the done batches.
+        done = []
+
+        # FWD on worker 0.
+        input_data = workers[0].read_input.bind(inp)
+        for i in range(num_microbatches):
+            fwd_queues[0].append(input_data)
+
+        while len(done) < num_microbatches:
+            for i, worker in enumerate(workers):
+                if fwd_counter[i] > 0 and fwd_queues[i]:
+                    b = fwd_queues[i].pop(0)
+                    b = worker.fwd.bind(b)
+                    if i < num_workers - 1:
+                        fwd_queues[i + 1].append(b)
+                        # Use NCCL channel for communication between workers.
+                        b.with_tensor_transport(transport="nccl")
+                    else:
+                        bwd_queues[i].append(b)
+                    fwd_counter[i] -= 1
+                elif bwd_queues[i]:
+                    b = bwd_queues[i].pop(0)
+                    b = worker.bwd.bind(b)
+                    if i > 0:
+                        bwd_queues[i - 1].append(b)
+                        # Use NCCL channel for communication between workers.
+                        b.with_tensor_transport(transport="nccl")
+                    else:
+                        done.append(b)
+                    fwd_counter[i] += 1
+        dag = ray.dag.MultiOutputNode(done)
+    compiled_dag = dag.experimental_compile()
+    return compiled_dag
+
+
+@pytest.mark.parametrize("ray_start_regular", [{"num_gpus": 2}], indirect=True)
+@pytest.mark.parametrize("single_fetch", [True, False])
+def test_simulate_pp_2workers_2batches_1f1b(
+    ray_start_regular, single_fetch, monkeypatch
+):
+    """
+    This test simulates a simple 1F1B pipeline parallelism for training with
+    2 workers and 2 batches.
+
+    w1: fwd_b1  fwd_b2          bwd_b1          bwd_b2
+    w2:         fwd_b1  bwd_b1  fwd_b2  bwd_b2
+
+    The communication between workers is done using NCCL. The communication
+    within the worker actor is done using IntraProcessChannel.
+    """
+    if not USE_GPU:
+        pytest.skip("NCCL tests require GPUs")
+
+    w1 = Worker.remote()
+    w2 = Worker.remote()
+
+    with InputNode() as inp:
+        w1_input = w1.read_input.bind(inp)
+        batch_1 = w1.fwd.bind(w1_input)
+        batch_1.with_tensor_transport(transport="nccl")
+        batch_2 = w1.fwd.bind(w1_input)
+        batch_2.with_tensor_transport(transport="nccl")
+        batch_1 = w2.fwd.bind(batch_1)
+        batch_1 = w2.bwd.bind(batch_1)
+        batch_1.with_tensor_transport(transport="nccl")
+        batch_2 = w2.fwd.bind(batch_2)
+        batch_1 = w1.bwd.bind(batch_1)
+        batch_2 = w2.bwd.bind(batch_2)
+        batch_2.with_tensor_transport(transport="nccl")
+        batch_2 = w1.bwd.bind(batch_2)
+        dag = MultiOutputNode([batch_1, batch_2])
+    compiled_dag = dag.experimental_compile()
+
+    w1_expected_schedule = [
+        (0, _DAGNodeOperationType.READ),
+        (0, _DAGNodeOperationType.COMPUTE),
+        (0, _DAGNodeOperationType.WRITE),
+        (1, _DAGNodeOperationType.READ),
+        (1, _DAGNodeOperationType.COMPUTE),
+        (1, _DAGNodeOperationType.WRITE),
+        # `w1 (3, READ)` (P2P recv) is scheduled together with
+        # `w2 (1, WRITE)` (P2P send).
+        (3, _DAGNodeOperationType.READ),
+        (2, _DAGNodeOperationType.READ),
+        (2, _DAGNodeOperationType.COMPUTE),
+        (2, _DAGNodeOperationType.WRITE),
+        # `w1 (4, READ)` (P2P recv) is scheduled together with
+        # `w2 (3, WRITE)` (P2P send).
+        (4, _DAGNodeOperationType.READ),
+        (3, _DAGNodeOperationType.COMPUTE),
+        (3, _DAGNodeOperationType.WRITE),
+        (4, _DAGNodeOperationType.COMPUTE),
+        (4, _DAGNodeOperationType.WRITE),
+    ]
+    w2_expected_schedule = [
+        (0, _DAGNodeOperationType.READ),
+        (0, _DAGNodeOperationType.COMPUTE),
+        (0, _DAGNodeOperationType.WRITE),
+        (1, _DAGNodeOperationType.READ),
+        (1, _DAGNodeOperationType.COMPUTE),
+        (1, _DAGNodeOperationType.WRITE),
+        (2, _DAGNodeOperationType.READ),
+        (2, _DAGNodeOperationType.COMPUTE),
+        (2, _DAGNodeOperationType.WRITE),
+        (3, _DAGNodeOperationType.READ),
+        (3, _DAGNodeOperationType.COMPUTE),
+        (3, _DAGNodeOperationType.WRITE),
+    ]
+    w1_schedule = compiled_dag.actor_to_execution_schedule[w1]
+    w2_schedule = compiled_dag.actor_to_execution_schedule[w2]
+
+    for schedule, expected_schedule in zip(
+        [w1_schedule, w2_schedule], [w1_expected_schedule, w2_expected_schedule]
+    ):
+        assert len(schedule) == len(expected_schedule)
+        for i, operation in enumerate(schedule):
+            assert operation.exec_task_idx == expected_schedule[i][0]
+            assert operation.type == expected_schedule[i][1]
+
+    tensor_cpu = torch.zeros(10, 10)
+    tensor_cuda = tensor_cpu.to("cuda:0")
+    refs = compiled_dag.execute(tensor_cuda)
+
+    if single_fetch:
+        assert len(refs) == 2
+        for ref in refs:
+            tensor = ray.get(ref)
+            assert torch.equal(tensor.cpu(), tensor_cpu)
+    else:
+        tensors = ray.get(refs)
+        assert len(tensors) == 2
+        for tensor in tensors:
+            assert torch.equal(tensor.cpu(), tensor_cpu)
+
+
+@pytest.mark.parametrize("ray_start_regular", [{"num_gpus": 4}], indirect=True)
+def test_simulate_pp_4workers_8batches_1f1b(ray_start_regular, monkeypatch):
+    """
+    This test simulates a 1F1B pipeline parallelism for training with
+    4 workers and 8 batches.
+    """
+    if not USE_GPU:
+        pytest.skip("NCCL tests require GPUs")
+
+    num_workers, num_microbatches, num_lead_microbatches = 4, 8, 4
+    compiled_dag = generate_1f1b_dag(
+        num_workers, num_microbatches, num_lead_microbatches
+    )
+
+    tensor_cpu = torch.zeros(10, 10)
+    tensor_cuda = tensor_cpu.to("cuda:0")
+    tensors = ray.get(compiled_dag.execute(tensor_cuda))
+
+    assert len(tensors) == num_microbatches
+    for t in tensors:
+        assert torch.equal(t.cpu(), tensor_cpu)
+
+
+@pytest.mark.parametrize("ray_start_regular", [{"num_gpus": 3}], indirect=True)
+def test_three_actors_with_nccl_1(ray_start_regular):
+    """
+    Driver -> a.no_op -> b.no_op -> a.no_op_two -> Driver
+                      |          |
+                      -> c.no_op -
+    """
+    if not USE_GPU:
+        pytest.skip("NCCL tests require GPUs")
+
+    a = Worker.remote()
+    b = Worker.remote()
+    c = Worker.remote()
+
+    with InputNode() as inp:
+        dag = a.no_op.bind(inp)
+        dag.with_tensor_transport(transport="nccl")
+        branch1 = b.no_op.bind(dag)
+        branch1.with_tensor_transport(transport="nccl")
+        branch2 = c.no_op.bind(dag)
+        branch2.with_tensor_transport(transport="nccl")
+        dag = a.no_op_two.bind(branch1, branch2)
+
+    compiled_dag = dag.experimental_compile()
+
+    a_expected_schedule = [
+        (0, _DAGNodeOperationType.READ),
+        (0, _DAGNodeOperationType.COMPUTE),
+        (0, _DAGNodeOperationType.WRITE),
+        (1, _DAGNodeOperationType.READ),
+        (1, _DAGNodeOperationType.COMPUTE),
+        (1, _DAGNodeOperationType.WRITE),
+    ]
+    b_expected_schedule = [
+        (0, _DAGNodeOperationType.READ),
+        (0, _DAGNodeOperationType.COMPUTE),
+        (0, _DAGNodeOperationType.WRITE),
+    ]
+    c_expected_schedule = [
+        (0, _DAGNodeOperationType.READ),
+        (0, _DAGNodeOperationType.COMPUTE),
+        (0, _DAGNodeOperationType.WRITE),
+    ]
+    a_schedule = compiled_dag.actor_to_execution_schedule[a]
+    b_schedule = compiled_dag.actor_to_execution_schedule[b]
+    c_schedule = compiled_dag.actor_to_execution_schedule[c]
+
+    for schedule, expected_schedule in zip(
+        [a_schedule, b_schedule, c_schedule],
+        [a_expected_schedule, b_expected_schedule, c_expected_schedule],
+    ):
+        assert len(schedule) == len(expected_schedule)
+        for i, operation in enumerate(schedule):
+            assert operation.exec_task_idx == expected_schedule[i][0]
+            assert operation.type == expected_schedule[i][1]
+
+    tensor_cpu = torch.zeros(10, 10)
+    tensor_cuda = tensor_cpu.to("cuda:0")
+    ref = compiled_dag.execute(tensor_cuda)
+    tensors = ray.get(ref)
+
+    assert len(tensors) == 2
+    for t in tensors:
+        assert torch.equal(t.cpu(), tensor_cpu)
+
+
+@pytest.mark.parametrize("ray_start_regular", [{"num_gpus": 3}], indirect=True)
+@pytest.mark.parametrize("single_fetch", [True, False])
+def test_three_actors_with_nccl_2(ray_start_regular, single_fetch, monkeypatch):
+    """
+    Driver --> a.no_op -> b.no_op --> Driver
+            |                     |
+            -> b.no_op -> c.no_op -
+            |                     |
+            -> c.no_op -> a.no_op -
+    """
+    if not USE_GPU:
+        pytest.skip("NCCL tests require GPUs")
+
+    a = Worker.remote()
+    b = Worker.remote()
+    c = Worker.remote()
+
+    with InputNode() as inp:
+        branch1 = a.no_op.bind(inp)
+        branch1.with_tensor_transport(transport="nccl")
+        branch2 = b.no_op.bind(inp)
+        branch2.with_tensor_transport(transport="nccl")
+        branch3 = c.no_op.bind(inp)
+        branch3.with_tensor_transport(transport="nccl")
+        dag = MultiOutputNode(
+            [
+                a.no_op.bind(branch3),
+                b.no_op.bind(branch1),
+                c.no_op.bind(branch2),
+            ]
+        )
+
+    compiled_dag = dag.experimental_compile()
+
+    a_expected_schedule = [
+        (0, _DAGNodeOperationType.READ),
+        (0, _DAGNodeOperationType.COMPUTE),
+        (0, _DAGNodeOperationType.WRITE),
+        (1, _DAGNodeOperationType.READ),
+        (1, _DAGNodeOperationType.COMPUTE),
+        (1, _DAGNodeOperationType.WRITE),
+    ]
+    b_expected_schedule = [
+        # `b (1, READ)` (P2P recv) is scheduled together with
+        # `a (0, WRITE)` (P2P send).
+        (1, _DAGNodeOperationType.READ),
+        (0, _DAGNodeOperationType.READ),
+        (0, _DAGNodeOperationType.COMPUTE),
+        (0, _DAGNodeOperationType.WRITE),
+        (1, _DAGNodeOperationType.COMPUTE),
+        (1, _DAGNodeOperationType.WRITE),
+    ]
+    c_expected_schedule = [
+        # `c (1, READ)` (P2P recv) is scheduled together with
+        # `a (0, WRITE)` (P2P send).
+        (1, _DAGNodeOperationType.READ),
+        (0, _DAGNodeOperationType.READ),
+        (0, _DAGNodeOperationType.COMPUTE),
+        (0, _DAGNodeOperationType.WRITE),
+        (1, _DAGNodeOperationType.COMPUTE),
+        (1, _DAGNodeOperationType.WRITE),
+    ]
+
+    a_schedule = compiled_dag.actor_to_execution_schedule[a]
+    b_schedule = compiled_dag.actor_to_execution_schedule[b]
+    c_schedule = compiled_dag.actor_to_execution_schedule[c]
+
+    for schedule, expected_schedule in zip(
+        [a_schedule, b_schedule, c_schedule],
+        [a_expected_schedule, b_expected_schedule, c_expected_schedule],
+    ):
+        assert len(schedule) == len(expected_schedule)
+        for i, operation in enumerate(schedule):
+            assert operation.exec_task_idx == expected_schedule[i][0]
+            assert operation.type == expected_schedule[i][1]
+
+    tensor_cpu = torch.zeros(10, 10)
+    tensor_cuda = tensor_cpu.to("cuda:0")
+    refs = compiled_dag.execute(tensor_cuda)
+
+    if single_fetch:
+        assert len(refs) == 3
+        for ref in refs:
+            tensor = ray.get(ref)
+            assert torch.equal(tensor.cpu(), tensor_cpu)
+    else:
+        tensors = ray.get(refs)
+        assert len(tensors) == 3
+        for tensor in tensors:
+            assert torch.equal(tensor.cpu(), tensor_cpu)
+
+
+@pytest.mark.parametrize("ray_start_regular", [{"num_gpus": 3}], indirect=True)
+@pytest.mark.parametrize("overlap_gpu_communication", [True, False])
+def test_overlap_gpu_communication(ray_start_regular, overlap_gpu_communication):
+    """
+    Driver --> sender1.send -> receiver.recv --> Driver
+            |                                |
+            -> sender2.send -> receiver.recv -
+    """
+    if not USE_GPU:
+        pytest.skip("NCCL tests require GPUs")
+
+    sender1 = Worker.remote()
+    sender2 = Worker.remote()
+    receiver = Worker.remote()
+
+    shape = (10000,)
+    dtype = torch.float16
+
+    with InputNode() as inp:
+        branch1 = sender1.send.bind(shape, dtype, inp)
+
+        branch1 = branch1.with_tensor_transport(
+            transport="nccl", _static_shape=True, _direct_return=True
+        )
+        branch1 = receiver.recv.bind(branch1)
+
+        branch2 = sender2.send.bind(shape, dtype, inp)
+        branch2 = branch2.with_tensor_transport(
+            transport="nccl", _static_shape=True, _direct_return=True
+        )
+        branch2 = receiver.recv.bind(branch2)
+        dag = MultiOutputNode([branch1, branch2])
+
+    # Test normal execution.
+    compiled_dag = dag.experimental_compile(
+        _overlap_gpu_communication=overlap_gpu_communication
+    )
+
+    # Check receiver schedule
+    expected_no_overlap_schedule = [
+        (0, _DAGNodeOperationType.READ),
+        # `receiver (1, READ)` (P2P recv) is scheduled together with
+        # `sender2 (0, WRITE)` (P2P send).
+        (1, _DAGNodeOperationType.READ),
+        (0, _DAGNodeOperationType.COMPUTE),
+        (0, _DAGNodeOperationType.WRITE),
+        (1, _DAGNodeOperationType.COMPUTE),
+        (1, _DAGNodeOperationType.WRITE),
+    ]
+    expected_overlap_schedule = [
+        (0, _DAGNodeOperationType.READ),
+        # `receiver (1, READ)` (P2P recv) is scheduled together with
+        # `sender2 (0, WRITE)` (P2P send).
+        (1, _DAGNodeOperationType.READ),
+        (0, _DAGNodeOperationType.COMPUTE),
+        (0, _DAGNodeOperationType.WRITE),
+        (1, _DAGNodeOperationType.COMPUTE),
+        (1, _DAGNodeOperationType.WRITE),
+    ]
+    if overlap_gpu_communication:
+        expected_receiver_schedule = expected_overlap_schedule
+    else:
+        expected_receiver_schedule = expected_no_overlap_schedule
+
+    receiver_schedule = compiled_dag.actor_to_execution_schedule[receiver]
+
+    assert len(receiver_schedule) == len(expected_receiver_schedule)
+    for i, operation in enumerate(receiver_schedule):
+        assert operation.exec_task_idx == expected_receiver_schedule[i][0]
+        assert operation.type == expected_receiver_schedule[i][1]
+
+    compiled_dag.teardown()
+
+
+if __name__ == "__main__":
+    if os.environ.get("PARALLEL_CI"):
+        sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
+    else:
+        sys.exit(pytest.main(["-sv", __file__]))