Repository: Lightning-AI/pytorch-lightning
Pull request: [bugfix] Prevent deepcopy of dataloaders / Trainer in SWA Callback
## What does this PR do?

<!--
Please include a summary of the change and which issue is fixed.
 Please also include relevant motivation and context.
 List any dependencies that are required for this change.

If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

The following links the related issue to the PR (https://docs.github.com/en/free-pro-team@latest/github/managing-your-work-on-github/linking-a-pull-request-to-an-issue#linking-a-pull-request-to-an-issue-using-a-keyword)
-->

This PR introduces a context manager to prevent deecopying trainer / dataloaders from the LightningModule.
This PR uses this context manager to prevent with stochastic weight averaging
This PR improve ModelPruning save_checkpoint to avoid a deepcopy.

Fixes https://github.com/PyTorchLightning/pytorch-lightning/issues/7089

### Does your PR introduce any breaking changes ? If yes, please list them.

No.

<!-- In the CHANGELOG, separate each item in the unreleased section by a blank line to reduce collisions -->

## Before submitting
- [ ] Was this discussed/approved via a GitHub issue? (not for typos and docs)
- [ ] Did you read the [contributor guideline](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/.github/CONTRIBUTING.md), **Pull Request** section?
- [ ] Did you make sure your PR does only one thing, instead of bundling different changes together?
- [ ] Did you make sure to update the documentation with your changes? (if necessary)
- [ ] Did you write any new necessary tests? (not for typos and docs)
- [ ] Did you verify new and existing tests pass locally with your changes?
- [x] Did you update the [CHANGELOG](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/CHANGELOG.md)? (not for typos, docs, test updates, or internal minor changes/refactorings)
- [x] Did you list all the breaking changes introduced by this pull request?

<!-- For CHANGELOG separate each item in the unreleased section by a blank line to reduce collisions -->

## PR review
Anyone in the community is free to review the PR once the tests have passed.
Before you start reviewing make sure you have read [Review guidelines](https://github.com/PyTorchLightning/pytorch-lightning/wiki/Review-guidelines). In short, see the following bullet-list:

 - [x] Is this pull request ready for review? (if not, please submit in draft mode)
 - [x] Check that all items from **Before submitting** are resolved
 - [x] Make sure the title is self-explanatory and the description concisely explains the PR
 - [x] Add labels and milestones (and optionally projects) to the PR so it can be classified

## Did you have fun?
Make sure you had fun coding 🙃


Problem statement:
Using SWA callback triggers useless deepcopy of data loaders
## 🐛 Bug

Using SWA callback triggers a deepcopy of the whole PyTorch Lightning module ("pl_module") at the on_before_accelerator_backend_setup hook (lines 140-142 of pytorch_lightning/callbacks/swa.py file). The PyTorch Lightning has the dataloaders as attributes, hence using deepcopy leads to the (useless) duplication of training and validation data loaders. 


### Expected behavior

No copy of the data loaders when initializing the _average_model attribute of the StochasticWeightAveraging callback. 


### Additional context

If data loaders occupy more than half of RAM, this deepcopy will lead to memory overflow.

Hints from issue discussion:
Dear @alexandrebone,

The `dataloaders` are wrapped into functions and it shouldn't cause an increase of memory during the deepcopy.

However, I am not sure about dataset. I will check.

Best,
T.c


@tchaton ping :)
I also ran in OOM with SWA which does not run OOM when SWA is turned off. 
I guess in my case it was the deepcopy of the whole model, though.

Commit message: Simplify test

Patch to review:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 842c41f179073..7a9b2a2b7a73c 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -316,6 +316,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Use XLA utility API to move data to CPU (Single TPU core) ([#8078](https://github.com/PyTorchLightning/pytorch-lightning/pull/8078))
 
 
+- `deepcopy(LightningModule)` no longer copies the trainer or dataloaders/datamodule ([#8472](https://github.com/PyTorchLightning/pytorch-lightning/pull/8472))
+
+
 ### Deprecated
 
 
diff --git a/pytorch_lightning/core/lightning.py b/pytorch_lightning/core/lightning.py
index 90a96819061c0..86a0f9262ed05 100644
--- a/pytorch_lightning/core/lightning.py
+++ b/pytorch_lightning/core/lightning.py
@@ -1998,3 +1998,10 @@ def get_from_queue(self, queue: torch.multiprocessing.SimpleQueue) -> None:
         self.trainer.callback_metrics.update(
             apply_to_collection(callback_metrics, np.ndarray, lambda x: torch.tensor(x))
         )
+
+    def __getstate__(self) -> Dict[str, Any]:
+        # prevent copying data
+        skipped_keys = {
+            "trainer", "train_dataloader", "val_dataloader", "test_dataloader", "predict_dataloader", "datamodule"
+        }
+        return {k: v if k not in skipped_keys else None for k, v in self.__dict__.items()}
diff --git a/tests/callbacks/test_stochastic_weight_avg.py b/tests/callbacks/test_stochastic_weight_avg.py
index 94d7dcf84486a..32b2f5256d372 100644
--- a/tests/callbacks/test_stochastic_weight_avg.py
+++ b/tests/callbacks/test_stochastic_weight_avg.py
@@ -19,7 +19,7 @@
 from torch import nn
 from torch.utils.data import DataLoader
 
-from pytorch_lightning import Trainer
+from pytorch_lightning import LightningModule, Trainer
 from pytorch_lightning.accelerators import Accelerator
 from pytorch_lightning.plugins import DDPSpawnPlugin
 from pytorch_lightning.utilities import _TORCH_GREATER_EQUAL_1_6
@@ -217,3 +217,29 @@ def configure_optimizers(self):
         assert trainer.callbacks[0]._swa_lrs == (1e-3 if use_callbacks else 0.1)
     else:
         assert all(not isinstance(cb, StochasticWeightAveraging) for cb in trainer.callbacks)
+
+
+def test_swa_deepcopy(tmpdir):
+    """Test to ensure SWA Callback doesn't deepcopy dataloaders and datamodule potentially leading to OOM"""
+
+    class TestSWA(StochasticWeightAveraging):
+
+        def __init__(self, *args, **kwargs):
+            super().__init__(*args, **kwargs)
+            self.on_before_accelerator_backend_setup_called = False
+
+        def on_before_accelerator_backend_setup(self, trainer: 'Trainer', pl_module: 'LightningModule'):
+            super().on_before_accelerator_backend_setup(trainer, pl_module)
+            assert self._average_model.train_dataloader is not pl_module.train_dataloader
+            assert self._average_model.trainer is None
+            self.on_before_accelerator_backend_setup_called = True
+
+    model = BoringModel()
+    swa = TestSWA()
+    trainer = Trainer(
+        default_root_dir=tmpdir,
+        callbacks=swa,
+        fast_dev_run=True,
+    )
+    trainer.fit(model)
+    assert swa.on_before_accelerator_backend_setup_called


Code review comments:
diff_hunk: @@ -1998,3 +1998,10 @@ def get_from_queue(self, queue: torch.multiprocessing.SimpleQueue) -> None:
         self.trainer.callback_metrics.update(
             apply_to_collection(callback_metrics, np.ndarray, lambda x: torch.tensor(x))
         )
+
+    def __getstate__(self) -> Dict[str, Any]:
+        # prevent copying data
+        skipped_keys = {
+            "trainer", "train_dataloader", "val_dataloader", "test_dataloader", "predict_dataloader", "datamodule"
path: pytorch_lightning/core/lightning.py
text: @user1:
Why do you skip the dataloader ones which are functions, not properties or attributes?

@user2:
Just guessing, but if the user provides the loaders to the trainer, these functions are patched with callable classes also containing the loader object.

@author:
This doesn't play nice with ddp_spawn and ddp_cpu. I needed to make this optional through a context manager.