chore: import upstream snapshot with attribution

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

@@ -0,0 +1,75 @@
+# coding: utf-8
+import os
+import sys
+
+import pytest
+import torch
+
+import ray
+import ray.cluster_utils
+from ray.dag import InputNode, MultiOutputNode
+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"
+
+
+@pytest.mark.parametrize("ray_start_regular", [{"num_gpus": 2}], indirect=True)
+def test_multi_args_simulate_pp(ray_start_regular):
+    if not USE_GPU:
+        pytest.skip("NCCL tests require GPUs")
+
+    @ray.remote(num_cpus=0, num_gpus=1)
+    class Worker:
+        def __init__(self):
+            pass
+
+        def forward(self, data):
+            return data
+
+        def backward(self, data):
+            return data
+
+    NUM_MICROBATCHES = 2
+    w0 = Worker.remote()
+    w1 = Worker.remote()
+    with InputNode() as dag_input:
+        dag_outs = []
+        for microbatch_idx in range(NUM_MICROBATCHES):
+            microbatch = dag_input[microbatch_idx]
+            stage_fwd_out = w0.forward.bind(microbatch)
+            stage_fwd_out.with_tensor_transport(transport="nccl")
+            stage_fwd_out = w1.forward.bind(stage_fwd_out)
+            dag_outs.append(stage_fwd_out)
+
+        grad_out = dag_input[NUM_MICROBATCHES]
+        for _ in range(NUM_MICROBATCHES):
+            stage_bwd_out = w1.backward.bind(grad_out)
+            stage_bwd_out.with_tensor_transport(transport="nccl")
+            stage_bwd_out = w0.backward.bind(stage_bwd_out)
+            dag_outs.append(stage_bwd_out)
+
+        dag = MultiOutputNode(dag_outs)
+    compiled_dag = dag.experimental_compile()
+
+    tensor_cpu_list = [torch.zeros(1, i + 1) for i in range(3)]
+    tensor_cuda_list = [t.to("cuda:0") for t in tensor_cpu_list]
+    ref = compiled_dag.execute(
+        tensor_cuda_list[0], tensor_cuda_list[1], tensor_cuda_list[2]
+    )
+    tensors = ray.get(ref)
+
+    assert len(tensors) == 4
+    assert torch.equal(tensors[0], tensor_cpu_list[0])
+    assert torch.equal(tensors[1], tensor_cpu_list[1])
+    assert torch.equal(tensors[2], tensor_cpu_list[2])
+    assert torch.equal(tensors[3], tensor_cpu_list[2])
+
+
+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__]))