Repository: Lightning-AI/pytorch-lightning
Pull request: Fix LightningModule step methods bypassing DDP wrapper in Fabric
## What does this PR do?

Fixes #16856

Context:
In PyTorch, if you wrap your module into e.g. DistributedDataParallel
```py
model = ...
ddp_model = torch.nn.parallel.DistributedDataParallel(model, ...)
```
everything that you want to pass through that model **needs** to go through `ddp_model.forward`! Like so:
```py
out = ddp_model(batch)
```
All of the following would silently work but are **incorrect**, as it would render DDP completely useless:
```py
out = model(batch)
out = ddp_model.module(batch)
out = model.training_step(batch, batch_idx)
```
However, especially the usage in the last statement is what we want the user experience to be when working with a LightningModule in a Fabric training loop. Hence, we need to somehow simulate that `model.training_step(batch, batch_idx)` behaves like `ddp_model(batch, batch_idx)` without the user having to change anything in their code.

This PR does exactly that using black magic 300 IQ big brain logic. The code for this solution is a bit complicated to understand, but it is concise and confined to a single place in the _FabricModule. Alternatives we considered would be brittle and spread all over the strategies.








cc @borda @carmocca @justusschock @awaelchli

Problem statement:
Calls to LightningModule's `*_step` methods in Fabric bypasses DDP wrapper
### Bug description

In Fabric, we wrap the model in DDP, but all inputs going through the model need to go through forward(). If you call other methods on the module, like `training_step` etc., it bypasses the DDP wrapper. 

### How to reproduce the bug

```python
fabric = Fabric(devices=2)

model = fabric.setup(model)
out = model.training_step(x)
```


### Error messages and logs

For correctness, we should not bypass forward and let DDP run the necessary communication hooks. 

### Environment

<details>
  <summary>Current environment</summary>

```
#- Lightning Component (e.g. Trainer, LightningModule, LightningApp, LightningWork, LightningFlow):
#- PyTorch Lightning Version (e.g., 1.5.0):
#- Lightning App Version (e.g., 0.5.2):
#- PyTorch Version (e.g., 2.0):
#- Python version (e.g., 3.9):
#- OS (e.g., Linux):
#- CUDA/cuDNN version:
#- GPU models and configuration:
#- How you installed Lightning(`conda`, `pip`, source):
#- Running environment of LightningApp (e.g. local, cloud):
```

</details>


### More info

_No response_

cc @tchaton @carmocca @justusschock @awaelchli

Hints from issue discussion:
What I've reported here has now also been seen by a user on Discord: https://discord.com/channels/1077906959069626439/1097919761100316832/1097923294180692090

It is not possible for the user to just know they need to pass everything through forward and so I'll brainstorm a solution for this. One possibility is to wrap each user-method and call the ddp hooks ourselves: https://github.com/pytorch/pytorch/blob/06f19fdbe505be88e478e22a0e38bb15d5712c42/torch/nn/parallel/distributed.py#L1529-L1537
I investigated whether we can use our prepare_for_backward https://github.com/Lightning-AI/lightning/blob/26a549e87fba86fbec27a5231fb79540eb884b71/src/lightning/pytorch/overrides/distributed.py#L43 function in the fabric.backward() call, (like we have in the Trainer used for manual opt mode), and in theory it works but it requires the model as input and that's not available in fabric.backward() (and we don't want to force the user to pass it in). 

So I'm now investigating the original ugly idea which would introduce another wrapper between the FabricModule and the DDP wrapper.

