update non training steps in ddp and ddp_spawn

pytorch_lightning/plugins/training_type/ddp.py

@@ -290,16 +290,6 @@ def pre_configure_ddp(self):
             self._ddp_kwargs["find_unused_parameters"] = True
 
     def _wrap_model(self) -> None:
-        # skip warpping the model if we are not fitting as no gradients need to be exchanged
-        trainer_fn = self.lightning_module.trainer.state.fn
-        if trainer_fn != TrainerFn.FITTING:
-            self._model = (
-                LightningDistributedModule(self.model)
-                if not isinstance(self.model, (LightningDistributedModule))
-                else self.model
-            )
-            rank_zero_debug(f"In {trainer_fn} stage: Skipping wrapping the model with DistributedDataParallel")
-            return
         self._model = DistributedDataParallel(
             LightningDistributedModule(self.model), device_ids=self.determine_ddp_device_ids(), **self._ddp_kwargs
         )
@@ -318,6 +308,11 @@ def _register_ddp_hooks(self) -> None:
             )
 
     def configure_ddp(self):
+        # skip warpping the model if we are not fitting as no gradients need to be exchanged
+        trainer_fn = self.lightning_module.trainer.state.fn
+        if trainer_fn != TrainerFn.FITTING:
+            rank_zero_debug(f"In {trainer_fn} stage: Skipping wrapping the model with DistributedDataParallel")
+            return
         self.pre_configure_ddp()
         self._wrap_model()
         self._register_ddp_hooks()
