wehub-resource-sync
251 lines
7.7 KiB
Python
251 lines
7.7 KiB
Python
import pytest
|
|
import torch
|
|
|
|
import ray
|
|
from ray.train import RunConfig, ScalingConfig
|
|
from ray.train.constants import TORCH_PROCESS_GROUP_SHUTDOWN_TIMEOUT_S
|
|
from ray.train.examples.pytorch.torch_linear_example import (
|
|
train_func as linear_train_func,
|
|
)
|
|
from ray.train.torch import TorchConfig, TorchTrainer
|
|
from ray.train.torch.config import _is_backend_nccl
|
|
from ray.train.v2._internal.constants import HEALTH_CHECK_INTERVAL_S_ENV_VAR
|
|
from ray.train.v2.api.config import FailureConfig
|
|
from ray.train.v2.api.exceptions import WorkerGroupError
|
|
from ray.train.v2.torch.torchft_config import TorchftConfig
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def ray_start_4_cpus():
|
|
ray.init(num_cpus=4)
|
|
yield
|
|
ray.shutdown()
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def reduce_health_check_interval(monkeypatch):
|
|
monkeypatch.setenv(HEALTH_CHECK_INTERVAL_S_ENV_VAR, "0.2")
|
|
yield
|
|
|
|
|
|
def test_minimal(ray_start_4_cpus):
|
|
def train_func():
|
|
pass
|
|
|
|
trainer = TorchTrainer(train_func)
|
|
trainer.fit()
|
|
|
|
|
|
@pytest.mark.parametrize("num_workers", [1, 2])
|
|
def test_torch_linear(ray_start_4_cpus, num_workers):
|
|
def train_func(config):
|
|
result = linear_train_func(config)
|
|
assert len(result) == epochs
|
|
assert result[-1]["loss"] < result[0]["loss"]
|
|
|
|
num_workers = num_workers
|
|
epochs = 3
|
|
scaling_config = ScalingConfig(num_workers=num_workers)
|
|
config = {"lr": 1e-2, "hidden_size": 1, "batch_size": 4, "epochs": epochs}
|
|
trainer = TorchTrainer(
|
|
train_loop_per_worker=train_func,
|
|
train_loop_config=config,
|
|
scaling_config=scaling_config,
|
|
)
|
|
trainer.fit()
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"torch_config,expected_world_size",
|
|
[
|
|
(TorchConfig(backend="gloo", init_method="env"), 2),
|
|
(TorchConfig(backend="gloo", init_method="tcp"), 2),
|
|
(
|
|
TorchftConfig(
|
|
backend="gloo", init_method="env", lighthouse_kwargs={"min_replicas": 1}
|
|
),
|
|
1,
|
|
),
|
|
(
|
|
TorchftConfig(
|
|
backend="gloo", init_method="tcp", lighthouse_kwargs={"min_replicas": 1}
|
|
),
|
|
1,
|
|
),
|
|
],
|
|
)
|
|
def test_torch_start_shutdown(ray_start_4_cpus, torch_config, expected_world_size):
|
|
def check_process_group():
|
|
assert (
|
|
torch.distributed.is_initialized()
|
|
and torch.distributed.get_world_size() == expected_world_size
|
|
)
|
|
|
|
trainer = TorchTrainer(
|
|
train_loop_per_worker=check_process_group,
|
|
scaling_config=ScalingConfig(num_workers=2),
|
|
torch_config=torch_config,
|
|
)
|
|
trainer.fit()
|
|
|
|
|
|
@pytest.mark.parametrize("timeout_s", [5, 0])
|
|
def test_torch_process_group_shutdown_timeout(ray_start_4_cpus, monkeypatch, timeout_s):
|
|
"""Tests that we don't more than a predefined timeout
|
|
on Torch process group shutdown."""
|
|
|
|
monkeypatch.setenv(TORCH_PROCESS_GROUP_SHUTDOWN_TIMEOUT_S, str(timeout_s))
|
|
|
|
trainer = TorchTrainer(
|
|
train_loop_per_worker=lambda: None,
|
|
scaling_config=ScalingConfig(num_workers=2),
|
|
torch_config=TorchConfig(backend="gloo"),
|
|
)
|
|
# Even if shutdown times out (timeout_s=0),
|
|
# the training should complete successfully.
|
|
trainer.fit()
|
|
|
|
|
|
def test_torchft_linear(ray_start_4_cpus):
|
|
"""Test torchft linear training: loss goes down and models are equal across workers."""
|
|
|
|
from ray.train.v2.examples.pytorch.torchft_linear_example import (
|
|
train_func as torchft_linear_train_func,
|
|
)
|
|
|
|
@ray.remote
|
|
class WeightCollector:
|
|
def __init__(self):
|
|
self.weights = {}
|
|
|
|
def report(self, rank, weight, bias):
|
|
self.weights[rank] = {"weight": weight, "bias": bias}
|
|
|
|
def get_weights(self):
|
|
return self.weights
|
|
|
|
collector = WeightCollector.remote()
|
|
|
|
def train_func(config):
|
|
result = torchft_linear_train_func(config)
|
|
assert result[-1]["loss"] < result[0]["loss"]
|
|
world_rank = ray.train.get_context().get_world_rank()
|
|
ray.get(
|
|
config["collector"].report.remote(
|
|
world_rank, result[-1]["weight"], result[-1]["bias"]
|
|
)
|
|
)
|
|
|
|
trainer = TorchTrainer(
|
|
train_loop_per_worker=train_func,
|
|
train_loop_config={"collector": collector},
|
|
scaling_config=ScalingConfig(num_workers=2),
|
|
torch_config=TorchftConfig(
|
|
backend="gloo", lighthouse_kwargs={"min_replicas": 2}
|
|
),
|
|
)
|
|
result = trainer.fit()
|
|
assert result.error is None
|
|
|
|
# Check that models converged across workers.
|
|
weights = ray.get(collector.get_weights.remote())
|
|
assert len(weights) == 2
|
|
assert weights[0]["weight"] == pytest.approx(weights[1]["weight"], abs=1e-4)
|
|
assert weights[0]["bias"] == pytest.approx(weights[1]["bias"], abs=1e-4)
|
|
|
|
|
|
# TODO(tseah): Add test for lighthouse failures (e.g. lighthouse unreachable).
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"min_replicas,max_failures,expect_error,expected_train_fn_calls,expected_reports",
|
|
[
|
|
# TODO(tseah): enable this after we support training with 1/2 workers without
|
|
# trying to restart the replica group.
|
|
# (1, 0, False, 2),
|
|
# This continues training with 1 replica. It does not replay step 10 and its report.
|
|
# TODO(tseah): expected_reports should be 1 when we support training with 1/2 workers.
|
|
(1, 1, False, 3, 0),
|
|
# This errors immediately.
|
|
(2, 0, True, 2, 0),
|
|
# This stops training with 1 replica. It replays step 10 and its report.
|
|
(2, 1, False, 3, 2),
|
|
],
|
|
)
|
|
def test_torchft_linear_replica_failure(
|
|
ray_start_4_cpus,
|
|
min_replicas,
|
|
max_failures,
|
|
expect_error,
|
|
expected_train_fn_calls,
|
|
expected_reports,
|
|
):
|
|
"""Test torchft linear training behavior when a replica fails mid-training."""
|
|
|
|
from ray.train.v2.examples.pytorch.torchft_linear_example import (
|
|
train_func as torchft_linear_train_func,
|
|
)
|
|
|
|
num_workers = 2
|
|
# TODO(tseah): remove this check once we support training with 1/2 workers.
|
|
training_requires_all_workers = min_replicas == num_workers
|
|
|
|
@ray.remote
|
|
class Counter:
|
|
def __init__(self):
|
|
self.count = 0
|
|
|
|
def increment(self):
|
|
self.count += 1
|
|
return self.count
|
|
|
|
def get_count(self):
|
|
return self.count
|
|
|
|
counter = Counter.remote()
|
|
|
|
def train_fn(config):
|
|
counter.increment.remote()
|
|
return torchft_linear_train_func(config)
|
|
|
|
trainer = TorchTrainer(
|
|
train_loop_per_worker=train_fn,
|
|
train_loop_config={
|
|
"num_steps": 20,
|
|
"error_step": 10,
|
|
"error_rank": 0,
|
|
"num_replicas": min_replicas,
|
|
"training_requires_all_workers": training_requires_all_workers,
|
|
},
|
|
scaling_config=ScalingConfig(num_workers=num_workers),
|
|
torch_config=TorchftConfig(
|
|
backend="gloo", lighthouse_kwargs={"min_replicas": min_replicas}
|
|
),
|
|
run_config=RunConfig(failure_config=FailureConfig(max_failures=max_failures)),
|
|
)
|
|
if expect_error:
|
|
with pytest.raises(WorkerGroupError):
|
|
trainer.fit()
|
|
else:
|
|
result = trainer.fit()
|
|
assert len(result.best_checkpoints) == expected_reports
|
|
assert result.error is None
|
|
# Fewer train_fn calls indicate partial worker group restarts.
|
|
assert ray.get(counter.get_count.remote()) == expected_train_fn_calls
|
|
|
|
|
|
def test_is_backend_nccl():
|
|
assert _is_backend_nccl("nccl")
|
|
assert _is_backend_nccl("cuda:nccl")
|
|
assert _is_backend_nccl("cpu:gloo,cuda:nccl")
|
|
assert not _is_backend_nccl("gloo")
|
|
assert not _is_backend_nccl("cpu:nccl")
|
|
assert not _is_backend_nccl("cuda:gloo")
|
|
assert not _is_backend_nccl("cpu:gloo,cuda:gloo")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", "-x", __file__]))
|