Repository: Lightning-AI/pytorch-lightning
Pull request: Fix ShardedDataParallel has no attribute require_backward_grad_sync
## What does this PR do?

Fixes #6804 

## Before submitting
- [x] Was this discussed/approved via a GitHub issue? (not for typos and docs)
- [x] Did you read the [contributor guideline](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/.github/CONTRIBUTING.md), **Pull Request** section?
- [x] Did you make sure your PR does only one thing, instead of bundling different changes together?
- [x] Did you make sure to update the documentation with your changes? (if necessary)
- [x] 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)

<!-- 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:
'LightningShardedDataParallel' object has no attribute 'require_backward_grad_sync'
When trying to use sharded DDP with the latest Lightning version 1.2.6 and Fairscale 0.3.3, I get the following error. Note that I am also using manual optimization and disabling automatic, as I am training a GAN.

```
Traceback (most recent call last):
  File "run_train.py", line 41, in train
    trainer.fit(model, data_loader)
  File "/home/timbrooks/anaconda3/envs/human_gan_fairseq/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 499, in fit
    self.dispatch()
  File "/home/timbrooks/anaconda3/envs/human_gan_fairseq/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 546, in dispatch
    self.accelerator.start_training(self)
  File "/home/timbrooks/anaconda3/envs/human_gan_fairseq/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py", line 73, in start_training
    self.training_type_plugin.start_training(trainer)
  File "/home/timbrooks/anaconda3/envs/human_gan_fairseq/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/training_type_plugin.py", line 114, in start_training
    self._results = trainer.run_train()
  File "/home/timbrooks/anaconda3/envs/human_gan_fairseq/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 637, in run_train
    self.train_loop.run_training_epoch()
  File "/home/timbrooks/anaconda3/envs/human_gan_fairseq/lib/python3.8/site-packages/pytorch_lightning/trainer/training_loop.py", line 492, in run_training_epoch
    batch_output = self.run_training_batch(batch, batch_idx, dataloader_idx)
  File "/home/timbrooks/anaconda3/envs/human_gan_fairseq/lib/python3.8/site-packages/pytorch_lightning/trainer/training_loop.py", line 657, in run_training_batch
    self._curr_step_result = self.training_step(
  File "/home/timbrooks/anaconda3/envs/human_gan_fairseq/lib/python3.8/site-packages/pytorch_lightning/trainer/training_loop.py", line 293, in training_step
    training_step_output = self.trainer.accelerator.training_step(args)
  File "/home/timbrooks/anaconda3/envs/human_gan_fairseq/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py", line 156, in training_step
    return self.training_type_plugin.training_step(*args)
  File "/home/timbrooks/anaconda3/envs/human_gan_fairseq/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/ddp.py", line 297, in training_step
    return self.model(*args, **kwargs)
  File "/home/timbrooks/anaconda3/envs/human_gan_fairseq/lib/python3.8/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/timbrooks/anaconda3/envs/human_gan_fairseq/lib/python3.8/site-packages/fairscale/nn/data_parallel/sharded_ddp.py", line 223, in forward
    return self.module(*inputs, **kwargs)
  File "/home/timbrooks/anaconda3/envs/human_gan_fairseq/lib/python3.8/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/timbrooks/anaconda3/envs/human_gan_fairseq/lib/python3.8/site-packages/pytorch_lightning/overrides/base.py", line 48, in forward
    output = self.module.training_step(*inputs, **kwargs)
  File "/home/timbrooks/code/human_gan/models/human_gan.py", line 121, in training_step
    self._discriminator_step(real_image, keypoints)
  File "/home/timbrooks/code/human_gan/models/human_gan.py", line 182, in _discriminator_step
    self.manual_backward(loss, optimizer)
  File "/home/timbrooks/anaconda3/envs/human_gan_fairseq/lib/python3.8/site-packages/pytorch_lightning/core/lightning.py", line 1228, in manual_backward
    self.trainer.train_loop.backward(loss, optimizer=None, opt_idx=None, *args, **kwargs)
  File "/home/timbrooks/anaconda3/envs/human_gan_fairseq/lib/python3.8/site-packages/pytorch_lightning/trainer/training_loop.py", line 782, in backward
    self.trainer.accelerator.backward(result, optimizer, opt_idx, should_accumulate, *args, **kwargs)
  File "/home/timbrooks/anaconda3/envs/human_gan_fairseq/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py", line 254, in backward
    self.training_type_plugin.pre_backward(closure_loss, should_accumulate, optimizer, optimizer_idx)
  File "/home/timbrooks/anaconda3/envs/human_gan_fairseq/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/ddp.py", line 283, in pre_backward
    if not self.lightning_module.automatic_optimization and self.model.require_backward_grad_sync:
  File "/home/timbrooks/anaconda3/envs/human_gan_fairseq/lib/python3.8/site-packages/fairscale/nn/data_parallel/sharded_ddp.py", line 362, in __getattr__
    return getattr(self.module, name)
  File "/home/timbrooks/anaconda3/envs/human_gan_fairseq/lib/python3.8/site-packages/torch/nn/modules/module.py", line 947, in __getattr__
    raise AttributeError("'{}' object has no attribute '{}'".format(
AttributeError: 'LightningShardedDataParallel' object has no attribute 'require_backward_grad_sync'
```

Hints from issue discussion:
@SeanNaren is the correct fix here to just not access `self.model.require_backward_grad_sync` in the sharded plugin?

Commit message: test


test


test


update test


fix spawn


update plugins


revert


wip


sharded


fix name


c

