update

Author: deependujha

src/lightning/pytorch/callbacks/finetuning.py

@@ -106,6 +106,17 @@ def load_state_dict(self, state_dict: dict[str, Any]) -> None:
 
     @override
     def on_fit_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
+        # freeze the required modules before training
+        self.freeze_before_training(pl_module)
+
+        from lightning.pytorch.strategies import DeepSpeedStrategy
+
+        if isinstance(trainer.strategy, DeepSpeedStrategy):
+            raise NotImplementedError(
+                "The Finetuning callback does not support running with the DeepSpeed strategy."
+                " Choose a different strategy or disable the callback."
+            )
+
         # restore the param_groups created during the previous training.
         if self._restarting:
             named_parameters = dict(pl_module.named_parameters())
@@ -273,18 +284,6 @@ def unfreeze_and_add_param_group(
         if params:
             optimizer.add_param_group({"params": params, "lr": params_lr / denom_lr})
 
-    @override
-    def setup(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", stage: str) -> None:
-        self.freeze_before_training(pl_module)
-
-        from lightning.pytorch.strategies import DeepSpeedStrategy
-
-        if isinstance(trainer.strategy, DeepSpeedStrategy):
-            raise NotImplementedError(
-                "The Finetuning callback does not support running with the DeepSpeed strategy."
-                " Choose a different strategy or disable the callback."
-            )
-
     @staticmethod
     def _apply_mapping_to_param_groups(param_groups: list[dict[str, Any]], mapping: dict) -> list[dict[str, Any]]:
         output = []

src/lightning/pytorch/trainer/trainer.py

@@ -985,9 +985,9 @@ def _run(
         log.debug(f"{self.__class__.__name__}: preparing data")
         self._data_connector.prepare_data()
 
+        call._call_setup_hook(self)  # allow user to set up LightningModule in accelerator environment
         log.debug(f"{self.__class__.__name__}: configuring model")
         call._call_configure_model(self)
-        call._call_setup_hook(self)  # allow user to set up LightningModule in accelerator environment
 
         # check if we should delay restoring checkpoint till later
         if not self.strategy.restore_checkpoint_after_setup:

tests/tests_pytorch/models/test_hooks.py

@@ -472,11 +472,11 @@ def training_step(self, batch, batch_idx):
     expected = [
         {"name": "configure_callbacks"},
         {"name": "prepare_data"},
-        {"name": "configure_model"},
         {"name": "Callback.setup", "args": (trainer, model), "kwargs": {"stage": "fit"}},
         {"name": "setup", "kwargs": {"stage": "fit"}},
         # DeepSpeed needs the batch size to figure out throughput logging
         *([{"name": "train_dataloader"}] if using_deepspeed else []),
+        {"name": "configure_model"},
         {"name": "configure_optimizers"},
         {"name": "Callback.on_fit_start", "args": (trainer, model)},
         {"name": "on_fit_start"},
@@ -651,9 +651,9 @@ def test_trainer_model_hook_system_fit_no_val_and_resume_max_steps(tmp_path):
     expected = [
         {"name": "configure_callbacks"},
         {"name": "prepare_data"},
-        {"name": "configure_model"},
         {"name": "Callback.setup", "args": (trainer, model), "kwargs": {"stage": "fit"}},
         {"name": "setup", "kwargs": {"stage": "fit"}},
+        {"name": "configure_model"},
         {"name": "on_load_checkpoint", "args": (loaded_ckpt,)},
         {"name": "Callback.on_load_checkpoint", "args": (trainer, model, loaded_ckpt)},
         {"name": "Callback.load_state_dict", "args": ({"foo": True},)},
@@ -719,9 +719,9 @@ def test_trainer_model_hook_system_eval(tmp_path, override_on_x_model_train, bat
     expected = [
         {"name": "configure_callbacks"},
         {"name": "prepare_data"},
-        {"name": "configure_model"},
         {"name": "Callback.setup", "args": (trainer, model), "kwargs": {"stage": verb}},
         {"name": "setup", "kwargs": {"stage": verb}},
+        {"name": "configure_model"},
         {"name": "zero_grad"},
         *(hooks if batches else []),
         {"name": "Callback.teardown", "args": (trainer, model), "kwargs": {"stage": verb}},
@@ -746,9 +746,9 @@ def test_trainer_model_hook_system_predict(tmp_path):
     expected = [
         {"name": "configure_callbacks"},
         {"name": "prepare_data"},
-        {"name": "configure_model"},
         {"name": "Callback.setup", "args": (trainer, model), "kwargs": {"stage": "predict"}},
         {"name": "setup", "kwargs": {"stage": "predict"}},
+        {"name": "configure_model"},
         {"name": "zero_grad"},
         {"name": "predict_dataloader"},
         {"name": "train", "args": (False,)},