Set default strategy to ddp_fork in interactive environments (#13746)

This commit is contained in:
Adrian Wälchli 2022-07-22 21:34:30 +02:00 committed by GitHub
parent 9f51c07604
commit f6f06d4e42
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 16 additions and 1 deletions

View File

@ -87,6 +87,7 @@ from pytorch_lightning.utilities.imports import (
_HOROVOD_AVAILABLE,
_HPU_AVAILABLE,
_IPU_AVAILABLE,
_IS_INTERACTIVE,
_TORCH_GREATER_EQUAL_1_11,
_TPU_AVAILABLE,
)
@ -588,7 +589,9 @@ class AcceleratorConnector:
# TODO: lazy initialized device, then here could be self._strategy_flag = "single_device"
return SingleDeviceStrategy(device=device) # type: ignore
if len(self._parallel_devices) > 1:
return DDPSpawnStrategy.strategy_name
if _IS_INTERACTIVE:
return "ddp_fork"
return "ddp_spawn"
return DDPStrategy.strategy_name

View File

@ -424,6 +424,18 @@ def test_strategy_choice_ddp_spawn_cpu():
assert trainer.strategy.launcher._start_method == "spawn"
@RunIf(skip_windows=True)
@mock.patch("pytorch_lightning.trainer.connectors.accelerator_connector._IS_INTERACTIVE", True)
def test_strategy_choice_ddp_fork_in_interactive():
"""Test that when accelerator and strategy are unspecified, the connector chooses DDP Fork in interactive
environments by default."""
trainer = Trainer(devices=2)
assert isinstance(trainer.accelerator, CPUAccelerator)
assert isinstance(trainer.strategy, DDPSpawnStrategy)
assert isinstance(trainer.strategy.cluster_environment, LightningEnvironment)
assert trainer.strategy.launcher._start_method == "fork"
@RunIf(skip_windows=True)
def test_strategy_choice_ddp_fork_cpu():
trainer = Trainer(strategy="ddp_fork", accelerator="cpu", devices=2)