Fix double precision during evaluation (Lightning-AI#12983)

CHANGELOG.md

@@ -185,6 +185,12 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Fixed an issue causing zero-division error for empty dataloaders ([#12885](https://github.com/PyTorchLightning/pytorch-lightning/pull/12885))
 
 
+- Fixed `Trainer(precision=64)` during evaluation which now uses the wrapped precision module ([#12983](https://github.com/PyTorchLightning/pytorch-lightning/pull/12983))
+
+
+- Fixed an issue to use wrapped `LightningModule` for evaluation during `trainer.fit` for `BaguaStrategy` ([#12983](https://github.com/PyTorchLightning/pytorch-lightning/pull/12983))
+
+
 - Fixed mismatching default values for the types of some arguments in the DeepSpeed and Fully-Sharded strategies which made the CLI unable to use them ([#12989](https://github.com/PyTorchLightning/pytorch-lightning/pull/12989))
 
 

pytorch_lightning/strategies/ddp.py

@@ -342,20 +342,20 @@ def training_step(self, *args, **kwargs) -> STEP_OUTPUT:
 
     def validation_step(self, *args, **kwargs) -> Optional[STEP_OUTPUT]:
         with self.precision_plugin.val_step_context():
-            if isinstance(self.model, DistributedDataParallel):
+            if self.lightning_module.trainer.state.fn == TrainerFn.FITTING:
                 # used when calling `trainer.fit`
                 return self.model(*args, **kwargs)
             else:
                 # used when calling `trainer.validate`
-                return self.lightning_module.validation_step(*args, **kwargs)
+                return self.model.validation_step(*args, **kwargs)
 
     def test_step(self, *args, **kwargs) -> Optional[STEP_OUTPUT]:
         with self.precision_plugin.test_step_context():
-            return self.lightning_module.test_step(*args, **kwargs)
+            return self.model.test_step(*args, **kwargs)
 
     def predict_step(self, *args, **kwargs) -> STEP_OUTPUT:
         with self.precision_plugin.predict_step_context():
-            return self.lightning_module.predict_step(*args, **kwargs)
+            return self.model.predict_step(*args, **kwargs)
 
     def post_training_step(self):
         if not self.lightning_module.automatic_optimization:

pytorch_lightning/strategies/ddp_spawn.py

@@ -241,20 +241,20 @@ def training_step(self, *args, **kwargs) -> STEP_OUTPUT:
 
     def validation_step(self, *args, **kwargs) -> Optional[STEP_OUTPUT]:
         with self.precision_plugin.val_step_context():
-            if isinstance(self.model, DistributedDataParallel):
+            if self.lightning_module.trainer.state.fn == TrainerFn.FITTING:
                 # used when calling `trainer.fit`
                 return self.model(*args, **kwargs)
             else:
                 # used when calling `trainer.validate`
-                return self.lightning_module.validation_step(*args, **kwargs)
+                return self.model.validation_step(*args, **kwargs)
 
     def test_step(self, *args, **kwargs) -> Optional[STEP_OUTPUT]:
         with self.precision_plugin.test_step_context():
-            return self.lightning_module.test_step(*args, **kwargs)
+            return self.model.test_step(*args, **kwargs)
 
     def predict_step(self, *args, **kwargs) -> STEP_OUTPUT:
         with self.precision_plugin.predict_step_context():
-            return self.lightning_module.predict_step(*args, **kwargs)
+            return self.model.predict_step(*args, **kwargs)
 
     def post_training_step(self):
         if not self.lightning_module.automatic_optimization:

tests/plugins/test_double_plugin.py

@@ -163,6 +163,7 @@ def test_double_precision_ddp(tmpdir):
         log_every_n_steps=1,
     )
     trainer.fit(model)
+    trainer.validate(model)
 
 
 def test_double_precision_pickle(tmpdir):