Repository: Lightning-AI/pytorch-lightning
Pull request: Remember the eval mode of submodules when switching trainer stages
## What does this PR do?

Fixes #18930
Part of #16827

A common issue users are facing is that the loop calls `train()` on the LightningModule despite the user having frozen certain layers. For example,
- user finetunes a model and freezes certain layers they don't want to train
- user's model has a feature extractor they don't train

This leads to a surprise when the user finds out that their batch norm layers have changed statistics, even though they were set explicitly to `eval()` mode. To avoid this, the user has to learn that they should override the `on_validation_model_eval()` and `on_validation_model_train()` hooks in the module, but this is a detail difficult to find in our docs and get right. Most users who face this challenge end up on slack or GH to ask for help.

The PR makes the following changes to automate this for the user:
- The validation loop captures the `.training` mode of every submodule before calling `.eval()` now. When the validation loop ends, and before switching to training, it restores the `.training` mode on all submodules to what it was before. This ensures that layers the user has chosen to be in eval mode **remain** in eval mode!
- The fit-loop no longer calls `.train()` at the beginning with the same motivation: The user can now set a subset of their model to `.eval()` mode / freeze it explicitly in the LightningModule's `__init__` without doing acrobatics with hooks, and the Trainer will respect it and preserve it (see the added test). Note: This is not a breaking change, because PyTorch's default is to have a model in `.training=True` mode.


<!-- readthedocs-preview pytorch-lightning start -->
----
:books: Documentation preview :books:: https://pytorch-lightning--18951.org.readthedocs.build/en/18951/

<!-- readthedocs-preview pytorch-lightning end -->

cc @borda @justusschock @awaelchli

Problem statement:
Identify parts of the models that are set in '.eval()' mode during initialization and keep them in eval mode after validation epoch ends
### Description & Motivation

Currently, if I want a submodule of my model completely **frozen** during training, I cannot do so because even with 

1. Setting this submodule in `.eval()` mode in `__init__()`
2. Setting all `required_grad` of all parameters in this submodule to `False`
3. Using `torch.no_grad()` context when passing data through this module

If my submodule contains ANY batch-normalization layers, these will not stay in `.eval()` mode because the whole model (which contains this submodule) is set to `.train()` mode after validation epoch ends. Hence, these batch-normalization layers' running stats will change and affect the performance and training of the entire model.

This behavior is very easy to create a minimal example to see this effect.

### Pitch

Search for params that are set in eval mode during initialization and keep them in eval mode after validation epoch ends.

### Alternatives

_No response_

### Additional context

_No response_

cc @borda

Hints from issue discussion:

Commit message: docstring

Patch to review:
diff --git a/docs/source-pytorch/advanced/transfer_learning.rst b/docs/source-pytorch/advanced/transfer_learning.rst
index 102555b94a455..935bcd4bb8773 100644
--- a/docs/source-pytorch/advanced/transfer_learning.rst
+++ b/docs/source-pytorch/advanced/transfer_learning.rst
@@ -66,13 +66,13 @@ Example: Imagenet (Computer Vision)
             num_filters = backbone.fc.in_features
             layers = list(backbone.children())[:-1]
             self.feature_extractor = nn.Sequential(*layers)
+            self.feature_extractor.eval()
 
             # use the pretrained model to classify cifar-10 (10 image classes)
             num_target_classes = 10
             self.classifier = nn.Linear(num_filters, num_target_classes)
 
         def forward(self, x):
-            self.feature_extractor.eval()
             with torch.no_grad():
                 representations = self.feature_extractor(x).flatten(1)
             x = self.classifier(representations)
diff --git a/docs/source-pytorch/common/lightning_module.rst b/docs/source-pytorch/common/lightning_module.rst
index 797b25f852b92..3f3fd0758bc59 100644
--- a/docs/source-pytorch/common/lightning_module.rst
+++ b/docs/source-pytorch/common/lightning_module.rst
@@ -172,8 +172,7 @@ Under the hood, Lightning does the following (pseudocode):
 
 .. code-block:: python
 
-    # put model in train mode and enable gradient calculation
-    model.train()
+    # enable gradient calculation
     torch.set_grad_enabled(True)
 
     for batch_idx, batch in enumerate(train_dataloader):
