extract optimizer loop

[awaelchli]

What does this PR do?

Part of #7938

Extract an OptimizerLoop from the existing TrainingBatchLoop. This cleanly separates the code path for automatic and manual optimization. Automatic optimization now completely lives inside the OptimizerLoop.

Previously:

(details left out)

if lightning_module.automatic_optimization:
    for opt_idx, opt in enumerate(trainer.optimizers):
        self._run_optimization(batch, opt_idx, opt)
else:
     self._run_optimization(batch, 0, None)

The run_optimization had to handle both code paths for automatic and manual optimization, making it hard to follow, debug, and extend.

Now:

(details left out)

if lightning_module.automatic_optimization:
    self.optimizer.run(batch, optimizers)
else:
     self._run_manual_optimization(batch)

The code path branches here cleanly. The new OptimizerLoop handles automatic optimization completely and the manual optimization path is kept isolated. This greatly improves readability, debugging, and extensibility for loop customization.

A follow-up PR will remove some remaining traces of mixed interfaces, we are holding it out from this PR to make it more manageable to review.

Blocked by:

• remove calls to internal dev debugger in training- and eval loop #9188
• move block_ddp_sync_behaviour to utilities #9192
• move TrainingBatchLoop.build_kwargs to utilities #9198

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 🙃

[codecov]

Codecov Report

Merging #9191 (8475526) into master (35876bb) will increase coverage by 0%.
The diff coverage is 99%.

@@          Coverage Diff           @@
##           master   #9191   +/-   ##
======================================
  Coverage      92%     92%           
======================================
  Files         176     178    +2     
  Lines       14801   14871   +70     
======================================
+ Hits        13653   13723   +70     
  Misses       1148    1148           

[SeanNaren]

LGTM, great work!

[tchaton]

Great work !