From 5eab0deca989ef0936df179800c40518142546cd Mon Sep 17 00:00:00 2001 From: Danielle Pintz Date: Tue, 21 Dec 2021 23:31:25 -0800 Subject: [PATCH 1/6] follow ups to #10957 --- .../loops/optimization/optimizer_loop.py | 22 ++++++++----------- .../training_type/training_type_plugin.py | 4 ---- pytorch_lightning/profiler/pytorch.py | 1 - pytorch_lightning/profiler/xla.py | 3 +-- tests/profiler/test_profiler.py | 1 - 5 files changed, 10 insertions(+), 21 deletions(-) diff --git a/pytorch_lightning/loops/optimization/optimizer_loop.py b/pytorch_lightning/loops/optimization/optimizer_loop.py index ee4af19134cf3..7a8be458aa803 100644 --- a/pytorch_lightning/loops/optimization/optimizer_loop.py +++ b/pytorch_lightning/loops/optimization/optimizer_loop.py @@ -28,7 +28,7 @@ _extract_hiddens, check_finite_loss, ) -from pytorch_lightning.profiler import BaseProfiler, PassThroughProfiler +from pytorch_lightning.profiler import BaseProfiler from pytorch_lightning.trainer.progress import OptimizationProgress from pytorch_lightning.utilities import _AcceleratorType, AMPType from pytorch_lightning.utilities.exceptions import MisconfigurationException @@ -132,22 +132,18 @@ def __init__( self._step_fn = step_fn self._backward_fn = backward_fn self._zero_grad_fn = zero_grad_fn - self._profiler = PassThroughProfiler() if profiler is None else profiler def closure(self, *args: Any, **kwargs: Any) -> ClosureResult: - with self._profiler.profile("training_step_and_backward"): - step_output = self._step_fn() + step_output = self._step_fn() - if step_output.closure_loss is None: - self.warning_cache.warn( - "`training_step` returned `None`. If this was on purpose, ignore this warning..." - ) + if step_output.closure_loss is None: + self.warning_cache.warn("`training_step` returned `None`. If this was on purpose, ignore this warning...") - if self._zero_grad_fn is not None: - self._zero_grad_fn() + if self._zero_grad_fn is not None: + self._zero_grad_fn() - if self._backward_fn is not None and step_output.closure_loss is not None: - self._backward_fn(step_output.closure_loss) + if self._backward_fn is not None and step_output.closure_loss is not None: + self._backward_fn(step_output.closure_loss) return step_output @@ -400,7 +396,7 @@ def _optimizer_zero_grad(self, batch_idx: int, optimizer: torch.optim.Optimizer, optimizer: the current optimizer opt_idx: the index of the current optimizer """ - self.trainer._call_strategy_hook( + self.trainer._call_lightning_module_hook( "optimizer_zero_grad", self.trainer.current_epoch, batch_idx, optimizer, opt_idx ) self.optim_progress.optimizer.zero_grad.increment_completed() diff --git a/pytorch_lightning/plugins/training_type/training_type_plugin.py b/pytorch_lightning/plugins/training_type/training_type_plugin.py index 9d7bb51b9a773..5a3737a613e54 100644 --- a/pytorch_lightning/plugins/training_type/training_type_plugin.py +++ b/pytorch_lightning/plugins/training_type/training_type_plugin.py @@ -179,10 +179,6 @@ def optimizer_step( model = model or self.lightning_module self.precision_plugin.optimizer_step(model, optimizer, opt_idx, closure, **kwargs) - def optimizer_zero_grad(self, current_epoch: int, batch_idx: int, optimizer: Optimizer, opt_idx: int) -> None: - """Zeros all model parameter's gradients.""" - self.lightning_module.optimizer_zero_grad(current_epoch, batch_idx, optimizer, opt_idx) - def _setup_model_and_optimizers(self, model: Module, optimizers: List[Optimizer]) -> Tuple[Module, List[Optimizer]]: """Setup a model and multiple optimizers together. diff --git a/pytorch_lightning/profiler/pytorch.py b/pytorch_lightning/profiler/pytorch.py index 042a70966a476..8c542b8876811 100644 --- a/pytorch_lightning/profiler/pytorch.py +++ b/pytorch_lightning/profiler/pytorch.py @@ -195,7 +195,6 @@ def __call__(self, num_step: int) -> "ProfilerAction": class PyTorchProfiler(BaseProfiler): RECORD_FUNCTIONS = { - "training_step_and_backward", "training_step", "backward", "validation_step", diff --git a/pytorch_lightning/profiler/xla.py b/pytorch_lightning/profiler/xla.py index c89685bcad0be..be158f7be495e 100644 --- a/pytorch_lightning/profiler/xla.py +++ b/pytorch_lightning/profiler/xla.py @@ -53,9 +53,8 @@ class XLAProfiler(BaseProfiler): - STEP_FUNCTIONS = {"training_step_and_backward", "validation_step", "test_step", "predict_step"} + STEP_FUNCTIONS = {"validation_step", "test_step", "predict_step"} RECORD_FUNCTIONS = { - "training_step_and_backward", "training_step", "backward", "validation_step", diff --git a/tests/profiler/test_profiler.py b/tests/profiler/test_profiler.py index 9980adb4f41c7..204f622f1dbcc 100644 --- a/tests/profiler/test_profiler.py +++ b/tests/profiler/test_profiler.py @@ -313,7 +313,6 @@ def test_pytorch_profiler_trainer_ddp(tmpdir, pytorch_profiler): expected = {"[Strategy]DDPStrategy.validation_step"} if not _KINETO_AVAILABLE: expected |= { - "training_step_and_backward", "[Strategy]DDPStrategy.training_step", "[Strategy]DDPStrategy.backward", } From dad4ae1b231f7a107fb590bed3627b4866beb015 Mon Sep 17 00:00:00 2001 From: Danielle Pintz Date: Wed, 22 Dec 2021 20:29:16 -0800 Subject: [PATCH 2/6] addr comments --- pytorch_lightning/loops/optimization/optimizer_loop.py | 3 --- pytorch_lightning/profiler/__init__.py | 6 +++--- 2 files changed, 3 insertions(+), 6 deletions(-) diff --git a/pytorch_lightning/loops/optimization/optimizer_loop.py b/pytorch_lightning/loops/optimization/optimizer_loop.py index 7a8be458aa803..6334922f69f84 100644 --- a/pytorch_lightning/loops/optimization/optimizer_loop.py +++ b/pytorch_lightning/loops/optimization/optimizer_loop.py @@ -28,7 +28,6 @@ _extract_hiddens, check_finite_loss, ) -from pytorch_lightning.profiler import BaseProfiler from pytorch_lightning.trainer.progress import OptimizationProgress from pytorch_lightning.utilities import _AcceleratorType, AMPType from pytorch_lightning.utilities.exceptions import MisconfigurationException @@ -110,7 +109,6 @@ class Closure(AbstractClosure[ClosureResult]): Can be set to ``None`` to skip the backward operation. zero_grad_fn: A function that zeroes the gradients. Can be set to ``None`` to skip zero_grad, for example when accumulating gradients. - profiler: A profiler for profiling the actions of the passed in closure functions. Example: @@ -126,7 +124,6 @@ def __init__( step_fn: Callable[[], ClosureResult], backward_fn: Optional[Callable[[Tensor], None]] = None, zero_grad_fn: Optional[Callable[[], None]] = None, - profiler: Optional[BaseProfiler] = None, ): super().__init__() self._step_fn = step_fn diff --git a/pytorch_lightning/profiler/__init__.py b/pytorch_lightning/profiler/__init__.py index 58cee0c1d8af2..e67295c12b09e 100644 --- a/pytorch_lightning/profiler/__init__.py +++ b/pytorch_lightning/profiler/__init__.py @@ -141,9 +141,9 @@ def custom_processing_step(self, data): The profiler's results will be printed on the completion of ``{fit,validate,test,predict}``. -This profiler will record ``training_step_and_backward``, ``training_step``, ``backward``, +This profiler will record ``training_step``, ``backward``, ``validation_step``, ``test_step``, and ``predict_step`` by default. -The output below shows the profiling for the action ``training_step_and_backward``. +The output below shows the profiling for the action ``training_step``. The user can provide ``PyTorchProfiler(record_functions={...})`` to extend the scope of profiled functions. .. note:: @@ -156,7 +156,7 @@ def custom_processing_step(self, data): Profiler Report - Profile stats for: training_step_and_backward + Profile stats for: training_step --------------------- --------------- --------------- --------------- --------------- --------------- Name Self CPU total % Self CPU total CPU total % CPU total CPU time avg --------------------- --------------- --------------- --------------- --------------- --------------- From 1c53230f03c5d36f1fd55564a7988d9f3539d194 Mon Sep 17 00:00:00 2001 From: Carlos Mocholi Date: Thu, 23 Dec 2021 05:40:15 +0100 Subject: [PATCH 3/6] Fix call --- pytorch_lightning/loops/optimization/optimizer_loop.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/pytorch_lightning/loops/optimization/optimizer_loop.py b/pytorch_lightning/loops/optimization/optimizer_loop.py index 6334922f69f84..5a5625e28515d 100644 --- a/pytorch_lightning/loops/optimization/optimizer_loop.py +++ b/pytorch_lightning/loops/optimization/optimizer_loop.py @@ -273,10 +273,7 @@ def _make_closure(self, split_batch: Any, batch_idx: int, opt_idx: int, optimize step_fn = self._make_step_fn(split_batch, batch_idx, opt_idx) backward_fn = self._make_backward_fn(optimizer, opt_idx) zero_grad_fn = self._make_zero_grad_fn(batch_idx, opt_idx, optimizer) - - return Closure( - step_fn=step_fn, backward_fn=backward_fn, zero_grad_fn=zero_grad_fn, profiler=self.trainer.profiler - ) + return Closure(step_fn=step_fn, backward_fn=backward_fn, zero_grad_fn=zero_grad_fn) def _make_step_fn(self, split_batch: Any, batch_idx: int, opt_idx: int) -> Callable[[], ClosureResult]: """Build the step function that runs the `training_step` and processes its output.""" From 6b3133a8a54cc001fc2ee45f3971de70d6537b4d Mon Sep 17 00:00:00 2001 From: Danielle Pintz Date: Thu, 23 Dec 2021 11:01:19 -0800 Subject: [PATCH 4/6] revert strategy change --- pytorch_lightning/loops/optimization/optimizer_loop.py | 2 +- pytorch_lightning/strategies/strategy.py | 4 ++++ 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/pytorch_lightning/loops/optimization/optimizer_loop.py b/pytorch_lightning/loops/optimization/optimizer_loop.py index 5a5625e28515d..35d6b98e3ce11 100644 --- a/pytorch_lightning/loops/optimization/optimizer_loop.py +++ b/pytorch_lightning/loops/optimization/optimizer_loop.py @@ -390,7 +390,7 @@ def _optimizer_zero_grad(self, batch_idx: int, optimizer: torch.optim.Optimizer, optimizer: the current optimizer opt_idx: the index of the current optimizer """ - self.trainer._call_lightning_module_hook( + self.trainer._call_strategy_hook( "optimizer_zero_grad", self.trainer.current_epoch, batch_idx, optimizer, opt_idx ) self.optim_progress.optimizer.zero_grad.increment_completed() diff --git a/pytorch_lightning/strategies/strategy.py b/pytorch_lightning/strategies/strategy.py index 3855a7972d8d0..fbf6642cc18ba 100644 --- a/pytorch_lightning/strategies/strategy.py +++ b/pytorch_lightning/strategies/strategy.py @@ -177,6 +177,10 @@ def optimizer_step( model = model or self.lightning_module self.precision_plugin.optimizer_step(model, optimizer, opt_idx, closure, **kwargs) + def optimizer_zero_grad(self, current_epoch: int, batch_idx: int, optimizer: Optimizer, opt_idx: int) -> None: + """Zeros all model parameter's gradients.""" + self.lightning_module.optimizer_zero_grad(current_epoch, batch_idx, optimizer, opt_idx) + def _setup_model_and_optimizers(self, model: Module, optimizers: List[Optimizer]) -> Tuple[Module, List[Optimizer]]: """Setup a model and multiple optimizers together. From ffd7a4dc6de49f0e12a3dca9fbf2b67206cf9304 Mon Sep 17 00:00:00 2001 From: Danielle Pintz Date: Thu, 23 Dec 2021 11:07:51 -0800 Subject: [PATCH 5/6] fix space --- pytorch_lightning/strategies/strategy.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pytorch_lightning/strategies/strategy.py b/pytorch_lightning/strategies/strategy.py index fbf6642cc18ba..fe9093838c157 100644 --- a/pytorch_lightning/strategies/strategy.py +++ b/pytorch_lightning/strategies/strategy.py @@ -177,7 +177,7 @@ def optimizer_step( model = model or self.lightning_module self.precision_plugin.optimizer_step(model, optimizer, opt_idx, closure, **kwargs) - def optimizer_zero_grad(self, current_epoch: int, batch_idx: int, optimizer: Optimizer, opt_idx: int) -> None: + def optimizer_zero_grad(self, current_epoch: int, batch_idx: int, optimizer: Optimizer, opt_idx: int) -> None: """Zeros all model parameter's gradients.""" self.lightning_module.optimizer_zero_grad(current_epoch, batch_idx, optimizer, opt_idx) From e1cdf7a163884c8ee61f3bf0f68b6c81f78a6401 Mon Sep 17 00:00:00 2001 From: Danielle Pintz Date: Thu, 23 Dec 2021 11:15:10 -0800 Subject: [PATCH 6/6] changelog --- CHANGELOG.md | 3 +++ 1 file changed, 3 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index a6c7bbe4c6b40..8c733a522eb71 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -346,6 +346,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/). - Removed `Strategy.init_optimizers` in favor of `Strategy.setup_optimizers` ([#11236](https://github.com/PyTorchLightning/pytorch-lightning/pull/11236)) + +- Removed `profile("training_step_and_backward")` ([#11222](https://github.com/PyTorchLightning/pytorch-lightning/pull/11222)) + ### Fixed - Fixed security vulnerabilities CVE-2020-1747 and CVE-2020-14343 caused by the `PyYAML` dependency ([#11099](https://github.com/PyTorchLightning/pytorch-lightning/pull/11099))