Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Reorder configure_model #19060

Merged
merged 5 commits into from
Dec 5, 2023
Merged

Reorder configure_model #19060

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
bad9863
Reorder `configure_model`
carmocca Nov 23, 2023
34366f3
Hook order test
carmocca Nov 23, 2023
a7d92a2
Merge branch 'master' into carmocca-patch-1
carmocca Nov 23, 2023
1610be9
Update docs
carmocca Nov 23, 2023
8a2f2fa
Update lightning_module.rst
carmocca Nov 24, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 5 additions & 10 deletions docs/source-pytorch/common/lightning_module.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1065,22 +1065,17 @@ for more information.

.. code-block:: python

# runs on every device: devices can be GPUs, TPUs, ...
def fit(self):
if global_rank == 0:
# prepare data is called on GLOBAL_ZERO only
prepare_data()

configure_callbacks()

with parallel(devices):
# devices can be GPUs, TPUs, ...
train_on_device(model)

if local_rank == 0:
prepare_data()

def train_on_device(model):
# called PER DEVICE
setup("fit")
configure_model()
configure_optimizers()

on_fit_start()

# the sanity check runs here
Expand Down
5 changes: 2 additions & 3 deletions src/lightning/pytorch/trainer/trainer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -947,15 +947,14 @@ def _run(
self.__setup_profiler()

call._call_setup_hook(self) # allow user to setup lightning_module in accelerator environment
log.debug(f"{self.__class__.__name__}: configuring model")
call._call_configure_model(self)

# check if we should delay restoring checkpoint till later
if not self.strategy.restore_checkpoint_after_setup:
log.debug(f"{self.__class__.__name__}: restoring module and callbacks from checkpoint path: {ckpt_path}")
self._checkpoint_connector._restore_modules_and_callbacks(ckpt_path)

log.debug(f"{self.__class__.__name__}: configuring model")
call._call_configure_model(self)

# reset logger connector
self._logger_connector.reset_results()
self._logger_connector.reset_metrics()
Expand Down
4 changes: 2 additions & 2 deletions tests/tests_pytorch/models/test_hooks.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -576,10 +576,10 @@ def test_trainer_model_hook_system_fit_no_val_and_resume_max_epochs(tmpdir):
{"name": "prepare_data"},
{"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},)},
{"name": "configure_model"},
{"name": "configure_optimizers"},
{"name": "Callback.on_fit_start", "args": (trainer, model)},
{"name": "on_fit_start"},
Expand Down Expand Up @@ -654,10 +654,10 @@ def test_trainer_model_hook_system_fit_no_val_and_resume_max_steps(tmpdir):
{"name": "prepare_data"},
{"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},)},
{"name": "configure_model"},
{"name": "configure_optimizers"},
{"name": "Callback.on_fit_start", "args": (trainer, model)},
{"name": "on_fit_start"},
Expand Down