Commit message: [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Patch to review:
diff --git a/src/lightning/fabric/fabric.py b/src/lightning/fabric/fabric.py
index 096131ada0f28..5ebb04bc9cc71 100644
--- a/src/lightning/fabric/fabric.py
+++ b/src/lightning/fabric/fabric.py
@@ -363,7 +363,7 @@ def backward(self, tensor: Tensor, *args: Any, model: Optional[_FabricModule] =
                 # requires to attach the current `DeepSpeedEngine` for the `_FabricOptimizer.step` call.
                 self._strategy._deepspeed_engine = module
 
-        self._precision.backward(tensor, module, *args, **kwargs)
+        self._strategy.backward(tensor, module, *args, **kwargs)
 
     def clip_gradients(
         self,
diff --git a/src/lightning/fabric/wrappers.py b/src/lightning/fabric/wrappers.py
index b2c3705d47a5c..6f59d801d742b 100644
--- a/src/lightning/fabric/wrappers.py
+++ b/src/lightning/fabric/wrappers.py
@@ -15,6 +15,7 @@
 from typing import Any, Callable, Dict, Generator, Iterator, Mapping, Optional, overload, TypeVar, Union
 
 import torch
+from lightning_utilities import WarningCache
 from lightning_utilities.core.apply_func import apply_to_collection
 from torch import nn as nn
 from torch import Tensor
@@ -28,7 +29,9 @@
 from lightning.fabric.utilities.data import _set_sampler_epoch
 from lightning.fabric.utilities.device_dtype_mixin import _DeviceDtypeModuleMixin
 from lightning.fabric.utilities.types import Optimizable
+from lightning.fabric.utilities.warnings import PossibleUserWarning
 
+warning_cache = WarningCache()
 T_destination = TypeVar("T_destination", bound=Dict[str, Any])
 
 
@@ -132,7 +135,35 @@ def state_dict(
     def load_state_dict(self, state_dict: Mapping[str, Any], strict: bool = True) -> _IncompatibleKeys:
         return self._original_module.load_state_dict(state_dict=state_dict, strict=strict)
 
+    def _redirection_through_forward(self, method_name: str) -> Callable:
+        assert method_name != "forward"
+        original_forward = self._original_module.forward
+
+        def wrapped_forward(*args: Any, **kwargs: Any) -> Any:
+            # Unpatch ourselves immediately before calling the method `method_name`
+            # because itself may want to call the real `forward`
+            self._original_module.forward = original_forward
+            # Call the actual method e.g. `.training_step(...)`
+            return getattr(self._original_module, method_name)(*args, **kwargs)
+
+        # We make the caller "unknowingly" send their arguments through the forward_module's `__call__`.
+        # We expect that the `forward_module` will eventually call `original_module.forward`, which we
+        # have patched to redirect back to `original_module.method_name()`.
+        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 call_forward_module
+
     def __getattr__(self, item: Any) -> Any:
+        if (
+            item in ("training_step", "validation_step", "test_step", "predict_step")
+            and self._forward_module != self._original_module
+        ):
+            # Special support for `LightningModule`, to prevent bypassing DDP's forward
+            return self._redirection_through_forward(item)
+
         try:
             # __getattr__ gets called as a last resort if the attribute does not exist
             # call nn.Module's implementation first
@@ -140,7 +171,16 @@ def __getattr__(self, item: Any) -> Any:
         except AttributeError:
             # If the attribute is not available on the _FabricModule wrapper, redirect to the wrapped nn.Module
             original_module = super().__getattr__("_original_module")
-            return getattr(original_module, item)
+
+            attr = getattr(original_module, item)
+            if inspect.ismethod(attr) and self._forward_module != self._original_module:
+                warning_cache.warn(
+                    f"You are calling the method `{type(self._original_module).__name__}.{item}()` from outside the"
+                    " model. This will bypass the wrapper from the strategy and result in incorrect behavior in"
+                    f" `.backward()`. You should pass your inputs through `{type(self._original_module)}.forward()`.",
+                    category=PossibleUserWarning,
+                )
+            return attr
 
 
 class _FabricDataLoader:
diff --git a/tests/tests_fabric/test_wrappers.py b/tests/tests_fabric/test_wrappers.py
index acd4f41e6e3e9..87cc4a3271104 100644
--- a/tests/tests_fabric/test_wrappers.py
+++ b/tests/tests_fabric/test_wrappers.py
@@ -16,6 +16,7 @@
 
 import pytest
 import torch
+from lightning_utilities.test.warning import no_warning_call
 from torch.utils.data import BatchSampler, DistributedSampler
 from torch.utils.data.dataloader import DataLoader
 
@@ -59,13 +60,43 @@ def __init__(self):
     fabric_module = _FabricModule(wrapped_module, Mock(), original_module=original_module)
     assert fabric_module.attribute == 1
     assert fabric_module.layer is original_module.layer
-    assert fabric_module.method() == 2
     assert fabric_module.forward.__self__.__class__ == _FabricModule
 
     with pytest.raises(AttributeError):
         _ = fabric_module.not_exists
 
 
+def test_fabric_module_method_lookup():
+    """Test that access to methods warns about inproper use when a wrapper from a strategy is involved."""
+    from lightning.fabric.wrappers import warning_cache
+
+    class OriginalModule(torch.nn.Module):
+        def method(self):
+            return 100
+
+    class ModuleWrapper(torch.nn.Module):
+        def __init__(self, module):
+            super().__init__()
+            self.wrapped = module
+
+    # Regular case: forward_module == original_module -> no warnings
+    original_module = OriginalModule()
+    fabric_module = _FabricModule(forward_module=original_module, precision=Mock(), original_module=original_module)
+    warning_cache.clear()
+    with no_warning_call(UserWarning):
+        assert fabric_module.method() == 100
+    assert not warning_cache
+
+    # Special case: original module wrapped by forward module: -> warn
+    original_module = OriginalModule()
+    wrapped_module = ModuleWrapper(original_module)
+    fabric_module = _FabricModule(forward_module=wrapped_module, precision=Mock(), original_module=original_module)
+    warning_cache.clear()
+    with pytest.warns(UserWarning, match=r"You are calling the method `OriginalModule.method\(\)` from outside the"):
+        assert fabric_module.method() == 100
+    warning_cache.clear()
+
+
 def test_fabric_module_state_dict_access():
     """Test that state_dict access passes through to the original module."""
 
@@ -353,3 +384,64 @@ def test_is_wrapped():
     assert not is_wrapped(dataloader)
     wrapped = _FabricDataLoader(dataloader)
     assert is_wrapped(wrapped)
+
+
+def test_step_method_redirection():
+    """Test that the FabricModule redirects the special `LightningModule.*_step` methods through the forward-
+    module."""
+
+    class DDP(torch.nn.Module):
+        def __init__(self, module):
+            super().__init__()
+            self.module = module
+
+        def forward(self, *args, **kwargs):
+            return self.module(*args, **kwargs)
+
+    class LightningModule(torch.nn.Module):
+        def forward(self):
+            return "forward_return"
+
+        def training_step(self, arg, kwarg=None):
+            assert self() == "forward_return"
+            assert arg == "train_arg"
+            assert kwarg == "train_kwarg"
+            return "training_step_return"
+
+        def validation_step(self, arg, kwarg=None):
+            assert self() == "forward_return"
+            assert arg == "val_arg"
+            assert kwarg == "val_kwarg"
+            return "validation_step_return"
+
+        def normal_method(self):
+            pass
+
+    original_module = LightningModule()
+    forward_module = DDP(original_module)
+    fabric_module = _FabricModule(forward_module=forward_module, precision=Mock(), original_module=original_module)
+
+    # Regular methods on the original_module are visible and identiacal on the fabric_module ...
+    assert fabric_module.normal_method == original_module.normal_method
+
+    # ... but speical methods like training_step get redirected to the forward_module
+    assert fabric_module.training_step.__name__ == "call_forward_module"
+    assert fabric_module.validation_step.__name__ == "call_forward_module"
+    assert fabric_module.test_step.__name__ == "call_forward_module"
+    assert fabric_module.predict_step.__name__ == "call_forward_module"
+
+    with pytest.raises(AttributeError, match="has no attribute 'predict_step'"):
+        # A special method that does not exist will raise its AttributeError when being called
+        fabric_module.predict_step()
+
+    # The forward method on the original module remains untouched
+    assert original_module.forward.__name__ == "forward"
+
+    # The special methods get redirected correctly to produce the expected output
+    assert fabric_module.training_step("train_arg", kwarg="train_kwarg") == "training_step_return"
+    assert fabric_module.validation_step("val_arg", kwarg="val_kwarg") == "validation_step_return"
+
+    # Special case: forward_module == original_module -> no special treatment applied
+    fabric_module = _FabricModule(forward_module=original_module, precision=Mock(), original_module=original_module)
+    assert fabric_module.training_step == original_module.training_step
+    assert fabric_module.validation_step == original_module.validation_step


Code review comments:
diff_hunk: @@ -353,3 +384,64 @@ def test_is_wrapped():
     assert not is_wrapped(dataloader)
     wrapped = _FabricDataLoader(dataloader)
     assert is_wrapped(wrapped)
+
+
+def test_step_method_redirection():
+    """Test that the FabricModule redirects the special `LightningModule.*_step` methods through the forward-
+    module."""
+
+    class DDP(torch.nn.Module):
+        def __init__(self, module):
+            super().__init__()
+            self.module = module
+
+        def forward(self, *args, **kwargs):
+            return self.module(*args, **kwargs)
+
+    class LightningModule(torch.nn.Module):
+        def forward(self):
+            return "forward_return"
+
+        def training_step(self, arg, kwarg=None):
+            assert self() == "forward_return"
+            assert arg == "train_arg"
+            assert kwarg == "train_kwarg"
+            return "training_step_return"
+
+        def validation_step(self, arg, kwarg=None):
+            assert self() == "forward_return"
+            assert arg == "val_arg"
+            assert kwarg == "val_kwarg"
+            return "validation_step_return"
+
+        def normal_method(self):
+            pass
+
+    original_module = LightningModule()
+    forward_module = DDP(original_module)
+    fabric_module = _FabricModule(forward_module=forward_module, precision=Mock(), original_module=original_module)
+
+    # Regular methods on the original_module are visible and identiacal on the fabric_module ...
+    assert fabric_module.normal_method == original_module.normal_method
+
+    # ... but speical methods like training_step get redirected to the forward_module
path: tests/tests_fabric/test_wrappers.py
text: @author:
```suggestion
    # ... but special methods like training_step get redirected to the forward_module
```