Enable precision autocast for LightningModule step methods in Fabric (#17439)

(cherry picked from commit d9b4ebd)

Author: awaelchli

src/lightning/fabric/CHANGELOG.md

@@ -9,7 +9,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Changed
 
--
+
+- Enable precision autocast for LightningModule step methods in Fabric ([#17439](https://github.com/Lightning-AI/lightning/pull/17439))
 
 
 ### Fixed

src/lightning/fabric/wrappers.py

@@ -154,7 +154,7 @@ def wrapped_forward(*args: Any, **kwargs: Any) -> Any:
         def call_forward_module(*args: Any, **kwargs: Any) -> Any:
             # Patch the original_module's forward so we can redirect the arguments back to the real method
             self._original_module.forward = wrapped_forward
-            return self._forward_module(*args, **kwargs)
+            return self.forward(*args, **kwargs)
 
         return call_forward_module
 

tests/tests_fabric/test_wrappers.py

@@ -21,6 +21,7 @@
 from torch.utils.data.dataloader import DataLoader
 
 from lightning.fabric.fabric import Fabric
+from lightning.fabric.plugins import Precision
 from lightning.fabric.utilities.device_dtype_mixin import _DeviceDtypeModuleMixin
 from lightning.fabric.wrappers import _FabricDataLoader, _FabricModule, _FabricOptimizer, is_wrapped
 from tests_fabric.helpers.runif import RunIf
@@ -417,9 +418,10 @@ def validation_step(self, arg, kwarg=None):
         def normal_method(self):
             pass
 
+    precision = Mock(wraps=Precision())
     original_module = LightningModule()
     forward_module = DDP(original_module)
-    fabric_module = _FabricModule(forward_module=forward_module, precision=Mock(), original_module=original_module)
+    fabric_module = _FabricModule(forward_module=forward_module, precision=precision, original_module=original_module)
 
     # Regular methods on the original_module are visible and identical on the fabric_module ...
     assert fabric_module.normal_method == original_module.normal_method
@@ -441,6 +443,7 @@ def normal_method(self):
     assert fabric_module.training_step("train_arg", kwarg="train_kwarg") == "training_step_return"
     assert fabric_module.training_step("train_arg", kwarg="train_kwarg") == "training_step_return"  # call 2nd time
     assert fabric_module.validation_step("val_arg", kwarg="val_kwarg") == "validation_step_return"
+    precision.forward_context.assert_called()
 
     # The forward method remains untouched/unpatched after the special methods have been called
     assert original_module.forward.__name__ == "forward"