[accelerator][BugFix] Resolve some test for 1 gpu

pytorch_lightning/accelerators/accelerator_connector.py

@@ -116,11 +116,11 @@ def __init__(
  self.parallel_device_ids = device_parser.parse_gpu_ids(self.gpus)
  self.root_gpu = device_parser.determine_root_gpu_device(self.parallel_device_ids)
 
- self.handle_given_plugins(plugins)
-
  self.set_distributed_mode()
  self.configure_slurm_ddp()
 
+ self.handle_given_plugins(plugins)
+
  self.accelerator = self.select_accelerator()
 
  # override dist backend when using tpus
@@ -148,6 +148,7 @@ def __init__(
 
  def handle_given_plugins(self, plugins: Optional[Sequence]):
  if plugins is None:
+ self._cluster_environment = self.select_cluster_environment()
  return
 
  if not isinstance(plugins, Sequence):
@@ -481,7 +482,7 @@ def set_distributed_mode(self):
  # for DDP overwrite nb processes by requested GPUs
  if (
  self._device_type == DeviceType.GPU
- and self._distrib_type in (DistributedType.DDP, DistributedType.DDP_SPAWN, DistributedType.DDP2)
+ and self._distrib_type in (DistributedType.DDP, DistributedType.DDP_SPAWN)
  ):
  self.num_processes = self.num_gpus
 

pytorch_lightning/plugins/precision/apex_amp.py

@@ -11,7 +11,7 @@
 # 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.
-from typing import List, Tuple
+from typing import List, Tuple, Callable
 
 import torch
 from torch.optim import Optimizer
@@ -71,7 +71,7 @@ def backward(
  # do backward pass
  # TODO: not entirely sure, why we need this
  if model is not None and isinstance(model, LightningModule):
- model.backward(closure_loss, optimizer, opt_idx)
+ model.backward(closure_loss, optimizer, opt_idx, **kwargs)
 
  # TODO: avoid dev_debugger and track these calls with mock
  model.trainer.dev_debugger.track_event('AMP', str(AMPType.APEX))
@@ -90,6 +90,16 @@ def backward(
  closure_loss = closure_loss.detach()
  return closure_loss
 
+ def pre_optimizer_step(
+ self, pl_module: LightningModule, optimizer: Optimizer, optimizer_idx: int, closure: Callable, **kwargs
+ ) -> bool:
+ """Hook to do something before each optimizer step."""
+ # Apex: Amp does not support closure use with optimizers
+ closure()
+ optimizer.step()
+ return False
+
+
  def configure_apex(
  self,
  amp: object,

pytorch_lightning/trainer/trainer.py

@@ -458,6 +458,7 @@ def fit(
  # ----------------------------
  # SET UP TRAINING
  # ----------------------------
+ self.call_hook("on_before_accelerator_backend_setup", model)
  self.accelerator_backend.setup(self, model)
  self.setup_trainer(model)
 
@@ -469,7 +470,6 @@ def fit(
 
  # plugin will setup training (e.g. ddp will launch child processes)
  # TODO: the old setup is now called "pre_training", where should this hook be called now?
- self.call_hook("on_before_accelerator_backend_setup", model)
  self.training_type_plugin.pre_training()
  self.precision_plugin.pre_training()
 

tests/accelerators/legacy/test_accelerator_connector.py

@@ -75,7 +75,7 @@ def test_accelerator_choice_ddp_spawn(cuda_available_mock, device_count_mock):
  assert isinstance(trainer.training_type_plugin.cluster_environment, TorchElasticEnvironment)
 
 
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="test requires GPU")
+@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
 @mock.patch.dict(
  os.environ, {
  "CUDA_VISIBLE_DEVICES": "0,1",
@@ -89,13 +89,12 @@ def test_accelerator_choice_ddp_slurm():
 
  class CB(Callback):
 
- def on_fit_start(self, trainer, pl_module):
+ def on_before_accelerator_backend_setup(self, trainer, pl_module):
  assert trainer.use_ddp
  assert trainer.accelerator_connector.is_slurm_managing_tasks
  assert isinstance(trainer.accelerator_backend, GPUAccelerator)
  assert isinstance(trainer.training_type_plugin, DDPPlugin)
  assert isinstance(trainer.training_type_plugin.cluster_environment, SLURMEnvironment)
- assert trainer.training_type_plugin.task_idx == 10
  assert trainer.training_type_plugin.cluster_environment.local_rank() == 10
  raise SystemExit()
 
@@ -127,13 +126,12 @@ def test_accelerator_choice_ddp2_slurm(device_count_mock):
 
  class CB(Callback):
 
- def on_fit_start(self, trainer, pl_module):
+ def on_before_accelerator_backend_setup(self, trainer, pl_module):
  assert trainer.use_ddp2
  assert trainer.accelerator_connector.is_slurm_managing_tasks
  assert isinstance(trainer.accelerator_backend, GPUAccelerator)
  assert isinstance(trainer.training_type_plugin, DDP2Plugin)
  assert isinstance(trainer.training_type_plugin.cluster_environment, SLURMEnvironment)
- assert trainer.training_type_plugin.task_idx == 10
  assert trainer.training_type_plugin.cluster_environment.local_rank() == 10
 
  raise SystemExit()
@@ -157,12 +155,11 @@ def test_accelerator_choice_ddp_te(device_count_mock):
 
  class CB(Callback):
 
- def on_fit_start(self, trainer, pl_module):
+ def on_before_accelerator_backend_setup(self, trainer, pl_module):
  assert trainer.use_ddp
  assert isinstance(trainer.accelerator_backend, GPUAccelerator)
  assert isinstance(trainer.training_type_plugin, DDPPlugin)
  assert isinstance(trainer.training_type_plugin.cluster_environment, TorchElasticEnvironment)
- assert trainer.training_type_plugin.task_idx == 10
  assert trainer.training_type_plugin.cluster_environment.local_rank() == 10
  raise SystemExit()
 
@@ -185,12 +182,11 @@ def test_accelerator_choice_ddp2_te(device_count_mock):
 
  class CB(Callback):
 
- def on_fit_start(self, trainer, pl_module):
+ def on_before_accelerator_backend_setup(self, trainer, pl_module):
  assert trainer.use_ddp2
  assert isinstance(trainer.accelerator_backend, GPUAccelerator)
  assert isinstance(trainer.training_type_plugin, DDP2Plugin)
  assert isinstance(trainer.training_type_plugin.cluster_environment, TorchElasticEnvironment)
- assert trainer.training_type_plugin.task_idx == 10
  assert trainer.training_type_plugin.cluster_environment.local_rank() == 10
  raise SystemExit()
 
@@ -216,12 +212,11 @@ def test_accelerator_choice_ddp_cpu_te(device_count_mock):
 
  class CB(Callback):
 
- def on_fit_start(self, trainer, pl_module):
+ def on_before_accelerator_backend_setup(self, trainer, pl_module):
  assert trainer.use_ddp
  assert isinstance(trainer.accelerator_backend, CPUAccelerator)
  assert isinstance(trainer.training_type_plugin, DDPPlugin)
  assert isinstance(trainer.training_type_plugin.cluster_environment, TorchElasticEnvironment)
- assert trainer.training_type_plugin.task_idx == 10
  assert trainer.training_type_plugin.cluster_environment.local_rank() == 10
  raise SystemExit()
 
@@ -251,7 +246,7 @@ def test_accelerator_choice_ddp_cpu_slurm(device_count_mock):
 
  class CB(Callback):
 
- def on_fit_start(self, trainer, pl_module):
+ def on_before_accelerator_backend_setup(self, trainer, pl_module):
  assert trainer.use_ddp
  assert trainer.accelerator_connector.is_slurm_managing_tasks
  assert isinstance(trainer.accelerator_backend, CPUAccelerator)
@@ -293,7 +288,7 @@ def master_address(self):
 
  class CB(Callback):
 
- def on_fit_start(self, trainer, pl_module):
+ def on_before_accelerator_backend_setup(self, trainer, pl_module):
  assert trainer.use_ddp
  assert isinstance(trainer.accelerator_backend, CPUAccelerator)
  assert isinstance(trainer.training_type_plugin, DDPPlugin)
@@ -362,7 +357,7 @@ def test_dist_backend_accelerator_mapping(device_count_mock):
 
  class CB(Callback):
 
- def on_fit_start(self, trainer, pl_module):
+ def on_before_accelerator_backend_setup(self, trainer, pl_module):
  assert isinstance(trainer.accelerator_backend, CPUAccelerator)
  assert isinstance(trainer.training_type_plugin, DDPPlugin)
  raise SystemExit()

tests/callbacks/test_callbacks.py

@@ -53,8 +53,8 @@ def test_trainer_callback_system(torch_save):
  assert callback_mock.method_calls == [
  call.on_init_start(trainer),
  call.on_init_end(trainer),
- call.on_fit_start(trainer, model),
  call.on_before_accelerator_backend_setup(trainer, model),
+ call.on_fit_start(trainer, model),
  call.setup(trainer, model, 'fit'),
  call.on_pretrain_routine_start(trainer, model),
  call.on_pretrain_routine_end(trainer, model),
@@ -108,8 +108,8 @@ def test_trainer_callback_system(torch_save):
  assert callback_mock.method_calls == [
  call.on_init_start(trainer),
  call.on_init_end(trainer),
- call.on_fit_start(trainer, model),
  call.on_before_accelerator_backend_setup(trainer, model),
+ call.on_fit_start(trainer, model),
  call.setup(trainer, model, 'test'),
  call.on_test_start(trainer, model),
  call.on_test_epoch_start(trainer, model),

tests/deprecated_api/test_remove_1-4.py

@@ -163,7 +163,7 @@ def configure_ddp(self):
  assert isinstance(self.model.module, LightningDistributedModule)
 
 
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="test requires GPU machine")
+@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
 @pytest.mark.skipif(sys.platform == "win32", reason="DDP not available on windows")
 def test_v1_4_0_deprecated_lightning_distributed_data_parallel(tmpdir):
  model = BoringModel()

tests/plugins/legacy/test_rpc_plugin.py

@@ -33,7 +33,7 @@ def test_rpc_choice(tmpdir, ddp_backend, gpus, num_processes):
 
  class CB(Callback):
 
- def on_fit_start(self, trainer, pl_module):
+ def on_before_accelerator_backend_setup(self, trainer, pl_module):
  assert isinstance(trainer.training_type_plugin, RPCPlugin)
  raise RuntimeError('finished plugin check')
 

tests/plugins/test_apex_plugin.py

@@ -30,7 +30,7 @@ def test_amp_choice_default_ddp_cpu(tmpdir, ddp_backend, gpus, num_processes):
 
  class CB(Callback):
 
- def on_fit_start(self, trainer, pl_module):
+ def on_before_accelerator_backend_setup(self, trainer, pl_module):
  assert isinstance(trainer.precision_plugin, ApexMixedPrecisionPlugin)
  raise SystemExit()
 
@@ -72,7 +72,7 @@ class MyApexPlugin(ApexMixedPrecisionPlugin):
 
  class CB(Callback):
 
- def on_fit_start(self, trainer, pl_module):
+ def on_before_accelerator_backend_setup(self, trainer, pl_module):
  assert isinstance(trainer.precision_plugin, MyApexPlugin)
  raise SystemExit()
 

tests/plugins/test_sharded_plugin.py

@@ -21,7 +21,7 @@ def test_sharded_ddp_choice(tmpdir, accelerator):
 
  class CB(Callback):
 
- def on_fit_start(self, trainer, pl_module):
+ def on_before_accelerator_backend_setup(self, trainer, pl_module):
  if accelerator == 'ddp_sharded':
  assert isinstance(trainer.accelerator_backend.training_type_plugin, DDPShardedPlugin)
  elif accelerator == 'ddp_sharded_spawn':
@@ -68,7 +68,7 @@ def test_ddp_choice_sharded_amp(tmpdir, accelerator):
 
  class CB(Callback):
 
- def on_fit_start(self, trainer, pl_module):
+ def on_before_accelerator_backend_setup(self, trainer, pl_module):
  assert isinstance(trainer.accelerator_backend.precision_plugin, ShardedNativeMixedPrecisionPlugin)
  raise SystemExit()
 

tests/trainer/optimization/test_manual_optimization.py

@@ -346,7 +346,7 @@ def training_step(self, batch, batch_idx, optimizer_idx):
  # ensure we forward the correct params to the optimizer
  # without retain_graph we can't do multiple backward passes
  self.manual_backward(loss_2, opt_b, retain_graph=True)
- self.manual_backward(loss_2, opt_a, retain_graph=True)
+ self.manual_backward(loss_2, opt_a)
 
  assert self.layer.weight.grad is not None
  opt_b.step()