wehub-resource-sync
470 lines
16 KiB
Python
470 lines
16 KiB
Python
# 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__]))
|