wehub-resource-sync
216 lines
7.4 KiB
Python
216 lines
7.4 KiB
Python
import time
|
|
from typing import List
|
|
|
|
import pytest
|
|
import torch
|
|
from torch.nn.parallel import DistributedDataParallel
|
|
from torch.utils.data import DataLoader, DistributedSampler
|
|
|
|
import ray
|
|
from ray.train import RunConfig, ScalingConfig
|
|
from ray.train.examples.pytorch.torch_linear_example import LinearDataset
|
|
from ray.train.torch import TorchTrainer
|
|
from ray.train.v2._internal.execution.callback import WorkerGroupCallback
|
|
from ray.train.v2._internal.execution.worker_group import Worker
|
|
from ray.train.v2.api.exceptions import WorkerGroupError
|
|
|
|
|
|
def test_torch_trainer_cuda_initialization(ray_start_4_cpus_2_gpus):
|
|
"""Test that Torch CUDA initialization works with TorchTrainer.
|
|
|
|
This test verifies that PyTorch can properly initialize CUDA on multiple
|
|
workers before the training context is set up, ensuring that GPU resources
|
|
are available and accessible across all training workers.
|
|
|
|
See https://github.com/ray-project/ray/pull/56509 for more details.
|
|
"""
|
|
|
|
def train_func():
|
|
"""Empty training function for this initialization test.
|
|
|
|
Since we're only testing CUDA initialization, the actual training
|
|
logic is not needed for this test case.
|
|
"""
|
|
pass
|
|
|
|
def init_torch():
|
|
"""Trigger (lazy) initialization of CUDA."""
|
|
torch.cuda.is_available()
|
|
|
|
class InitTorchCallback(WorkerGroupCallback):
|
|
"""Callback to initialize PyTorch CUDA before training begins.
|
|
|
|
Implements before_init_train_context because this is where torch is typically imported,
|
|
ensuring that the CUDA environment is properly initialized.
|
|
"""
|
|
|
|
def before_init_train_context(self, workers: List[Worker]):
|
|
"""Execute CUDA initialization on all workers."""
|
|
futures = []
|
|
for worker in workers:
|
|
futures.append(worker.execute_async(init_torch))
|
|
ray.get(futures)
|
|
return {}
|
|
|
|
callback = InitTorchCallback()
|
|
|
|
trainer = TorchTrainer(
|
|
train_func,
|
|
scaling_config=ScalingConfig(num_workers=2, use_gpu=True),
|
|
run_config=RunConfig(callbacks=[callback]),
|
|
)
|
|
|
|
trainer.fit()
|
|
|
|
|
|
@pytest.mark.parametrize("num_gpus_per_worker", [0.5, 1, 2])
|
|
def test_torch_get_devices(ray_start_2x2_gpu_cluster, num_gpus_per_worker):
|
|
# cluster setups: 2 nodes, 2 gpus per node
|
|
# `CUDA_VISIBLE_DEVICES` is set to "0,1" on node 1 and node 2
|
|
if num_gpus_per_worker == 0.5:
|
|
# worker gpu topology:
|
|
# 4 workers on node 1, 4 workers on node 2
|
|
# `ray.get_gpu_ids()` returns [0], [0], [1], [1] on node 1
|
|
# and [0], [0], [1], [1] on node 2
|
|
expected_devices_per_rank = [[0], [0], [1], [1], [0], [0], [1], [1]]
|
|
elif num_gpus_per_worker == 1:
|
|
# worker gpu topology:
|
|
# 2 workers on node 1, 2 workers on node 2
|
|
# `ray.get_gpu_ids()` returns [0], [1] on node 1 and [0], [1] on node 2
|
|
expected_devices_per_rank = [[0], [1], [0], [1]]
|
|
elif num_gpus_per_worker == 2:
|
|
# worker gpu topology:
|
|
# 1 workers on node 1, 1 workers on node 2
|
|
# `ray.get_gpu_ids()` returns {0, 1} on node 1 and {0, 1} on node 2
|
|
# and `device_id` returns the one index from each set.
|
|
# So total count of devices should be 2.
|
|
expected_devices_per_rank = [[0, 1], [0, 1]]
|
|
else:
|
|
raise RuntimeError(
|
|
"New parameter for this test has been added without checking that the "
|
|
"correct devices have been returned."
|
|
)
|
|
|
|
def train_fn():
|
|
assert torch.cuda.current_device() == ray.train.torch.get_device().index
|
|
devices = sorted([device.index for device in ray.train.torch.get_devices()])
|
|
rank = ray.train.get_context().get_world_rank()
|
|
assert devices == expected_devices_per_rank[rank]
|
|
|
|
trainer = TorchTrainer(
|
|
train_fn,
|
|
scaling_config=ray.train.ScalingConfig(
|
|
num_workers=int(4 / num_gpus_per_worker),
|
|
use_gpu=True,
|
|
resources_per_worker={"GPU": num_gpus_per_worker},
|
|
),
|
|
)
|
|
trainer.fit()
|
|
|
|
|
|
def test_torch_prepare_model(ray_start_4_cpus_2_gpus):
|
|
"""Tests if ``prepare_model`` correctly wraps in DDP."""
|
|
|
|
def train_fn():
|
|
model = torch.nn.Linear(1, 1)
|
|
|
|
# Wrap in DDP.
|
|
model = ray.train.torch.prepare_model(model)
|
|
|
|
# Make sure model is wrapped in DDP.
|
|
assert isinstance(model, DistributedDataParallel)
|
|
|
|
# Make sure model is on cuda.
|
|
assert next(model.parameters()).is_cuda
|
|
|
|
trainer = TorchTrainer(
|
|
train_fn, scaling_config=ScalingConfig(num_workers=2, use_gpu=True)
|
|
)
|
|
trainer.fit()
|
|
|
|
|
|
class LinearDatasetDict(LinearDataset):
|
|
"""Modifies the LinearDataset to return a Dict instead of a Tuple."""
|
|
|
|
def __getitem__(self, index):
|
|
return {"x": self.x[index, None], "y": self.y[index, None]}
|
|
|
|
|
|
class NonTensorDataset(LinearDataset):
|
|
"""Modifies the LinearDataset to also return non-tensor objects."""
|
|
|
|
def __getitem__(self, index):
|
|
return {"x": self.x[index, None], "y": 2}
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"dataset", (LinearDataset, LinearDatasetDict, NonTensorDataset)
|
|
)
|
|
def test_torch_prepare_dataloader(ray_start_4_cpus_2_gpus, dataset):
|
|
data_loader = DataLoader(dataset(a=1, b=2, size=10))
|
|
|
|
def train_fn():
|
|
wrapped_data_loader = ray.train.torch.prepare_data_loader(data_loader)
|
|
|
|
# Check that DistributedSampler has been added to the data loader.
|
|
assert isinstance(wrapped_data_loader.sampler, DistributedSampler)
|
|
|
|
# Make sure you can properly iterate through the DataLoader.
|
|
# Case where the dataset returns a tuple or list from __getitem__.
|
|
if isinstance(dataset, LinearDataset):
|
|
for batch in wrapped_data_loader:
|
|
x = batch[0]
|
|
y = batch[1]
|
|
|
|
# Make sure the data is on the correct device.
|
|
assert x.is_cuda and y.is_cuda
|
|
# Case where the dataset returns a dict from __getitem__.
|
|
elif isinstance(dataset, LinearDatasetDict):
|
|
for batch in wrapped_data_loader:
|
|
for x, y in zip(batch["x"], batch["y"]):
|
|
# Make sure the data is on the correct device.
|
|
assert x.is_cuda and y.is_cuda
|
|
|
|
elif isinstance(dataset, NonTensorDataset):
|
|
for batch in wrapped_data_loader:
|
|
for x, y in zip(batch["x"], batch["y"]):
|
|
# Make sure the data is on the correct device.
|
|
assert x.is_cuda and y == 2
|
|
|
|
trainer = TorchTrainer(
|
|
train_fn, scaling_config=ScalingConfig(num_workers=2, use_gpu=True)
|
|
)
|
|
trainer.fit()
|
|
|
|
|
|
def test_torch_fail_on_nccl_timeout(ray_start_4_cpus_2_gpus):
|
|
"""Tests that TorchTrainer raises exception on NCCL timeouts."""
|
|
|
|
def train_fn():
|
|
model = torch.nn.Linear(1, 1)
|
|
model = ray.train.torch.prepare_model(model)
|
|
|
|
# Rank 0 worker will never reach the collective operation.
|
|
# NCCL should timeout.
|
|
if ray.train.get_context().get_world_rank() == 0:
|
|
while True:
|
|
time.sleep(100)
|
|
|
|
torch.distributed.barrier()
|
|
|
|
trainer = TorchTrainer(
|
|
train_fn,
|
|
scaling_config=ScalingConfig(num_workers=2, use_gpu=True),
|
|
torch_config=ray.train.torch.TorchConfig(timeout_s=2),
|
|
)
|
|
|
|
# Training should fail and not hang.
|
|
with pytest.raises(WorkerGroupError):
|
|
trainer.fit()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", "-x", __file__]))
|