lightning/tests/models/test_horovod.py

457 lines
15 KiB
Python

# Copyright The PyTorch Lightning team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import json
import os
import shlex
import subprocess
import sys
from unittest.mock import patch
import numpy as np
import pytest
import torch
from sklearn.metrics import accuracy_score
from torch import optim
from torchmetrics.classification.accuracy import Accuracy
import tests.helpers.pipelines as tpipes
import tests.helpers.utils as tutils
from pytorch_lightning import Trainer
from pytorch_lightning.accelerators import CPUAccelerator
from pytorch_lightning.utilities import _HOROVOD_AVAILABLE
from pytorch_lightning.utilities.exceptions import MisconfigurationException
from tests.helpers import BoringModel
from tests.helpers.advanced_models import BasicGAN
from tests.helpers.runif import RunIf
if _HOROVOD_AVAILABLE:
import horovod
import horovod.torch as hvd
@RunIf(min_gpus=1, horovod=True)
@pytest.mark.xfail(reason="FIXME(@Borda): nccl is not available in the GPU image")
def test_nccl_is_available_on_gpu_environment():
from tests.helpers.runif import _HOROVOD_NCCL_AVAILABLE
# the GPU environment should always install Horovod NCCL
assert _HOROVOD_NCCL_AVAILABLE
# This script will run the actual test model training in parallel
TEST_SCRIPT = os.path.join(os.path.dirname(__file__), "data", "horovod", "train_default_model.py")
def _run_horovod(trainer_options):
"""Execute the training script across multiple workers in parallel."""
devices = trainer_options.get("devices", 1)
tutils.reset_seed()
# TODO: Find out why coverage breaks CI.
# append = '-a' if '.coverage' in os.listdir(_PROJECT_ROOT) else ''
# str(num_processes), sys.executable, '-m', 'coverage', 'run', '--source', 'pytorch_lightning', append,
cmdline = [
"horovodrun",
"-np",
str(devices),
sys.executable,
TEST_SCRIPT,
"--trainer-options",
shlex.quote(json.dumps(trainer_options)),
]
if trainer_options.get("accelerator", "cpu") == "gpu":
cmdline += ["--on-gpu"]
exit_code = subprocess.call(" ".join(cmdline), shell=True, env=os.environ.copy())
assert exit_code == 0
@RunIf(skip_windows=True, horovod=True, skip_49370=True)
def test_horovod_cpu(tmpdir):
"""Test Horovod running multi-process on CPU."""
trainer_options = dict(
default_root_dir=str(tmpdir),
gradient_clip_val=1.0,
enable_progress_bar=False,
max_epochs=1,
limit_train_batches=0.4,
limit_val_batches=0.2,
strategy="horovod",
)
_run_horovod(trainer_options)
@RunIf(skip_windows=True, horovod=True, skip_49370=True)
def test_horovod_cpu_accumulate_grad_batches(tmpdir):
trainer_options = dict(
default_root_dir=tmpdir,
enable_progress_bar=False,
max_epochs=1,
limit_train_batches=4,
limit_val_batches=0,
accumulate_grad_batches=2,
strategy="horovod",
)
_run_horovod(trainer_options)
@RunIf(skip_windows=True, horovod=True, skip_49370=True)
def test_horovod_cpu_clip_grad_by_value(tmpdir):
"""Test Horovod running multi-process on CPU."""
trainer_options = dict(
default_root_dir=str(tmpdir),
gradient_clip_val=1.0,
gradient_clip_algorithm="value",
enable_progress_bar=False,
max_epochs=1,
limit_train_batches=0.4,
limit_val_batches=0.2,
strategy="horovod",
)
_run_horovod(trainer_options)
@RunIf(skip_windows=True, horovod=True, skip_49370=True)
def test_horovod_cpu_implicit(tmpdir):
"""Test Horovod without specifying a backend, inferring from env set by `horovodrun`."""
trainer_options = dict(
default_root_dir=str(tmpdir),
gradient_clip_val=1.0,
enable_progress_bar=False,
max_epochs=1,
limit_train_batches=0.4,
limit_val_batches=0.2,
)
_run_horovod(trainer_options)
@RunIf(min_gpus=2, skip_windows=True, horovod_nccl=True)
def test_horovod_multi_gpu(tmpdir):
"""Test Horovod with multi-GPU support."""
trainer_options = dict(
default_root_dir=str(tmpdir),
gradient_clip_val=1.0,
enable_progress_bar=False,
max_epochs=1,
limit_train_batches=0.4,
limit_val_batches=0.2,
accelerator="gpu",
devices=2,
strategy="horovod",
)
_run_horovod(trainer_options)
@RunIf(min_gpus=2, skip_windows=True, horovod_nccl=True)
def test_horovod_multi_gpu_accumulate_grad_batches(tmpdir):
trainer_options = dict(
default_root_dir=tmpdir,
enable_progress_bar=False,
max_epochs=1,
limit_train_batches=4,
limit_val_batches=0,
accumulate_grad_batches=2,
accelerator="gpu",
devices=2,
strategy="horovod",
)
_run_horovod(trainer_options)
@RunIf(horovod=True, skip_windows=True)
def test_horovod_raises_unsupported_accumulate_grad_batches(tmpdir):
"""Ensure MisConfigurationException for different `accumulate_grad_batches` at different epochs for Horovod
Strategy on multi-gpus."""
model = BoringModel()
trainer = Trainer(
default_root_dir=tmpdir,
enable_progress_bar=False,
accumulate_grad_batches={0: 4, 2: 2},
accelerator="auto",
devices=1,
strategy="horovod",
)
with pytest.raises(MisconfigurationException, match="Horovod.*does not support.*accumulate_grad_batches"):
trainer.fit(model)
@RunIf(min_gpus=2, skip_windows=True, horovod_nccl=True)
def test_horovod_multi_gpu_grad_by_value(tmpdir):
"""Test Horovod with multi-GPU support."""
trainer_options = dict(
default_root_dir=str(tmpdir),
gradient_clip_val=1.0,
gradient_clip_algorithm="value",
enable_progress_bar=False,
max_epochs=1,
limit_train_batches=0.4,
limit_val_batches=0.2,
accelerator="gpu",
devices=2,
strategy="horovod",
)
_run_horovod(trainer_options)
# todo: need to be fixed :]
# https://discuss.pytorch.org/t/torch-cuda-amp-vs-nvidia-apex/74994
# Check with (tgaddair) on Horovod issues if this feature is needed
@pytest.mark.skip(reason="TODO: Horovod currently doesn't work with Apex")
@RunIf(min_gpus=2, skip_windows=True, amp_apex=True, horovod_nccl=True)
def test_horovod_apex(tmpdir):
"""Test Horovod with multi-GPU support using apex amp."""
trainer_options = dict(
default_root_dir=str(tmpdir),
gradient_clip_val=1.0,
enable_progress_bar=False,
max_epochs=1,
limit_train_batches=0.4,
limit_val_batches=0.2,
accelerator="gpu",
devices=2,
strategy="horovod",
amp_backend="apex",
precision=16,
)
_run_horovod(trainer_options)
@RunIf(min_gpus=2, skip_windows=True, horovod_nccl=True)
def test_horovod_amp(tmpdir):
"""Test Horovod with multi-GPU support using native amp."""
trainer_options = dict(
default_root_dir=str(tmpdir),
gradient_clip_val=1.0,
enable_progress_bar=False,
max_epochs=1,
limit_train_batches=0.4,
limit_val_batches=0.2,
accelerator="gpu",
devices=2,
strategy="horovod",
amp_backend="native",
precision=16,
)
_run_horovod(trainer_options)
@RunIf(min_gpus=2, skip_windows=True, horovod_nccl=True)
def test_horovod_gather(tmpdir):
"""Test Horovod with multi-GPU support using native amp."""
trainer_options = dict(
default_root_dir=str(tmpdir),
gradient_clip_val=1.0,
enable_progress_bar=False,
max_epochs=1,
limit_train_batches=0.4,
limit_val_batches=0.2,
accelerator="gpu",
devices=2,
strategy="horovod",
)
_run_horovod(trainer_options)
@RunIf(min_gpus=1, skip_windows=True, horovod_nccl=True)
def test_horovod_transfer_batch_to_gpu(tmpdir):
class TestTrainingStepModel(BoringModel):
def training_step(self, batch, *args, **kwargs):
assert str(batch.device) != "cpu"
return super().training_step(batch, *args, **kwargs)
def validation_step(self, batch, *args, **kwargs):
assert str(batch.device) != "cpu"
return super().validation_step(batch, *args, **kwargs)
model = TestTrainingStepModel()
trainer_options = dict(
default_root_dir=str(tmpdir),
enable_progress_bar=False,
max_epochs=1,
limit_train_batches=0.4,
limit_val_batches=0.2,
accelerator="gpu",
devices=1,
strategy="horovod",
)
tpipes.run_model_test_without_loggers(trainer_options, model)
@RunIf(skip_windows=True, horovod=True)
def test_horovod_multi_optimizer(tmpdir):
model = BasicGAN()
# fit model
trainer = Trainer(
default_root_dir=str(tmpdir),
enable_progress_bar=False,
max_epochs=1,
limit_train_batches=0.4,
limit_val_batches=0.2,
strategy="horovod",
)
trainer.fit(model)
assert trainer.state.finished, f"Training failed with {trainer.state}"
assert len(trainer.optimizers) == 2
for i, optimizer in enumerate(trainer.optimizers):
assert hasattr(optimizer, "synchronize"), "optimizer has not been wrapped into DistributedOptimizer"
def get_model_params(model):
return set(list(model.parameters()))
def get_optimizer_params(optimizer):
return {p for group in optimizer.param_groups for p in group.get("params", [])}
assert get_model_params(model.generator) != get_model_params(model.discriminator)
assert get_model_params(model.generator) == get_optimizer_params(trainer.optimizers[0])
assert get_model_params(model.discriminator) == get_optimizer_params(trainer.optimizers[1])
# todo: need to be fixed :]
@pytest.mark.skip(reason="TODO: CI agent.jobstatus=Succeeded: Permission denied")
@RunIf(skip_windows=True, horovod=True)
def test_result_reduce_horovod(tmpdir):
"""Make sure result logging works with Horovod.
This test mirrors tests/core/test_results.py::_ddp_test_fn
"""
tutils.reset_seed()
tutils.set_random_main_port()
def hvd_test_fn():
path_here = os.path.abspath(os.path.dirname(__file__))
path_root = os.path.abspath(os.path.join(path_here, "..", ".."))
sys.path.insert(0, os.path.abspath(path_root))
class TestModel(BoringModel):
def training_step(self, batch, batch_idx):
self.training_step_called = True
tensor = torch.tensor([1.0])
self.log("test_tensor", tensor, sync_dist=True, reduce_fx="sum", on_step=True, on_epoch=True)
res = self._results
# Check that `tensor` is summed across all ranks automatically
assert (
res["test_tensor"].item() == hvd.size()
), "Result-Log does not work properly with Horovod and Tensors"
def training_epoch_end(self, outputs) -> None:
assert len(outputs) == 0
model = TestModel()
model.val_dataloader = None
trainer = Trainer(
default_root_dir=tmpdir,
limit_train_batches=2,
limit_val_batches=2,
max_epochs=1,
log_every_n_steps=1,
enable_model_summary=False,
logger=False,
)
trainer.fit(model)
horovod.run(hvd_test_fn, np=2)
# todo: need to be fixed :]
@pytest.mark.skip(reason="TODO: CI agent.jobstatus=Succeeded: Permission denied")
@RunIf(skip_windows=True, horovod=True, num_gpus=2)
def test_accuracy_metric_horovod():
num_batches = 10
batch_size = 16
threshold = 0.5
def sk_metric(preds, target):
sk_preds = (preds.view(-1).numpy() >= threshold).astype(np.uint8)
sk_target = target.view(-1).numpy()
return accuracy_score(y_true=sk_target, y_pred=sk_preds)
preds = torch.rand(num_batches, batch_size)
target = torch.randint(high=2, size=(num_batches, batch_size))
def _compute_batch():
trainer = Trainer(fast_dev_run=True, strategy="horovod", logger=False)
assert isinstance(trainer.accelerator, CPUAccelerator)
# TODO: test that we selected the correct training_type_plugin based on horovod flags
metric = Accuracy(
compute_on_step=True,
dist_sync_on_step=True,
dist_sync_fn=trainer.strategy.all_gather,
threshold=threshold,
)
for i in range(hvd.rank(), num_batches, hvd.size()):
batch_result = metric(preds[i], target[i])
if hvd.rank() == 0:
dist_preds = torch.stack([preds[i + r] for r in range(hvd.size())])
dist_target = torch.stack([target[i + r] for r in range(hvd.size())])
sk_batch_result = sk_metric(dist_preds, dist_target)
assert np.allclose(batch_result.numpy(), sk_batch_result)
# check on all batches on all ranks
result = metric.compute()
assert isinstance(result, torch.Tensor)
total_preds = torch.stack([preds[i] for i in range(num_batches)])
total_target = torch.stack([target[i] for i in range(num_batches)])
sk_result = sk_metric(total_preds, total_target)
assert np.allclose(result.numpy(), sk_result)
horovod.run(_compute_batch, np=2)
@RunIf(skip_windows=True, horovod=True)
def test_horovod_multi_optimizer_with_scheduling_stepping(tmpdir):
class TestModel(BoringModel):
def training_step(self, batch, batch_idx, optimizer_idx):
return super().training_step(batch, batch_idx)
def configure_optimizers(self):
optimizer1 = optim.Adam(self.parameters(), lr=0.1)
optimizer2 = optim.Adam(self.parameters(), lr=0.1)
lr_scheduler1 = optim.lr_scheduler.StepLR(optimizer1, 1, gamma=0.1)
lr_scheduler2 = optim.lr_scheduler.StepLR(optimizer2, 1, gamma=0.1)
return [optimizer1, optimizer2], [lr_scheduler1, lr_scheduler2]
model = TestModel()
model.training_epoch_end = None
num_workers = 8
init_lr = 0.1 * num_workers
with patch("horovod.torch.size", return_value=8):
# fit model
trainer = Trainer(
default_root_dir=tmpdir, max_epochs=1, limit_val_batches=0.5, limit_train_batches=0.2, strategy="horovod"
)
trainer.fit(model)
adjusted_lr1 = [pg["lr"] for pg in trainer.optimizers[0].param_groups][0]
adjusted_lr2 = [pg["lr"] for pg in trainer.optimizers[1].param_groups][0]
# Called ones after end of epoch with gamma=0.1
assert pytest.approx(init_lr * 0.1) == adjusted_lr1
# Called every 3 steps, meaning for 1 epoch of 11 batches, it is called 3 times with gamma=0.1
assert pytest.approx(init_lr * 0.1) == adjusted_lr2