@@ -398,16 +393,28 @@ def reduce(self, tensor, group: Optional[Any] = None, reduce_op: Union[ReduceOp,
         return tensor
 
     def training_step(self, *args, **kwargs):
-        return self.model(*args, **kwargs)
+        if isinstance(self.model, DistributedDataParallel):
+            return self.model(*args, **kwargs)
+        else:
+            return self.model.training_step(*args, **kwargs)
 
     def validation_step(self, *args, **kwargs):
-        return self.model(*args, **kwargs)
+        if isinstance(self.model, DistributedDataParallel):
+            return self.model(*args, **kwargs)
+        else:
+            return self.model.validation_step(*args, **kwargs)
 
     def test_step(self, *args, **kwargs):
-        return self.model(*args, **kwargs)
+        if isinstance(self.model, DistributedDataParallel):
+            return self.model(*args, **kwargs)
+        else:
+            return self.model.test_step(*args, **kwargs)
 
     def predict_step(self, *args, **kwargs):
-        return self.model(*args, **kwargs)
+        if isinstance(self.model, DistributedDataParallel):
+            return self.model(*args, **kwargs)
+        else:
+            return self.model.predict_step(*args, **kwargs)
 
     def post_training_step(self):
         if not self.lightning_module.automatic_optimization:

pytorch_lightning/plugins/training_type/ddp_spawn.py

@@ -235,16 +235,6 @@ def pre_configure_ddp(self):
             self._ddp_kwargs["find_unused_parameters"] = True
 
     def _wrap_model(self) -> None:
-        # skip warpping the model if we are not fitting as no gradients need to be exchanged
-        trainer_fn = self.lightning_module.trainer.state.fn
-        if trainer_fn != TrainerFn.FITTING:
-            self._model = (
-                LightningDistributedModule(self.model)
-                if not isinstance(self.model, (LightningDistributedModule))
-                else self.model
-            )
-            rank_zero_debug(f"In {trainer_fn} stage: Skipping wrapping the model with DistributedDataParallel")
-            return
         self._model = DistributedDataParallel(
             LightningDistributedModule(self.model), device_ids=self.determine_ddp_device_ids(), **self._ddp_kwargs
         )
@@ -261,6 +251,11 @@ def _register_ddp_hooks(self) -> None:
             )
 
     def configure_ddp(self):
+        # skip warpping the model if we are not fitting as no gradients need to be exchanged
+        trainer_fn = self.lightning_module.trainer.state.fn
+        if trainer_fn != TrainerFn.FITTING:
+            rank_zero_debug(f"In {trainer_fn} stage: Skipping wrapping the model with DistributedDataParallel")
+            return
         self.pre_configure_ddp()
         self._wrap_model()
         self._register_ddp_hooks()
@@ -370,16 +365,28 @@ def reduce(self, tensor, group: Optional[Any] = None, reduce_op: Union[ReduceOp,
         return tensor
 
     def training_step(self, *args, **kwargs):
-        return self.model(*args, **kwargs)
+        if isinstance(self.model, DistributedDataParallel):
+            return self.model(*args, **kwargs)
+        else:
+            return self.model.training_step(*args, **kwargs)
 
     def validation_step(self, *args, **kwargs):
-        return self.model(*args, **kwargs)
+        if isinstance(self.model, DistributedDataParallel):
+            return self.model(*args, **kwargs)
+        else:
+            return self.model.validation_step(*args, **kwargs)
 
     def test_step(self, *args, **kwargs):
-        return self.model(*args, **kwargs)
+        if isinstance(self.model, DistributedDataParallel):
+            return self.model(*args, **kwargs)
+        else:
+            return self.model.test_step(*args, **kwargs)
 
     def predict_step(self, *args, **kwargs):
-        return self.model(*args, **kwargs)
+        if isinstance(self.model, DistributedDataParallel):
+            return self.model(*args, **kwargs)
+        else:
+            return self.model.predict_step(*args, **kwargs)
 
     def post_training_step(self):
         if not self.lightning_module.automatic_optimization:

pytorch_lightning/plugins/training_type/sharded.py

@@ -36,6 +36,11 @@ class DDPShardedPlugin(DDPPlugin):
     _REDUCE_BUFFER_SIZE_DEFAULT = 2 ** 23  # 8M
 
     def configure_ddp(self):
+        # skip warpping the model if we are not fitting as no gradients need to be exchanged
+        trainer_fn = self.lightning_module.trainer.state.fn
+        if trainer_fn != TrainerFn.FITTING:
+            rank_zero_debug(f"In {trainer_fn} stage: Skipping wrapping the model with ShardedDataParallel")
+            return
         self._wrap_optimizers()
         self._wrap_model()
         setattr(self._model, "require_backward_grad_sync", False)
@@ -67,12 +72,6 @@ def _wrap_optimizers(self):
         self._reinit_optimizers_with_oss()
 
     def _wrap_model(self) -> None:
-        # skip warpping the model if we are not fitting as no gradients need to be exchanged
-        trainer_fn = self.lightning_module.trainer.state.fn
-        if trainer_fn != TrainerFn.FITTING:
-            self._model = self.model
-            rank_zero_debug(f"In {trainer_fn} stage: Skipping wrapping the model with ShardedDataParallel")
-            return
         self._model = ShardedDataParallel(
             LightningShardedDataParallel(self.model),
             sharded_optimizer=self.lightning_module.trainer.optimizers,

pytorch_lightning/plugins/training_type/sharded_spawn.py

@@ -35,6 +35,11 @@ class DDPSpawnShardedPlugin(DDPSpawnPlugin):
     """Optimizer sharded training provided by FairScale."""
 
     def configure_ddp(self):
+        # skip warpping the model if we are not fitting as no gradients need to be exchanged
+        trainer_fn = self.lightning_module.trainer.state.fn
+        if trainer_fn != TrainerFn.FITTING:
+            rank_zero_debug(f"In {trainer_fn} stage: Skipping wrapping the model with ShardedDataParallel")
+            return
         self._wrap_optimizers()
         self._wrap_model()
         setattr(self._model, "require_backward_grad_sync", False)
@@ -56,12 +61,6 @@ def _wrap_optimizers(self):
         self._reinit_optimizers_with_oss()
 
     def _wrap_model(self):
-        # skip warpping the model if we are not fitting as no gradients need to be exchanged
-        trainer_fn = self.lightning_module.trainer.state.fn
-        if trainer_fn != TrainerFn.FITTING:
-            self._model = self.model
-            rank_zero_debug(f"In {trainer_fn} stage: Skipping wrapping the model with ShardedDataParallel")
-            return
         self._model = ShardedDataParallel(
             LightningShardedDataParallel(self.model), sharded_optimizer=self.lightning_module.trainer.optimizers
         )