def select_strategy(self) -> Strategy:
if isinstance(self.distributed_backend, Accelerator) and self.distributed_backend.strategy is not None:
plugin = self.distributed_backend.strategy
elif self.use_ddp2:
plugin = DDP2Strategy(parallel_devices=self.parallel_devices, cluster_environment=self.cluster_environment)
elif self.use_ddp and self.use_deepspeed:
plugin = DeepSpeedStrategy(
cluster_environment=self.select_cluster_environment(), parallel_devices=self.parallel_devices
)
elif self.use_ddp:
use_slurm_ddp = self.use_ddp and self._is_slurm_managing_tasks()
use_torchelastic_ddp = self.use_ddp and TorchElasticEnvironment.detect()
use_kubeflow_ddp = self.use_ddp and KubeflowEnvironment.detect()
use_ddp_spawn = self._strategy_type == _StrategyType.DDP_SPAWN
use_ddp_cpu_spawn = use_ddp_spawn and self.use_cpu
use_tpu_spawn = self.use_tpu and self._strategy_type == _StrategyType.TPU_SPAWN
use_ddp_cpu_torch_elastic = use_ddp_cpu_spawn and TorchElasticEnvironment.detect()
use_ddp_cpu_kubeflow = use_ddp_cpu_spawn and KubeflowEnvironment.detect()
use_ddp_cpu_slurm = use_ddp_cpu_spawn and self._is_slurm_managing_tasks()
use_ddp_sharded = self._strategy_type == _StrategyType.DDP_SHARDED
use_ddp_sharded_spawn = self._strategy_type == _StrategyType.DDP_SHARDED_SPAWN
use_ddp_fully_sharded = self._strategy_type == _StrategyType.DDP_FULLY_SHARDED
if use_tpu_spawn:
ddp_strategy_cls = TPUSpawnStrategy
elif use_ddp_sharded:
ddp_strategy_cls = DDPShardedStrategy
elif use_ddp_sharded_spawn:
ddp_strategy_cls = DDPSpawnShardedStrategy
elif (
use_ddp_cpu_slurm
or use_slurm_ddp
or use_ddp_cpu_torch_elastic
or use_torchelastic_ddp
or use_kubeflow_ddp
or use_ddp_cpu_kubeflow
):
ddp_strategy_cls = DDPStrategy
elif use_ddp_spawn or use_ddp_cpu_spawn:
ddp_strategy_cls = DDPSpawnStrategy
elif use_ddp_fully_sharded:
ddp_strategy_cls = DDPFullyShardedStrategy
else:
ddp_strategy_cls = DDPStrategy
plugin = ddp_strategy_cls(
parallel_devices=self.parallel_devices, cluster_environment=self.cluster_environment
)
elif self.use_dp:
plugin = DataParallelStrategy(parallel_devices=self.parallel_devices)
elif self.use_horovod:
plugin = HorovodStrategy(parallel_devices=self.parallel_devices)
elif self.use_tpu and isinstance(self.tpu_cores, list):
plugin = SingleTPUStrategy(self.tpu_id)
elif self.use_ipu:
plugin = IPUStrategy(parallel_devices=self.parallel_devices)
else:
single_gpu_ordinal = device_parser.determine_root_gpu_device(self.parallel_device_ids)
plugin = SingleDeviceStrategy(device=single_gpu_ordinal if self.use_gpu else "cpu")
return plugin
def test_unsupported_ddp2_strategy():
with pytest.raises(
TypeError,
match=
"The `DDP2Strategy`/`DDP2Plugin` is no longer supported in v1.7 and will be"
):
DDP2Strategy()
with pytest.raises(
TypeError,
match=
"The `DDP2Strategy`/`DDP2Plugin` is no longer supported in v1.7 and will be"
):
DDP2Plugin()
with pytest.raises(ValueError,
match="The DDP2 strategy is no longer supported."):
Trainer(strategy="ddp2")
@mock.patch.dict(
os.environ,
{
"CUDA_VISIBLE_DEVICES": "0,1",
"SLURM_NTASKS": "2",
"SLURM_JOB_NAME": "SOME_NAME",
"SLURM_NODEID": "0",
"SLURM_PROCID": "1",
"SLURM_LOCALID": "1",
},
)
@mock.patch("torch.cuda.set_device")
@mock.patch("torch.cuda.device_count", return_value=2)
@mock.patch("pytorch_lightning.strategies.DDPStrategy.setup_distributed",
autospec=True)
@pytest.mark.parametrize("strategy", ["ddp2", DDP2Strategy()])
def test_strategy_choice_ddp2_slurm(set_device_mock, device_count_mock,
setup_distributed_mock, strategy):
trainer = Trainer(fast_dev_run=True, strategy=strategy, gpus=2)
assert trainer._accelerator_connector._is_slurm_managing_tasks()
assert isinstance(trainer.accelerator, GPUAccelerator)
assert isinstance(trainer.strategy, DDP2Strategy)
assert isinstance(trainer.strategy.cluster_environment, SLURMEnvironment)
assert trainer.strategy.cluster_environment.local_rank() == 1
assert trainer.strategy.local_rank == 1
@mock.patch.dict(
os.environ,
{
"CUDA_VISIBLE_DEVICES": "0,1",