Fix ShardedDataParallel has no attribute require_backward_grad_sync (#6915)

Co-authored-by: Kaushik B <[email protected]>

Author: awaelchli

CHANGELOG.md

@@ -237,6 +237,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Fixed `--gpus` default for parser returned by `Trainer.add_argparse_args` ([#6898](https://github.com/PyTorchLightning/pytorch-lightning/pull/6898))
 
 
+- Fixed `AttributeError for `require_backward_grad_sync` when running manual optimization with sharded plugin ([#6915](https://github.com/PyTorchLightning/pytorch-lightning/pull/6915))
+
+
 ## [1.2.7] - 2021-04-06
 
 ### Fixed

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

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

tests/plugins/test_sharded_plugin.py

@@ -278,3 +278,32 @@ 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(
+        default_root_dir=tmpdir,
+        accelerator=accelerator,
+        fast_dev_run=2,
+        gpus=2,
+    )
+    trainer.fit(model)