PoC: Accelerator refactor (#5743)

* restoring the result from subprocess * fix queue.get() order for results * add missing "block_backward_sync" context manager * add missing "block_backward_sync" context manager * fix sync_batchnorm * fix supported gpu-ids for tuple * fix clip gradients and inf recursion * accelerator selection: added cluster_environment plugin * fix torchelastic test * fix reduce early stopping decision for DDP * fix tests: callbacks, conversion to lightning optimizer * fix lightning optimizer does not pickle * fix setting benchmark and deterministic option * fix slurm amp test * fix prepare_data test and determine node_rank * fix retrieving last path when testing * remove obsolete plugin argument * fix test: test_trainer_config * fix torchscript tests * fix trainer.model access * move properties * fix test_transfer_batch_hook * fix auto_select_gpus * fix omegaconf test * fix test that needs to simulate slurm ddp * add horovod plugin * fix test with named arguments * clean up whitespace * fix datamodules test * remove old accelerators * fix naming * move old plugins * move to plugins * create precision subpackage * create training_type subpackage * fix all new import errors * fix wrong arguments order passed to test * fix LR finder * Added sharded training type and amp plugin * Move clip grad to precision plugin * Added sharded spawn, select accelerators based on distributed_backend + enable custom fp16 plugin automatically * Fix import issue, attempting to fix tests * Fix initial test * Reflect hook logic from master, should wrap model after move to device * Optional state consolidation, since master has optimizers not wrapped * change attribute for instance test * reset optimizers optimizers are not used in main process, so state would be wrong. * legacy * imports in accel * legacy2 * trainer imports * fix import errors after rebase * move hook to new setup location * provide unwrapping logic * fix trainer callback system * added ddp2 implementation * fix imports .legacy * move plugins * restore legacy * drop test.py from root * add tpu accelerator and plugins * fixes * fix lightning optimizer merge * reset bugreportmodel * unwrapping * step routing forward * model access * unwrap * opt * integrate distrib_type * sync changes * sync * fixes * add forgotten generators * add missing logic * update * import * missed imports * import fixes * isort * mv f * changelog * format * move helper to parallel plugin * d * add world size * clean up * duplicate * activate ddp_sharded and tpu * set nvidia flags * remove unused colab var * use_tpu <-> on_tpu attrs * make some ddp_cpu and clusterplugin tests pass * Ref/accelerator connector (#5742) * final cleanup Co-authored-by: Adrian Wälchli <aedu.waelchli@gmail.com> * connector cleanup Co-authored-by: Adrian Wälchli <aedu.waelchli@gmail.com> * trainer cleanup Co-authored-by: Adrian Wälchli <aedu.waelchli@gmail.com> * accelerator cleanup + missing logic in accelerator connector Co-authored-by: Adrian Wälchli <aedu.waelchli@gmail.com> * add missing changes to callbacks Co-authored-by: Adrian Wälchli <aedu.waelchli@gmail.com> * reflect accelerator changes to lightning module Co-authored-by: Adrian Wälchli <aedu.waelchli@gmail.com> * clean cluster envs Co-authored-by: Adrian Wälchli <aedu.waelchli@gmail.com> * cleanup plugins Co-authored-by: Adrian Wälchli <aedu.waelchli@gmail.com> * add broadcasting Co-authored-by: Adrian Wälchli <aedu.waelchli@gmail.com> * yapf * remove plugin connector Co-authored-by: Adrian Wälchli <aedu.waelchli@gmail.com> * plugins * manual optimization * update optimizer routing * add rank to torchelastic * fix memory mixed precision * setstate on trainer for pickling in ddp spawn * add predict method * add back commented accelerator code * adapt test for sync_batch_norm to new plugin * fix deprecated tests * fix ddp cpu choice when no num_processes are given * yapf format * skip a memory test that cannot pass anymore * fix pickle error in spawn plugin * x * avoid * x * fix cyclic import in docs build * add support for sharded * update typing * add sharded and sharded_spawn to distributed types * make unwrap model default * refactor LightningShardedDataParallel similar to LightningDistributedDataParallel * update sharded spawn to reflect changes * update sharded to reflect changes * Merge 1.1.5 changes * fix merge * fix merge * yapf isort * fix merge * yapf isort * fix indentation in test * copy over reinit scheduler implementation from dev1.2 * fix apex tracking calls with dev_debugger * reduce diff to dev1.2, clean up * fix trainer config test when gpus>0 and num_processes >0 and ddp_cpu * sort plugin tests legacy/new * fix error handling for amp on cpu * fix merge fix merge fix merge * [Feat] Resolve manual_backward (#5837) * resolve manual_backward * resolve flake8 * update * resolve for ddp_spawn * resolve flake8 * resolve flake8 * resolve flake8 Co-authored-by: Ubuntu <ubuntu@ip-172-31-88-60.ec2.internal> * fix tests/accelerator tests on cpu * [BugFix] Resolve manual optimization (#5852) * resolve manual_optimization * update * update Co-authored-by: Ubuntu <ubuntu@ip-172-31-88-60.ec2.internal> * Remove copy trainer parameters to happen earlier within the loop and add safe guard to get ref model (#5856) * resovle a bug * Accelerator refactor sharded rpc (#5854) * rpc branch * merge * update handling of rpc * make devices etc. Optional in RPC * set devices etc. later if necessary * remove devices from sequential * make devices optional in rpc * fix import * uncomment everything * fix cluster selection Co-authored-by: Ubuntu <ubuntu@ip-172-31-88-60.ec2.internal> * resolve bug * fix assert in rpc test * resolve a test * fix docs compilation * accelerator refactor - fix for sharded parity test (#5866) * fix memory issue with ddp_spawn * x x x x x x x x x * x * Remove DDP2 as this does not apply * Add missing pre optimizer hook to ensure lambda closure is called * fix apex docstring * [accelerator][BugFix] Resolve some test for 1 gpu (#5863) * update * revert init * resolve a bug * update * resolve flake8 * update * update * update * revert init * resolve a bug * update * resolve flake8 * update * update * update * update * update * revert init * resolve a bug * update * resolve flake8 * update * update * update * revert init * update * resolve flake8 * update * update * update * update * update * all_gather * update * make plugins work, add misconfig for RPC * update * update * remove breaking test * resolve some tests * resolve flake8 * revert to ddp_spawn Co-authored-by: root <root@ip-172-31-88-60.ec2.internal> Co-authored-by: Ubuntu <ubuntu@ip-172-31-88-60.ec2.internal> Co-authored-by: Justus Schock <justus.schock@rwth-aachen.de> * yapf isort * resolve flake8 * fix apex doctests * fix apex doctests 2 * resolve docs * update drone * clean env * update * update * update * update * merge * Fix RPC related tests, clean out old API, update for new accelerator API [skip ci] (#5881) * Fix RPC related tests, clean out old API, update for new accelerator API * Move tests out of legacy folder, update paths and names * Update test_remove_1-4.py * Expose properties for tpu cores/gpus/num_gpus * Add root GPU property * Move properties to properties.py * move tests that were previously in drone * Fix root GPU property (#5908) * Move root GPU to property, remove horovod set as this is handled in horovod plugin, ensure we mock correctly to set GPU accelerator * Add missing tests back * fix best model path transfer when no checkpoint callback available * Fix setup hook order [wip] (#5858) * Call trainer setup hook before accelerator setup * Add test case * add new test * typo * fix callback order in test Co-authored-by: tchaton <thomas@grid.ai> Co-authored-by: Adrian Wälchli <aedu.waelchli@gmail.com> * rename ddp sequential -> rpc sequential for special test * revert * fix stupid merge problem * Use property in connector for sampler (#5913) * merge the import conflicts * fix spawning of processes in slurm * [wip] Fix some bugs for TPU [skip ci] (#5878) * fixed for single tpu * fixed spawn * fixed spawn * update * update * wip * resolve bugs * resolve bug * update on comment * removed decorator * resolve comments * set to 4 * update * update * need cleaning * update * update * update * resolve flake8 * resolve bugs * exclude broadcast * resolve bugs * change test * update * update * skip if meet fails * properly raise trace * update * add catch * wrap test * resolve typo * update * typo Co-authored-by: Lezwon Castelino <lezwon@gmail.com> Co-authored-by: Your Name <you@example.com> * resolve some tests * update * fix imports * update * resolve flake8 * update azure pipeline * skip a sharded test on cpu that requires a gpu * resolve tpus * resolve bug * resolve flake8 * update * updat utils * revert permission change on files * suggestions from carlos Co-authored-by: Carlos Mocholí <carlossmocholi@gmail.com> * remove unrelated formatting changes * remove incomplete comment * Update pytorch_lightning/accelerators/__init__.py Co-authored-by: Carlos Mocholí <carlossmocholi@gmail.com> * remove unrelated formatting change * add types * warn 1.7 ddp manual backward only if ddp kwarg unset * yapf + isort * pep8 unused imports * fix cyclic import in docs * Apply suggestions from code review * typer in accelerator.py * typo * Apply suggestions from code review * formatting * update on comments * update typo * Update pytorch_lightning/trainer/properties.py Co-authored-by: Adrian Wälchli <aedu.waelchli@gmail.com> * update * suggestion from code review * suggestion from code review Co-authored-by: Adrian Wälchli <aedu.waelchli@gmail.com> Co-authored-by: SeanNaren <sean@grid.ai> Co-authored-by: Jirka Borovec <jirka.borovec@seznam.cz> Co-authored-by: chaton <thomas@grid.ai> Co-authored-by: Ubuntu <ubuntu@ip-172-31-88-60.ec2.internal> Co-authored-by: Sean Naren <sean.narenthiran@gmail.com> Co-authored-by: root <root@ip-172-31-88-60.ec2.internal> Co-authored-by: Lezwon Castelino <lezwon@gmail.com> Co-authored-by: Your Name <you@example.com> Co-authored-by: Carlos Mocholí <carlossmocholi@gmail.com> Co-authored-by: Jirka Borovec <Borda@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Lightning-AI · Feb 12, 2021 · da6dbc8 · da6dbc8
1 parent 309ce7a
commit da6dbc8
Show file tree

Hide file tree

Showing 92 changed files with 1,685 additions and 1,678 deletions.
diff --git a/.gitignore b/.gitignore
@@ -151,6 +151,6 @@ wandb
 
 # dataset generated from bolts in examples.
 cifar-10-batches-py
-
+*.pt
 # ctags
-tags
+tags
diff --git a/benchmarks/test_sharded_parity.py b/benchmarks/test_sharded_parity.py
@@ -15,14 +15,13 @@
 import os
 import platform
 import time
-from typing import Type, Union
+from typing import Type
 
 import pytest
 import torch
 
 from pytorch_lightning import seed_everything, Trainer
-from pytorch_lightning.plugins.legacy.ddp_plugin import DDPPlugin
-from pytorch_lightning.plugins.legacy.sharded_plugin import DDPShardedPlugin
+from pytorch_lightning.plugins import DDPSpawnShardedPlugin
 from pytorch_lightning.utilities import _FAIRSCALE_AVAILABLE, _NATIVE_AMP_AVAILABLE
 from tests.accelerators.legacy import DDPLauncher
 from tests.helpers.boring_model import BoringModel, RandomDataset
@@ -34,8 +33,6 @@
 def test_ddp_sharded_plugin_correctness_one_gpu():
  plugin_parity_test(
  gpus=1,
- accelerator='ddp_spawn',
- plugin=DDPShardedPlugin(),
  model_cls=SeedTrainLoaderModel,
  )
 
@@ -48,8 +45,6 @@ def test_ddp_sharded_plugin_correctness_amp_one_gpu():
  plugin_parity_test(
  gpus=1,
  precision=16,
- accelerator='ddp_spawn',
- plugin=DDPShardedPlugin(),
  model_cls=SeedTrainLoaderModel,
  )
 
@@ -61,8 +56,6 @@ def test_ddp_sharded_plugin_correctness_amp_one_gpu():
 def test_ddp_sharded_plugin_correctness_multi_gpu():
  plugin_parity_test(
  gpus=2,
- accelerator='ddp_spawn',
- plugin=DDPShardedPlugin(),
  model_cls=SeedTrainLoaderModel,
  max_percent_speed_diff=0.25, # todo: Increase speed diff since only 2 GPUs sharding 2 optimizers
  )
@@ -76,8 +69,6 @@ def test_ddp_sharded_plugin_correctness_amp_multi_gpu():
  plugin_parity_test(
  gpus=2,
  precision=16,
- accelerator='ddp_spawn',
- plugin=DDPShardedPlugin(),
  model_cls=SeedTrainLoaderModel,
  max_percent_speed_diff=0.25, # todo: Increase speed diff since only 2 GPUs sharding 2 optimizers
  )
@@ -91,8 +82,6 @@ def test_ddp_string_sharded_plugin_correctness_amp_multi_gpu():
  plugin_parity_test(
  gpus=2,
  precision=16,
- accelerator='ddp_spawn',
- plugin='ddp_sharded',
  model_cls=SeedTrainLoaderModel,
  max_percent_speed_diff=0.25, # todo: Increase speed diff since only 2 GPUs sharding 2 optimizers
  )
@@ -108,8 +97,6 @@ def test_ddp_sharded_plugin_correctness_multi_gpu_ddp(tmpdir, args=None):
  plugin_parity_test(
  gpus=args.gpus,
  precision=args.precision,
- accelerator=args.accelerator,
- plugin=DDPShardedPlugin(),
  model_cls=SeedTrainLoaderModel,
  )
 
@@ -124,8 +111,6 @@ def test_ddp_sharded_plugin_correctness_amp_multi_gpu_ddp(tmpdir, args=None):
  plugin_parity_test(
  gpus=args.gpus,
  precision=args.precision,
- accelerator=args.accelerator,
- plugin=DDPShardedPlugin(),
  model_cls=SeedTrainLoaderModel,
  )
 
@@ -139,9 +124,7 @@ def test_ddp_sharded_plugin_correctness_multi_gpu_multi_optim():
  Ensures same results using multiple optimizers across multiple GPUs
  """
  plugin_parity_test(
- plugin=DDPShardedPlugin(),
  gpus=2,
- accelerator='ddp_spawn',
  model_cls=SeedTrainLoaderMultipleOptimizersModel,
  max_percent_speed_diff=0.25, # todo: Increase speed diff since only 2 GPUs sharding 2 optimizers
  )
@@ -156,9 +139,7 @@ def test_ddp_sharded_plugin_correctness_multi_gpu_multi_optim_manual(tmpdir):
  Ensures using multiple optimizers across multiple GPUs with manual optimization
  """
  plugin_parity_test(
- plugin=DDPShardedPlugin(),
  gpus=2,
- accelerator='ddp_spawn',
  model_cls=SeedTrainLoaderManualModel,
  max_percent_speed_diff=0.25, # todo: Increase speed diff since only 2 GPUs sharding 2 optimizers
  )
@@ -260,9 +241,7 @@ def record_ddp_fit_model_stats(trainer, model, use_cuda):
 
 def plugin_parity_test(
  model_cls: Type[SeedTrainLoaderModel],
- plugin: Union[str, DDPPlugin],
  seed: int = 42,
- accelerator: str = 'ddp_spawn',
  gpus: int = 0,
  precision: int = 32,
  max_percent_speed_diff: float = 0.1,
@@ -273,9 +252,7 @@ def plugin_parity_test(
 
  Args:
  model_cls: Model class to use for test.
- plugin: Plugin to parity test.
  seed: Seed for generators. Note that this does not handle the seed for data-loading on multi-process.
- accelerator: Accelerator type for test.
  gpus: Number of GPUS to enable.
  precision: Whether to use AMP or normal FP32 training.
  max_percent_speed_diff: The maximum speed difference compared to normal DDP training.
@@ -293,7 +270,7 @@ def plugin_parity_test(
  max_epochs=1,
  gpus=gpus,
  precision=precision,
- accelerator=accelerator,
+ accelerator='ddp_spawn',
  )
 
  max_memory_ddp, ddp_time = record_ddp_fit_model_stats(trainer=trainer, model=ddp_model, use_cuda=use_cuda)
@@ -307,9 +284,9 @@ def plugin_parity_test(
  max_epochs=1,
  gpus=gpus,
  precision=precision,
- accelerator=accelerator,
- plugins=[plugin],
+ accelerator='ddp_sharded_spawn',
  )
+ assert isinstance(trainer.training_type_plugin, DDPSpawnShardedPlugin)
 
  max_memory_custom, custom_model_time = record_ddp_fit_model_stats(
  trainer=trainer, model=custom_plugin_model, use_cuda=use_cuda

diff --git a/dockers/tpu-tests/tpu_test_cases.jsonnet b/dockers/tpu-tests/tpu_test_cases.jsonnet
@@ -21,7 +21,7 @@ local tputests = base.BaseTest {
  command: utils.scriptCommand(
  |||
  cd pytorch-lightning
- coverage run --source=pytorch_lightning -m pytest -v \
+ coverage run --source=pytorch_lightning -m pytest -v --capture=no \
  pytorch_lightning/utilities/xla_device_utils.py \
  tests/accelerators/legacy/test_tpu_backend.py \
  tests/models/test_tpu.py

diff --git a/docs/source/advanced/amp.rst b/docs/source/advanced/amp.rst
@@ -31,10 +31,10 @@ Native torch
 When using PyTorch 1.6+ Lightning uses the native amp implementation to support 16-bit.
 
 .. testcode::
- :skipif: not _APEX_AVAILABLE and not _NATIVE_AMP_AVAILABLE
+ :skipif: not _APEX_AVAILABLE and not _NATIVE_AMP_AVAILABLE or not torch.cuda.is_available()
 
  # turn on 16-bit
- trainer = Trainer(precision=16)
+ trainer = Trainer(precision=16, gpus=1)
 
 Apex 16-bit
 ^^^^^^^^^^^
@@ -73,7 +73,7 @@ Enable 16-bit
 ^^^^^^^^^^^^^
 
 .. testcode::
- :skipif: not _APEX_AVAILABLE and not _NATIVE_AMP_AVAILABLE
+ :skipif: not _APEX_AVAILABLE and not _NATIVE_AMP_AVAILABLE or not torch.cuda.is_available()
 
  # turn on 16-bit
  trainer = Trainer(amp_level='O2', precision=16)

diff --git a/docs/source/common/trainer.rst b/docs/source/common/trainer.rst
@@ -1178,13 +1178,13 @@ If used on TPU will use torch.bfloat16 but tensor printing
 will still show torch.float32.
 
 .. testcode::
- :skipif: not _APEX_AVAILABLE and not _NATIVE_AMP_AVAILABLE
+ :skipif: not _APEX_AVAILABLE and not _NATIVE_AMP_AVAILABLE or not torch.cuda.is_available()
 
  # default used by the Trainer
  trainer = Trainer(precision=32)
 
  # 16-bit precision
- trainer = Trainer(precision=16)
+ trainer = Trainer(precision=16, gpus=1)
 
 Example::
 

diff --git a/pytorch_lightning/accelerators/__init__.py b/pytorch_lightning/accelerators/__init__.py
@@ -1,5 +1,4 @@
 # 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
@@ -11,15 +10,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 pytorch_lightning.accelerators.legacy.accelerator import Accelerator # noqa: F401
-from pytorch_lightning.accelerators.legacy.cpu_accelerator import CPUAccelerator # noqa: F401
-from pytorch_lightning.accelerators.legacy.ddp2_accelerator import DDP2Accelerator # noqa: F401
-from pytorch_lightning.accelerators.legacy.ddp_accelerator import DDPAccelerator # noqa: F401
-from pytorch_lightning.accelerators.legacy.ddp_cpu_hpc_accelerator import DDPCPUHPCAccelerator # noqa: F401
-from pytorch_lightning.accelerators.legacy.ddp_cpu_spawn_accelerator import DDPCPUSpawnAccelerator # noqa: F401
-from pytorch_lightning.accelerators.legacy.ddp_hpc_accelerator import DDPHPCAccelerator # noqa: F401
-from pytorch_lightning.accelerators.legacy.ddp_spawn_accelerator import DDPSpawnAccelerator # noqa: F401
-from pytorch_lightning.accelerators.legacy.dp_accelerator import DataParallelAccelerator # noqa: F401
-from pytorch_lightning.accelerators.legacy.gpu_accelerator import GPUAccelerator # noqa: F401
-from pytorch_lightning.accelerators.legacy.horovod_accelerator import HorovodAccelerator # noqa: F401
-from pytorch_lightning.accelerators.legacy.tpu_accelerator import TPUAccelerator # noqa: F401
+from pytorch_lightning.accelerators.accelerator import Accelerator # noqa F401
+from pytorch_lightning.accelerators.cpu import CPUAccelerator # noqa F401
+from pytorch_lightning.accelerators.gpu import GPUAccelerator # noqa F401
+from pytorch_lightning.accelerators.tpu import TPUAccelerator # noqa F401