Rooba
/
lightning
mirror of https://github.com/Lightning-AI/lightning.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
src/pytorch_lightning/trainer/connectors/accelerator_connector.py
View File

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

12
tests/tests_pytorch/accelerators/test_accelerator_connector.py
View File

@ -424,6 +424,18 @@ def test_strategy_choice_ddp_spawn_cpu():
assert trainer.strategy.launcher._start_method == "spawn" 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) @RunIf(skip_windows=True)
def test_strategy_choice_ddp_fork_cpu(): def test_strategy_choice_ddp_fork_cpu():
trainer = Trainer(strategy="ddp_fork", accelerator="cpu", devices=2) trainer = Trainer(strategy="ddp_fork", accelerator="cpu", devices=2)