Patch to review:
diff --git a/pytorch_lightning/plugins/training_type/sharded.py b/pytorch_lightning/plugins/training_type/sharded.py
index 7536ef9b1d856..eaa9331d280b4 100644
--- a/pytorch_lightning/plugins/training_type/sharded.py
+++ b/pytorch_lightning/plugins/training_type/sharded.py
@@ -13,6 +13,9 @@
 # limitations under the License.
 from typing import Optional
 
+import torch
+from torch.optim import Optimizer
+
 from pytorch_lightning.core.lightning import LightningModule
 from pytorch_lightning.core.optimizer import is_lightning_optimizer
 from pytorch_lightning.plugins.training_type.ddp import DDPPlugin
@@ -33,6 +36,7 @@ def configure_ddp(self):
         self._model = ShardedDataParallel(
             LightningShardedDataParallel(self.model), sharded_optimizer=self.lightning_module.trainer.optimizers
         )
+        setattr(self._model, "require_backward_grad_sync", False)
 
     def _reinit_optimizers_with_oss(self):
         optimizers = self.lightning_module.trainer.optimizers
@@ -70,3 +74,9 @@ def _optim_state_dict(self, optimizer):
     @property
     def lightning_module(self) -> LightningModule:
         return unwrap_lightning_module_sharded(self._model)
+
+    def pre_backward(self, closure_loss: torch.Tensor, should_accumulate: bool, optimizer: Optimizer, opt_idx: int):
+        pass
+
+    def post_training_step(self):
+        pass
diff --git a/pytorch_lightning/plugins/training_type/sharded_spawn.py b/pytorch_lightning/plugins/training_type/sharded_spawn.py
index 7aadf797e160a..dc1e5e83843aa 100644
--- a/pytorch_lightning/plugins/training_type/sharded_spawn.py
+++ b/pytorch_lightning/plugins/training_type/sharded_spawn.py
@@ -13,6 +13,9 @@
 # limitations under the License.
 from typing import Optional
 
+import torch
+from torch.optim import Optimizer
+
 from pytorch_lightning.core.lightning import LightningModule
 from pytorch_lightning.plugins.training_type.ddp_spawn import DDPSpawnPlugin
 from pytorch_lightning.trainer.states import TrainerState
@@ -32,6 +35,7 @@ def configure_ddp(self):
         self._model = ShardedDataParallel(
             LightningShardedDataParallel(self.model), sharded_optimizer=self.lightning_module.trainer.optimizers
         )
+        setattr(self._model, "require_backward_grad_sync", False)
 
     def _reinit_optimizers_with_oss(self):
         optimizers = self.lightning_module.trainer.optimizers
@@ -65,3 +69,9 @@ def _optim_state_dict(self, optimizer):
     @property
     def lightning_module(self) -> LightningModule:
         return unwrap_lightning_module_sharded(self._model)
+
+    def pre_backward(self, closure_loss: torch.Tensor, should_accumulate: bool, optimizer: Optimizer, opt_idx: int):
+        pass
+
+    def post_training_step(self):
+        pass
diff --git a/tests/plugins/test_sharded_plugin.py b/tests/plugins/test_sharded_plugin.py
index 655e12f046e04..7903cddffdb8b 100644
--- a/tests/plugins/test_sharded_plugin.py
+++ b/tests/plugins/test_sharded_plugin.py
@@ -278,3 +278,31 @@ def test_ddp_sharded_plugin_test_multigpu(tmpdir, trainer_kwargs):
 
     trainer.validate(model)
     trainer.test(model)
+
+
+class ManualBoringModel(BoringModel):
+
+    def __init__(self):
+        super().__init__()
+        self.automatic_optimization = False
+
+    def training_step(self, batch, batch_idx):
+        opt = self.optimizers()
+        opt.zero_grad()
+        output = self(batch)
+        loss = self.loss(batch, output)
+        self.manual_backward(loss)
+        opt.step()
+        return {"loss": loss}
+
+
+@RunIf(skip_windows=True, special=True, fairscale=True, min_gpus=2)
+@pytest.mark.parametrize("accelerator", ["ddp_sharded", "ddp_sharded_spawn"])
+def test_ddp_sharded_plugin_manual_optimization(tmpdir, accelerator):
+    model = ManualBoringModel()
+    trainer = Trainer(
+        accelerator=accelerator,
+        fast_dev_run=2,
+        gpus=2,
+    )
+    trainer.fit(model)


Code review comments:
diff_hunk: @@ -278,3 +278,31 @@ def test_ddp_sharded_plugin_test_multigpu(tmpdir, trainer_kwargs):
 
     trainer.validate(model)
     trainer.test(model)
+
+
+class ManualBoringModel(BoringModel):
+
+    def __init__(self):
+        super().__init__()
+        self.automatic_optimization = False
+
+    def training_step(self, batch, batch_idx):
+        opt = self.optimizers()
+        opt.zero_grad()
+        output = self(batch)
+        loss = self.loss(batch, output)
+        self.manual_backward(loss)
+        opt.step()
+        return {"loss": loss}
+
+
+@RunIf(skip_windows=True, special=True, fairscale=True, min_gpus=2)
+@pytest.mark.parametrize("accelerator", ["ddp_sharded", "ddp_sharded_spawn"])
+def test_ddp_sharded_plugin_manual_optimization(tmpdir, accelerator):
+    model = ManualBoringModel()
+    trainer = Trainer(
+        accelerator=accelerator,
path: tests/plugins/test_sharded_plugin.py
text: @user1:
```suggestion
    trainer = Trainer(
    	default_root_dir=tmpdir
        accelerator=accelerator,
```