[accelerator][FeatBugFix] Improve manual optimization API (#5771)

* 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 * update on comments * 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 * update on comments * resolve some comments * update on comments * resolve test * add toggle_model * update * update on comments * update doc * typo * update * typo * remove space * update * update on comments Co-authored-by: Adrian Wälchli <aedu.waelchli@gmail.com> Co-authored-by: justusschock <justus.schock@posteo.de> Co-authored-by: SeanNaren <sean@grid.ai> Co-authored-by: Justus Schock <12886177+justusschock@users.noreply.github.com> Co-authored-by: Jirka Borovec <jirka.borovec@seznam.cz> Co-authored-by: Justus Schock <justus.schock@rwth-aachen.de> 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 16, 2021 · 6e79bef · 6e79bef
1 parent 6d1e055
commit 6e79bef
Show file tree

Hide file tree

Showing 14 changed files with 340 additions and 367 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -190,6 +190,10 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Refactored `EpochResultStore` ([#5522](https://github.com/PyTorchLightning/pytorch-lightning/pull/5522))
 
 
+- LightningOptimizer manual optimizer is more flexible and expose `toggle_model` ([#5771](https://github.com/PyTorchLightning/pytorch-lightning/pull/5771))
+
+
+
 ### Deprecated
 
 - Function `stat_scores_multiple_classes` is deprecated in favor of `stat_scores` ([#4839](https://github.com/PyTorchLightning/pytorch-lightning/pull/4839))

diff --git a/docs/source/common/optimizers.rst b/docs/source/common/optimizers.rst
@@ -21,46 +21,117 @@ Manual optimization
 For advanced research topics like reinforcement learning, sparse coding, or GAN research, it may be desirable
 to manually manage the optimization process. To do so, do the following:
 
-* Disable automatic optimization in Trainer: Trainer(automatic_optimization=False)
+* Override your LightningModule ``automatic_optimization`` property to return ``False``
 * Drop or ignore the optimizer_idx argument
-* Use `self.manual_backward(loss)` instead of `loss.backward()` to automatically scale your loss
+* Use `self.manual_backward(loss)` instead of `loss.backward()`.
+
+.. note:: This is only recommended for experts who need ultimate flexibility. Lightning will handle only precision and accelerators logic. The users are left with zero_grad, accumulated_grad_batches, model toggling, etc..
+
+.. warning:: Before 1.2, ``optimzer.step`` was calling ``zero_grad`` internally. From 1.2, it is left to the users expertize.
+
+.. tip:: To perform ``accumulate_grad_batches`` with one optimizer, you can do as such.
+
+.. tip:: ``self.optimizers()`` will return ``LightningOptimizer`` objects. You can access your own optimizer with ``optimizer.optimizer``. However, if you use your own optimizer to perform a step, Lightning won't be able to support accelerators and precision for you.
+
 
 .. code-block:: python
 
- def training_step(self, batch, batch_idx, optimizer_idx):
+ def training_step(batch, batch_idx, optimizer_idx):
+ opt = self.optimizers()
+
+ loss = self.compute_loss(batch)
+ self.manual_backward(loss)
+ opt.step()
 
- # 1. ignore optimizer_idx
- # 2. `use_pl_optimizer=True` means `opt_g` and `opt_d` will be of type `LightningOptimizer`
- # `LightningOptimizer` simply wrapped your optimizer and behave the same way !
- # When calling `optimizer.step`, `LightningOptimizer` will just handle TPU, AMP, accumulate_grad_batches, etc ... for you.
+ # accumulate gradient batches
+ if batch_idx % 2 == 0:
+ opt.zero_grad()
 
- # access your optimizers with `use_pl_optimizer=False` or `optimizer.optimizer` when using use_pl_optimizer=True
- # use_pl_optimizer=True is the default
- (opt_g, opt_d) = self.optimizers(use_pl_optimizer=True)
 
- # do anything you want
- loss_a = ...
+.. tip:: It is a good practice to provide the optimizer with a ``closure`` function that performs a ``forward`` and ``backward`` pass of your model. It is optional for most optimizers, but makes your code compatible if you switch to an optimizer which requires a closure.
 
- # use self.backward which will also handle scaling the loss when using amp
- self.manual_backward(loss_a, opt_g)
- opt_g.step()
+Here is the same example as above using a ``closure``.
 
+.. code-block:: python
+
+ def training_step(batch, batch_idx, optimizer_idx):
+ opt = self.optimizers()
+
+ def forward_and_backward():
+ loss = self.compute_loss(batch)
+ self.manual_backward(loss)
 
- # do anything you want
- loss_b = ...
+ opt.step(closure=forward_and_backward)
 
- # pass in any args that loss.backward() normally takes
- self.manual_backward(loss_b, opt_d, retain_graph=True)
- self.manual_backward(loss_b, opt_d)
- opt_d.step()
+ # accumulate gradient batches
+ if batch_idx % 2 == 0:
+ opt.zero_grad()
+
+
+.. code-block:: python
 
- # log losses
- self.log('loss_a', loss_a)
- self.log('loss_b', loss_b)
+ # Scenario for a GAN.
 
-.. note:: This is only recommended for experts who need ultimate flexibility
+ def training_step(...):
+ opt_gen, opt_dis = self.optimizers()
 
-Manual optimization does not yet support accumulated gradients but will be live in 1.1.0
+ # compute generator loss
+ loss_gen = self.compute_generator_loss(...)
+
+ # zero_grad needs to be called before backward
+ opt_gen.zero_grad()
+ self.manual_backward(loss_gen)
+ opt_gen.step()
+
+ # compute discriminator loss
+ loss_dis = self.compute_discriminator_loss(...)
+
+ # zero_grad needs to be called before backward
+ opt_dis.zero_grad()
+ self.manual_backward(loss_dis)
+ opt_dis.step()
+
+
+.. note:: ``LightningOptimizer`` provides a ``toggle_model`` function as a ``@context_manager`` for advanced users. It can be useful when performing gradient accumulation with several optimizers or training in a distributed setting.
+
+Here is an explanation of what it does:
+
+Considering the current optimizer as A and all other optimizers as B.
+Toggling means that all parameters from B exclusive to A will have their ``requires_grad`` attribute set to ``False``. Their original state will be restored when exiting the context manager.
+
+When performing gradient accumulation, there is no need to perform grad synchronization during the accumulation phase.
+Setting ``sync_grad`` to ``False`` will block this synchronization and improve your training speed.
+
+Here is an example on how to use it:
+
+.. code-block:: python
+
+
+ # Scenario for a GAN with gradient accumulation every 2 batches and optimized for multiple gpus.
+
+ def training_step(self, batch, batch_idx, ...):
+ opt_gen, opt_dis = self.optimizers()
+
+ accumulated_grad_batches = batch_idx % 2 == 0
+
+ # compute generator loss
+ def closure_gen():
+ loss_gen = self.compute_generator_loss(...)
+ self.manual_backward(loss_gen)
+ if accumulated_grad_batches:
+ opt_gen.zero_grad()
+
+ with opt_gen.toggle_model(sync_grad=accumulated_grad_batches):
+ opt_gen.step(closure=closure_gen)
+
+ def closure_dis():
+ loss_dis = self.compute_discriminator_loss(...)
+ self.manual_backward(loss_dis)
+ if accumulated_grad_batches:
+ opt_dis.zero_grad()
+
+ with opt_dis.toggle_model(sync_grad=accumulated_grad_batches):
+ opt_dis.step(closure=closure_dis)
 
 ------
 
@@ -166,7 +237,7 @@ returned as a dict which can contain the following keywords:
 * ``strict`` (optional): if set to ``True`` will enforce that value specified in ``monitor`` is available while trying
  to call ``scheduler.step()``, and stop training if not found. If ``False`` will only give a warning and continue training
  (without calling the scheduler).
-* ``name`` (optional): if using the :class:`~pytorch_lightning.callbacks.LearningRateMonitor` callback to monitor the 
+* ``name`` (optional): if using the :class:`~pytorch_lightning.callbacks.LearningRateMonitor` callback to monitor the
  learning rate progress, this keyword can be used to specify a specific name the learning rate should be logged as.
 
 .. testcode::
@@ -248,23 +319,6 @@ For example, here step optimizer A every 2 batches and optimizer B every 4 batch
  if batch_nb % 4 == 0 :
  optimizer.step(closure=closure)
 
-.. note:: When using ``Trainer(enable_pl_optimizer=True)``, ``.step`` accepts a boolean ``make_optimizer_step`` which can be used as follow.
-
-.. testcode::
-
- def optimizer_zero_grad(self, current_epoch, batch_idx, optimizer, opt_idx):
- optimizer.zero_grad()
-
- # Alternating schedule for optimizer steps (ie: GANs)
- def optimizer_step(self, current_epoch, batch_nb, optimizer, optimizer_idx, closure, on_tpu=False, using_native_amp=False, using_lbfgs=False):
- # update generator opt every 2 steps
- if optimizer_idx == 0:
- optimizer.step(closure=closure, make_optimizer_step=(batch_nb % 2) == 0)
-
- # update discriminator opt every 4 steps
- if optimizer_idx == 1:
- optimizer.step(closure=closure, make_optimizer_step=(batch_nb % 4) == 0)
-
 Here we add a learning-rate warm up
 
 .. testcode::

diff --git a/pytorch_lightning/accelerators/accelerator.py b/pytorch_lightning/accelerators/accelerator.py
@@ -238,7 +238,7 @@ def backward(
  self,
  closure_loss: torch.Tensor,
  optimizer: Optimizer,
- opt_idx: int,
+ optimizer_idx: int,
  should_accumulate: bool,
  *args,
  **kwargs,
@@ -247,17 +247,15 @@ def backward(
 
  Args:
  closure_loss: a tensor holding the loss value to backpropagate
- optimizer: the optimizer to do the step later on.
- opt_idx: the index of the optimizer
  should_accumulate: whether to accumulate gradients
  """
- self.training_type_plugin.pre_backward(closure_loss, should_accumulate, optimizer, opt_idx)
+ self.training_type_plugin.pre_backward(closure_loss, should_accumulate, optimizer, optimizer_idx)
 
  output = self.precision_plugin.backward(
- self.lightning_module, closure_loss, optimizer, opt_idx, should_accumulate, *args, **kwargs
+ self.lightning_module, closure_loss, optimizer, optimizer_idx, should_accumulate, *args, **kwargs
  )
 
- self.training_type_plugin.post_backward(closure_loss, should_accumulate, optimizer, opt_idx)
+ self.training_type_plugin.post_backward(closure_loss, should_accumulate, optimizer, optimizer_idx)
 
  return output
 

diff --git a/pytorch_lightning/core/lightning.py b/pytorch_lightning/core/lightning.py
@@ -1186,7 +1186,7 @@ def configure_optimizers(self):
  """
  rank_zero_warn("`configure_optimizers` must be implemented to be used with the Lightning Trainer")
 
- def manual_backward(self, loss: Tensor, optimizer: Optimizer, *args, **kwargs) -> None:
+ def manual_backward(self, loss: Tensor, optimizer: Optional[Optimizer] = None, *args, **kwargs) -> None:
  """
  Call this directly from your training_step when doing optimizations manually.
  By using this we can ensure that all the proper scaling when using 16-bit etc has been done for you
@@ -1207,12 +1207,18 @@ def training_step(...):
  self.manual_backward(loss, opt_a)
  opt_a.step()
  """
+ if optimizer is not None:
+ rank_zero_warn(
+ "`optimizer` argument to `manual_backward` is deprecated in v1.2 and will be removed in v1.4",
+ DeprecationWarning
+ )
+
  # make sure we're using manual opt
  self._verify_is_manual_optimization('manual_backward')
 
  # backward
  self._running_manual_backward = True
- self.trainer.train_loop.backward(loss, optimizer, -1, *args, **kwargs)
+ self.trainer.train_loop.backward(loss, optimizer=None, opt_idx=None, *args, **kwargs)
  self._running_manual_backward = False
 
  def backward(self, loss: Tensor, optimizer: Optimizer, optimizer_idx: int, *args, **kwargs) -> None: