Add profiling for on_load_checkpoint/on_save_checkpoint callback and LM hooks

[daniellepintz]

What does this PR do?

For the LM hooks we can simply call _call_lightning_module_hook to get profiling.
For the callback hooks, we have dedicated functions with special logic - I added profiling there.

Does your PR introduce any breaking changes? If yes, please list them.

Before submitting

• Was this discussed/approved via a GitHub issue? (not for typos and docs)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure your PR does only one thing, instead of bundling different changes together?
• Did you make sure to update the documentation with your changes? (if necessary)
• Did you write any new necessary tests? (not for typos and docs)
• Did you verify new and existing tests pass locally with your changes?
• Did you list all the breaking changes introduced by this pull request?
• Did you update the CHANGELOG? (not for typos, docs, test updates, or internal minor changes/refactorings)

PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing make sure you have read Review guidelines. In short, see the following bullet-list:

• Is this pull request ready for review? (if not, please submit in draft mode)
• Check that all items from Before submitting are resolved
• Make sure the title is self-explanatory and the description concisely explains the PR
• Add labels and milestones (and optionally projects) to the PR so it can be classified

Did you have fun?

Make sure you had fun coding 🙃

[akashkw]

Edit: #12149 (comment)

[daniellepintz]

@akashkw Thanks. I think I should be adding on_save_checkpoint and on_load_checkpoint to this dictionary in this test:

https://github.com/PyTorchLightning/pytorch-lightning/blob/b34d8673a9e77d25a826b7ada31d1204dfc818ab/tests/trainer/logging_/test_logger_connector.py#L212-L238

but when I do I still get an error :/ https://gist.github.com/daniellepintz/025a8870a63b40b83c6d042c0831de98

Does anyone have any insight on this?

[krshrimali]

@daniellepintz - Probably, you want to change the following 2 lines:

"on_save_checkpoint": None,
"on_load_checkpoint": None,

to:

"on_save_checkpoint": _LogOptions(
    allowed_on_step=(False, True), allowed_on_epoch=(False, True), default_on_step=True, default_on_epoch=False
),
"on_load_checkpoint": _LogOptions(
    allowed_on_step=(False, True), allowed_on_epoch=(False, True), default_on_step=True, default_on_epoch=False
),

[daniellepintz]

@krshrimali thanks! Just want to double check if we want to change the behavior and allow logging inside these hooks, cc @carmocca

[carmocca]

if we want to change the behavior and allow logging inside these hooks

Please don't. I'll take a look at the current failing tests.