@@ -1096,7 +1095,6 @@ for more information.
 
 
     def fit_loop():
-        model.train()
         torch.set_grad_enabled(True)
 
         on_train_epoch_start()
diff --git a/docs/source-pytorch/common/trainer.rst b/docs/source-pytorch/common/trainer.rst
index 98b30650bff23..37c279e311ae0 100644
--- a/docs/source-pytorch/common/trainer.rst
+++ b/docs/source-pytorch/common/trainer.rst
@@ -54,8 +54,7 @@ Here's the pseudocode for what the trainer does under the hood (showing the trai
 
 .. code-block:: python
 
-    # put model in train mode
-    model.train()
+    # enable grads
     torch.set_grad_enabled(True)
 
     losses = []
diff --git a/src/lightning/pytorch/CHANGELOG.md b/src/lightning/pytorch/CHANGELOG.md
index 99e477bee767a..b40489e7ea01b 100644
--- a/src/lightning/pytorch/CHANGELOG.md
+++ b/src/lightning/pytorch/CHANGELOG.md
@@ -11,12 +11,15 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Added `lightning.pytorch.callbacks.ThroughputMonitor` to track throughput and log it ([#18848](https://github.com/Lightning-AI/lightning/pull/18848))
 
+- The Trainer now restores the training mode set through `.train()` or `.eval()` on a submodule-level when switching from validation to training ([#18951](https://github.com/Lightning-AI/lightning/pull/18951))
 
 
 ### Changed
 
 - `seed_everything()` without passing in a seed no longer randomly selects a seed, and now defaults to `0` ([#18846](https://github.com/Lightning-AI/lightning/pull/18846))
 - `LightningCLI` no longer allows setting a normal class instance as default. A `lazy_instance` can be used instead ([#18822](https://github.com/Lightning-AI/lightning/pull/18822))
+- The `LightningModule.on_{validation,test,predict}_model_{eval,train}` now only get called if they are overridden by the user ([#18951](https://github.com/Lightning-AI/lightning/pull/18951))
+- The `Trainer.fit()` loop no longer calls `LightningModule.train()` at the start; it now preserves the user's configuration of frozen layers ([#18951](https://github.com/Lightning-AI/lightning/pull/18951))
 
 
 ### Deprecated
diff --git a/src/lightning/pytorch/core/hooks.py b/src/lightning/pytorch/core/hooks.py
index 1f8ff226c4a43..28aa9ec655ad3 100644
--- a/src/lightning/pytorch/core/hooks.py
+++ b/src/lightning/pytorch/core/hooks.py
@@ -158,23 +158,55 @@ def on_validation_model_zero_grad(self) -> None:
         self.zero_grad(**zero_grad_kwargs)
 
     def on_validation_model_eval(self) -> None:
-        """Sets the model to eval during the val loop."""
+        """Called when the validation loop starts.
+
+        The validation loop by default calls ``.eval()`` on the LightningModule before it starts. Override this hook
+        to change the behavior. See also :meth:`~lightning.pytorch.core.hooks.ModelHooks.on_validation_model_train`.
+
+        """
+        # The loop won't call this hook unless it is overridden. The line below is here in case the user calls super().
         self.trainer.model.eval()
 
     def on_validation_model_train(self) -> None:
-        """Sets the model to train during the val loop."""
-        self.trainer.model.train()
+        """Called when the validation loop ends.
 
-    def on_test_model_train(self) -> None:
-        """Sets the model to train during the test loop."""
+        The validation loop by default restores the `training` mode of the LightningModule to what it was before
+        starting validation. Override this hook to change the behavior. See also
+        :meth:`~lightning.pytorch.core.hooks.ModelHooks.on_validation_model_eval`.
+
+        """
+        # The loop won't call this hook unless it is overridden. The line below is here in case the user calls super().
         self.trainer.model.train()
 
     def on_test_model_eval(self) -> None:
-        """Sets the model to eval during the test loop."""
+        """Called when the test loop starts.
+
+        The test loop by default calls ``.eval()`` on the LightningModule before it starts. Override this hook
+        to change the behavior. See also :meth:`~lightning.pytorch.core.hooks.ModelHooks.on_test_model_train`.
+
+        """
+        # The loop won't call this hook unless it is overridden. The line below is here in case the user calls super().
         self.trainer.model.eval()
 
+    def on_test_model_train(self) -> None:
+        """Called when the test loop ends.
+
+        The test loop by default restores the `training` mode of the LightningModule to what it was before
+        starting testing. Override this hook to change the behavior. See also
+        :meth:`~lightning.pytorch.core.hooks.ModelHooks.on_test_model_eval`.
+
+        """
+        # The loop won't call this hook unless it is overridden. The line below is here in case the user calls super().
+        self.trainer.model.train()
+
     def on_predict_model_eval(self) -> None:
-        """Sets the model to eval during the predict loop."""
+        """Called when the predict loop starts.
+
+        The predict loop by default calls ``.eval()`` on the LightningModule before it starts. Override this hook
+        to change the behavior.
+
+        """
+        # The loop won't call this hook unless it is overridden. The line below is here in case the user calls super().
         self.trainer.model.eval()
 
     def on_train_epoch_start(self) -> None:
diff --git a/src/lightning/pytorch/loops/evaluation_loop.py b/src/lightning/pytorch/loops/evaluation_loop.py
index 269beede34205..a337e950d73b1 100644
--- a/src/lightning/pytorch/loops/evaluation_loop.py
+++ b/src/lightning/pytorch/loops/evaluation_loop.py
@@ -41,7 +41,7 @@
 from lightning.pytorch.utilities.combined_loader import CombinedLoader
 from lightning.pytorch.utilities.data import has_len_all_ranks
 from lightning.pytorch.utilities.exceptions import SIGTERMException
-from lightning.pytorch.utilities.model_helpers import is_overridden
+from lightning.pytorch.utilities.model_helpers import _ModuleMode, is_overridden
 from lightning.pytorch.utilities.signature_utils import is_param_in_hook_signature
 
 
@@ -72,6 +72,7 @@ def __init__(
         self._data_fetcher: Optional[_DataFetcher] = None
         self._seen_batches_per_dataloader: DefaultDict[int, int] = defaultdict(int)
         self._last_val_dl_reload_epoch = float("-inf")
+        self._module_mode = _ModuleMode()
 
     @property
     def num_dataloaders(self) -> int:
@@ -293,13 +294,21 @@ def _on_evaluation_model_eval(self) -> None:
         """Sets model to eval mode."""
         trainer = self.trainer
         hook_name = "on_test_model_eval" if trainer.testing else "on_validation_model_eval"
-        call._call_lightning_module_hook(trainer, hook_name)
+        if is_overridden(hook_name, trainer.lightning_module):
+            call._call_lightning_module_hook(trainer, hook_name)
+        else:
+            self._module_mode.capture(trainer.lightning_module)
+            assert trainer.model is not None
+            trainer.model.eval()
 
     def _on_evaluation_model_train(self) -> None:
-        """Sets model to train mode."""
+        """Undoes the eval mode."""
         trainer = self.trainer
         hook_name = "on_test_model_train" if trainer.testing else "on_validation_model_train"
-        call._call_lightning_module_hook(trainer, hook_name)
+        if is_overridden(hook_name, trainer.lightning_module):
+            call._call_lightning_module_hook(trainer, hook_name)
+        else:
+            self._module_mode.restore(trainer.lightning_module)
 
     def _on_evaluation_end(self, *args: Any, **kwargs: Any) -> None:
         """Runs ``on_{validation/test}_end`` hook."""
diff --git a/src/lightning/pytorch/loops/fit_loop.py b/src/lightning/pytorch/loops/fit_loop.py
index db9345f1ab306..e02fe70c4d72b 100644
--- a/src/lightning/pytorch/loops/fit_loop.py
+++ b/src/lightning/pytorch/loops/fit_loop.py
@@ -302,7 +302,6 @@ def setup_data(self) -> None:
     def reset(self) -> None:
         """Resets the internal state of this loop."""
         assert self.trainer.model is not None
-        self.trainer.model.train()
         torch.set_grad_enabled(True)
 
         if self.restarting:
diff --git a/src/lightning/pytorch/loops/prediction_loop.py b/src/lightning/pytorch/loops/prediction_loop.py
index b4c676024a355..5289e1b45c33c 100644
--- a/src/lightning/pytorch/loops/prediction_loop.py
+++ b/src/lightning/pytorch/loops/prediction_loop.py
@@ -38,6 +38,7 @@
 from lightning.pytorch.utilities.combined_loader import CombinedLoader
 from lightning.pytorch.utilities.data import has_len_all_ranks
 from lightning.pytorch.utilities.exceptions import MisconfigurationException
+from lightning.pytorch.utilities.model_helpers import _ModuleMode, is_overridden
 from lightning.pytorch.utilities.signature_utils import is_param_in_hook_signature
 from lightning.pytorch.utilities.types import _PREDICT_OUTPUT
 
@@ -61,6 +62,7 @@ def __init__(self, trainer: "pl.Trainer", inference_mode: bool = True) -> None:
         self._results = None  # for `trainer._results` access
         self._predictions: List[List[Any]] = []  # dataloaders x batches
         self._return_predictions = False
+        self._module_mode = _ModuleMode()
 
     @property
     def return_predictions(self) -> bool:
@@ -191,7 +193,7 @@ def reset(self) -> None:
     def on_run_start(self) -> None:
         """Calls ``_on_predict_model_eval``, ``_on_predict_start`` and ``_on_predict_epoch_start`` hooks."""
         self._verify_dataloader_idx_requirement()
-        call._call_lightning_module_hook(self.trainer, "on_predict_model_eval")
+        self._on_predict_model_eval()
         self._on_predict_start()
         self._on_predict_epoch_start()
 
@@ -199,6 +201,7 @@ def on_run_end(self) -> Optional[_PREDICT_OUTPUT]:
         """Calls ``on_predict_epoch_end`` and ``on_predict_end`` hooks and returns results from all dataloaders."""
         results = self._on_predict_epoch_end()
         self._on_predict_end()
+        self._on_predict_model_train()
         return results
 
     def teardown(self) -> None:
@@ -340,6 +343,18 @@ def _on_predict_start(self) -> None:
         call._call_lightning_module_hook(trainer, "on_predict_start")
         call._call_strategy_hook(trainer, "on_predict_start")
 
+    def _on_predict_model_eval(self) -> None:
+        trainer = self.trainer
+        if is_overridden("on_predict_model_eval", trainer.lightning_module):
+            call._call_lightning_module_hook(self.trainer, "on_predict_model_eval")
+        else:
+            self._module_mode.capture(trainer.lightning_module)
+            assert trainer.model is not None
+            trainer.model.eval()
+
+    def _on_predict_model_train(self) -> None:
+        self._module_mode.restore(self.trainer.lightning_module)
+
     def _on_predict_epoch_start(self) -> None:
         """Calls ``on_predict_epoch_start`` hooks."""
         trainer = self.trainer
diff --git a/src/lightning/pytorch/utilities/model_helpers.py b/src/lightning/pytorch/utilities/model_helpers.py
index 04b60db6d467b..bbbf174b46df8 100644
--- a/src/lightning/pytorch/utilities/model_helpers.py
+++ b/src/lightning/pytorch/utilities/model_helpers.py
@@ -14,7 +14,7 @@
 import inspect
 import os
 from types import MethodType
-from typing import TYPE_CHECKING, Any, Callable, Generic, Optional, Type, TypeVar
+from typing import TYPE_CHECKING, Any, Callable, Dict, Generic, Optional, Type, TypeVar
 
 from lightning_utilities.core.imports import RequirementCache
 from torch import nn
@@ -58,6 +58,24 @@ def get_torchvision_model(model_name: str, **kwargs: Any) -> nn.Module:
     return getattr(models, model_name)(**kwargs)
 
 
+class _ModuleMode:
+    """Captures the ``nn.Module.training`` (bool) mode of every submodule, and allows it to be restored later on."""
+
+    def __init__(self) -> None:
+        self.mode: Dict[str, bool] = {}
+
+    def capture(self, module: nn.Module) -> None:
+        for name, mod in module.named_modules():
+            self.mode[name] = mod.training
+
+    def restore(self, module: nn.Module) -> None:
+        for name, mod in module.named_modules():
+            if name not in self.mode:
+                # TODO: warning cache?
+                continue
+            mod.training = self.mode[name]
+
+
 def _check_mixed_imports(instance: object) -> None:
     old, new = "pytorch_" + "lightning", "lightning." + "pytorch"
     klass = type(instance)
diff --git a/src/lightning/pytorch/utilities/model_summary/model_summary.py b/src/lightning/pytorch/utilities/model_summary/model_summary.py
index 0ff060a0a8679..0e3c07317ad6d 100644
--- a/src/lightning/pytorch/utilities/model_summary/model_summary.py
+++ b/src/lightning/pytorch/utilities/model_summary/model_summary.py
@@ -26,6 +26,7 @@
 
 import lightning.pytorch as pl
 from lightning.fabric.utilities.imports import _TORCH_GREATER_EQUAL_2_0
+from lightning.pytorch.utilities.model_helpers import _ModuleMode
 from lightning.pytorch.utilities.rank_zero import WarningCache
 
 log = logging.getLogger(__name__)
@@ -288,7 +289,8 @@ def _forward_example_input(self) -> None:
         input_ = model._on_before_batch_transfer(input_)
         input_ = model._apply_batch_transfer_handler(input_)
 
-        mode = model.training
+        mode = _ModuleMode()
+        mode.capture(model)
         model.eval()
 
         forward_context = contextlib.nullcontext() if trainer is None else trainer.precision_plugin.forward_context()
@@ -300,7 +302,7 @@ def _forward_example_input(self) -> None:
                 model(**input_)
             else:
                 model(input_)
-        model.train(mode)  # restore mode of module
+        mode.restore(model)
 
     def _get_summary_data(self) -> List[Tuple[str, List[str]]]:
         """Makes a summary listing with:
diff --git a/tests/tests_pytorch/models/test_hooks.py b/tests/tests_pytorch/models/test_hooks.py
index 8b5836cd4d574..d06cadc86f52f 100644
--- a/tests/tests_pytorch/models/test_hooks.py
+++ b/tests/tests_pytorch/models/test_hooks.py
@@ -470,18 +470,14 @@ def training_step(self, batch, batch_idx):
         {"name": "Callback.on_sanity_check_start", "args": (trainer, model)},
         {"name": "val_dataloader"},
         {"name": "train", "args": (False,)},
-        {"name": "on_validation_model_eval"},
         {"name": "Callback.on_validation_start", "args": (trainer, model)},
         {"name": "on_validation_start"},
         *model._eval_epoch("validation", trainer, model, val_batches, "x", device=device),
         {"name": "Callback.on_validation_end", "args": (trainer, model)},
         {"name": "on_validation_end"},
-        {"name": "train", "args": (True,)},
-        {"name": "on_validation_model_train"},
         {"name": "Callback.on_sanity_check_end", "args": (trainer, model)},
         {"name": "train_dataloader"},
         # duplicate `train` because `_run_stage` calls it again in case validation wasn't run
-        {"name": "train", "args": (True,)},
         {"name": "Callback.on_train_start", "args": (trainer, model)},
         {"name": "on_train_start"},
         {"name": "Callback.on_train_epoch_start", "args": (trainer, model)},
@@ -490,14 +486,11 @@ def training_step(self, batch, batch_idx):
         {"name": "zero_grad", **({} if _TORCH_GREATER_EQUAL_2_0 else {"kwargs": {"set_to_none": True}})},
         {"name": "on_validation_model_zero_grad"},
         {"name": "train", "args": (False,)},
-        {"name": "on_validation_model_eval"},
         {"name": "Callback.on_validation_start", "args": (trainer, model)},
         {"name": "on_validation_start"},
         *model._eval_epoch("validation", trainer, model, val_batches, "x", device=device),
         {"name": "Callback.on_validation_end", "args": (trainer, model)},
         {"name": "on_validation_end"},
-        {"name": "train", "args": (True,)},
-        {"name": "on_validation_model_train"},
         {"name": "Callback.on_train_epoch_end", "args": (trainer, model)},
         {"name": "on_train_epoch_end"},  # before ModelCheckpoint because it's a "monitoring callback"
         # `ModelCheckpoint.save_checkpoint` is called here
@@ -570,7 +563,6 @@ def test_trainer_model_hook_system_fit_no_val_and_resume_max_epochs(tmpdir):
         {"name": "on_fit_start"},
         {"name": "zero_grad", **({} if _TORCH_GREATER_EQUAL_2_0 else {"kwargs": {"set_to_none": True}})},
         {"name": "train_dataloader"},
-        {"name": "train", "args": (True,)},
         {"name": "Callback.on_train_start", "args": (trainer, model)},
         {"name": "on_train_start"},
         {"name": "Callback.on_train_epoch_start", "args": (trainer, model)},
@@ -649,7 +641,6 @@ def test_trainer_model_hook_system_fit_no_val_and_resume_max_steps(tmpdir):
         {"name": "on_fit_start"},
         {"name": "zero_grad", **({} if _TORCH_GREATER_EQUAL_2_0 else {"kwargs": {"set_to_none": True}})},
         {"name": "train_dataloader"},
-        {"name": "train", "args": (True,)},
         {"name": "Callback.on_train_start", "args": (trainer, model)},
         {"name": "on_train_start"},
         {"name": "Callback.on_train_epoch_start", "args": (trainer, model)},
@@ -693,14 +684,11 @@ def test_trainer_model_hook_system_eval(tmpdir, batches, verb, noun, dataloader,
     hooks = [
         {"name": f"{dataloader}_dataloader"},
         {"name": "train", "args": (False,)},
-        {"name": f"on_{noun}_model_eval"},
         {"name": f"Callback.on_{noun}_start", "args": (trainer, model)},
         {"name": f"on_{noun}_start"},
         *model._eval_epoch(noun, trainer, model, batches, key, trainer.strategy.root_device),
         {"name": f"Callback.on_{noun}_end", "args": (trainer, model)},
         {"name": f"on_{noun}_end"},
-        {"name": "train", "args": (True,)},
-        {"name": f"on_{noun}_model_train"},
     ]
     expected = [
         {"name": "configure_callbacks"},
@@ -734,7 +722,6 @@ def test_trainer_model_hook_system_predict(tmpdir):
         {"name": "zero_grad", **({} if _TORCH_GREATER_EQUAL_2_0 else {"kwargs": {"set_to_none": True}})},
         {"name": "predict_dataloader"},
         {"name": "train", "args": (False,)},
-        {"name": "on_predict_model_eval"},
         {"name": "Callback.on_predict_start", "args": (trainer, model)},
         {"name": "on_predict_start"},
         {"name": "Callback.on_predict_epoch_start", "args": (trainer, model)},
@@ -896,3 +883,61 @@ def state_dict(self):
     _ = CustomHookedDataModule.load_from_checkpoint(ckpt_path, called=ldm_called)
     assert lm_called == [{"name": "on_load_checkpoint", "args": ({**saved_ckpt, "hyper_parameters": ANY},)}]
     assert ldm_called == [{"name": "load_state_dict", "args": (saved_ckpt[datamodule_state_dict_key],)}]
+
+
+def test_train_eval_mode_restored(tmp_path):
+    """Test that the trainer restores the `training` mode of all submodules to what it was before entering the loop."""
+
+    class MixedTrainModeModule(BoringModel):
+        def __init__(self):
+            super().__init__()
+            # A frozen submodule should keep its mode, regardless of whether we're training or not
+            self.frozen = torch.nn.Linear(2, 2)
+            self.frozen.eval()
+            self.frozen.requires_grad_(False)
+
+        def training_step(self, *args, **kwargs):
+            assert self.layer.weight.requires_grad
+            assert self.layer.training
+            assert not self.frozen.training
+            assert not self.frozen.weight.requires_grad
+            return super().training_step(*args, **kwargs)
+
+        def validation_step(self, *args, **kwargs):
+            assert self.layer.weight.requires_grad
+            assert not self.layer.training
+            assert not self.frozen.training
+            assert not self.frozen.weight.requires_grad
+            return super().validation_step(*args, **kwargs)
+
+        def test_step(self, *args, **kwargs):
+            assert self.layer.weight.requires_grad
+            assert not self.layer.training
+            assert not self.frozen.training
+            assert not self.frozen.weight.requires_grad
+            return super().test_step(*args, **kwargs)
+
+        def predict_step(self, *args, **kwargs):
+            assert self.layer.weight.requires_grad
+            assert not self.layer.training
+            assert not self.frozen.training
+            assert not self.frozen.weight.requires_grad
+            return super().predict_step(*args, **kwargs)
+
+    model = MixedTrainModeModule()
+    trainer = Trainer(
+        default_root_dir=tmp_path,
+        max_epochs=1,
+        val_check_interval=1,
+        limit_train_batches=3,
+        limit_val_batches=2,
+        limit_test_batches=2,
+        limit_predict_batches=2,
+        enable_progress_bar=False,
+        enable_model_summary=False,
+        enable_checkpointing=False,
+    )
+    trainer.fit(model)
+    trainer.validate(model)
+    trainer.test(model)
+    trainer.predict(model)
diff --git a/tests/tests_pytorch/utilities/test_model_helpers.py b/tests/tests_pytorch/utilities/test_model_helpers.py
index c742538bcbfbe..a7f8d83f36fc5 100644
--- a/tests/tests_pytorch/utilities/test_model_helpers.py
+++ b/tests/tests_pytorch/utilities/test_model_helpers.py
@@ -12,9 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import pytest
+import torch.nn
 from lightning.pytorch import LightningDataModule
 from lightning.pytorch.demos.boring_classes import BoringDataModule, BoringModel
-from lightning.pytorch.utilities.model_helpers import _restricted_classmethod, is_overridden
+from lightning.pytorch.utilities.model_helpers import _ModuleMode, _restricted_classmethod, is_overridden
 from lightning_utilities import module_available
 
 
@@ -67,3 +68,65 @@ def test_restricted_classmethod():
         RestrictedClass().restricted_cmethod()
 
     RestrictedClass.restricted_cmethod()  # no exception
+
+
+def test_module_mode():
+    class ChildChildModule(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.layer = torch.nn.Linear(2, 2)
+            self.dropout = torch.nn.Dropout()
+
+    class ChildModule(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.child = ChildChildModule()
+            self.dropout = torch.nn.Dropout()
+
+    class RootModule(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.child1 = ChildModule()
+            self.child2 = ChildModule()
+            self.norm = torch.nn.BatchNorm1d(2)
+
+    # Model with all submodules in the same mode
+    model = RootModule()
+    model.train()
+    mode = _ModuleMode()
+    mode.capture(model)
+    model.eval()
+    assert all(not m.training for m in model.modules())
+    mode.restore(model)
+    assert model.training
+    assert all(m.training for m in model.modules())
+    model.eval()
+
+    mode = _ModuleMode()
+    mode.capture(model)
+    model.eval()
+    assert all(not m.training for m in model.modules())
+    mode.restore(model)
+    assert all(not m.training for m in model.modules())
+    model.train()
+
+    # Model with submodules in different modes
+    model.norm.eval()
+    model.child1.eval()
+    model.child2.train()
+    model.child2.child.eval()
+    model.child2.child.layer.train()
+
+    mode = _ModuleMode()
+    mode.capture(model)
+    model.eval()
+    assert all(not m.training for m in model.modules())
+    mode.restore(model)
+    assert model.training
+    assert not model.norm.training
+    assert all(not m.training for m in model.child1.modules())
+    assert model.child2.training
+    assert model.child2.dropout.training
+    assert not model.child2.child.training
+    assert model.child2.child.layer.training
+    assert not model.child2.child.dropout.training


Code review comments:
diff_hunk: @@ -293,13 +294,21 @@ def _on_evaluation_model_eval(self) -> None:
         """Sets model to eval mode."""
         trainer = self.trainer
         hook_name = "on_test_model_eval" if trainer.testing else "on_validation_model_eval"
-        call._call_lightning_module_hook(trainer, hook_name)
+        if is_overridden(hook_name, trainer.lightning_module):
+            call._call_lightning_module_hook(trainer, hook_name)
+        else:
+            self._module_mode.capture(trainer.lightning_module)
+            assert trainer.model is not None
+            trainer.model.eval()
path: src/lightning/pytorch/loops/evaluation_loop.py
text: @user1:
Why don't we still call the hook unconditionally but capture the state before?

```suggestion
        self._module_mode.capture(trainer.lightning_module)
        call._call_lightning_module_hook(trainer, hook_name)
```