Some `LightningModule`/`LightningDataModule` hooks are not profiled

[rohitgr7]

🐛 Bug

These hooks are not profiled, which I think should be profiled:

'on_before_batch_transfer'
'transfer_batch_to_device'
'on_after_batch_transfer'
'configure_gradient_clipping'

To Reproduce

Expected behavior

Environment

• Lightning Component (e.g. Trainer, LightningModule, LightningApp, LightningWork, LightningFlow): LightningModule
• PyTorch Lightning Version (e.g., 1.5.0): master
• Lightning App Version (e.g., 0.5.2):
• PyTorch Version (e.g., 1.10):
• Python version (e.g., 3.9):
• OS (e.g., Linux):
• CUDA/cuDNN version:
• GPU models and configuration:
• How you installed PyTorch (conda, pip, source):
• If compiling from source, the output of torch.__config__.show():
• Running environment of LightningApp (e.g. local, cloud):
• Any other relevant information:

Additional context

cc @carmocca @kaushikb11 @ninginthecloud @rohitgr7 @nbcsm @guotuofeng

[carmocca]

'on_before_batch_transfer'
'transfer_batch_to_device'
'on_after_batch_transfer'

These should be easy to profile in https://github.com/Lightning-AI/lightning/blob/91dd6a68fb596d45914fc5d4fbbf2bad52e8399e/src/pytorch_lightning/core/module.py#L298-L303 as long as the _trainer is available.

'lr_scheduler_step'

I believe this one is profiled through this call: https://github.com/Lightning-AI/lightning/blob/91dd6a68fb596d45914fc5d4fbbf2bad52e8399e/src/pytorch_lightning/loops/epoch/training_epoch_loop.py#L482-L487

'backward'

Same, here: https://github.com/Lightning-AI/lightning/blob/91dd6a68fb596d45914fc5d4fbbf2bad52e8399e/src/pytorch_lightning/loops/optimization/optimizer_loop.py#L304

'configure_gradient_clipping'

This one could be added here: https://github.com/Lightning-AI/lightning/blob/91dd6a68fb596d45914fc5d4fbbf2bad52e8399e/src/pytorch_lightning/plugins/precision/precision_plugin.py#L185 by using trainer._call_lightning_module_hook

on_load_checkpoint

It should be getting profiled:

src/pytorch_lightning/trainer/connectors/checkpoint_connector.py:156:            self.trainer._call_lightning_datamodule_hook("on_load_checkpoint", self._loaded_checkpoint)
src/pytorch_lightning/trainer/connectors/checkpoint_connector.py:174:        self.trainer._call_lightning_module_hook("on_load_checkpoint", self._loaded_checkpoint)
src/pytorch_lightning/trainer/connectors/checkpoint_connector.py:255:        self.trainer._call_callbacks_on_load_checkpoint(self._loaded_checkpoint)

Why are these hooks not profiled?

I believe they all are profiled, but through the Strategy since it's the component that manages their execution, for example: https://github.com/Lightning-AI/lightning/blob/91dd6a68fb596d45914fc5d4fbbf2bad52e8399e/src/pytorch_lightning/loops/optimization/optimizer_loop.py#L407

[rohitgr7]

Thank you :)

Looks like I forgot to look for strategy calls and some edge cases that are hook specific.
Now we have just 4.