3/n Simplify spawn plugins: Merge pre_dispatch and setup logic

[awaelchli]

What does this PR do?

Part of #10987
Merge after #11149

The concept of "dispatch" no longer exists after #10896. This clean up merges two TTP setup calls that no longer need to be separated. The only minor inconvenience is that the on_fit_start call is inbetween the two, but on another note after #10896 the on_fit_start hook is no longer constraint to be exactly at this place.

Before:

# trainer.py, _run():

...

self.training_type_plugin.setup(self)  # WE WANT TO MERGE THIS ...
call_hook("on_fit_start")
self.training_type_plugin.pre_dispatch()  #  ... AND THIS

...

After:

# trainer.py, _run():

...

self.training_type_plugin.setup(self)  # PRE_DISPATCH BECOMES JUST SETUP()
call_hook("on_fit_start")
...

Follow-up work will include:

• Renaming restore_checkpoint_after_pre_dispatch
• Move the _log_hyperparams() call to earlier (can now be done after 2/n Simplify spawn plugins: Spawn immediately #10896)

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

Yes, anyone who implemented pre_dispatch() will have to move their code to setup().

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 update the CHANGELOG? (not for typos, docs, test updates, or internal minor changes/refactorings)

PR review

Anyone in the community is free to review the PR once the tests have passed.
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?

I made sure I had fun coding 🙃

Part of #1 (it's a lie, this is just here to avoid noisy GitHub bot)

cc @tchaton @justusschock @awaelchli @Borda

[four4fish]

Note: this will be a breaking change, I saw few use cases user overrides pre_dispatch

[ananthsub]

On naming, users might be confused with the strategy's setup running distinctly later than the lightning module, data module, and callbacks setup. It'll be especially confusing when we saysetup is the last hook that runs before checkpointing loading, but then needing to qualify which classes's setups were referring to

Do you think there's another name which captures what's being done in the strategy class here?

[four4fish]

On naming, users might be confused with the strategy's setup running distinctly later than the lightning module, data module, and callbacks setup. It'll be especially confusing when we saysetup is the last hook that runs before checkpointing loading, but then needing to qualify which classes's setups were referring to

Do you think there's another name which captures what's being done in the strategy class here?

@ananthsub that's a really good point. We have setup_distributed and setup_envirment to be the setup usually ppl think about. This setup() mainly related to set up optimizer, precision_plugin, and movement logic. I think training_setup()?

[tchaton]

LGTM !