Pass lr_scheduler to Accelerator.prepare

[sgugger]

This PR adds support for passing the learning rate scheduler to Accelerator.prepare. This only concerns learning rate schedulers that are called at each optimizer step, not schedulers that are called at the end of each epoch.

Why?

There are three reasons motivating this change:

• users don't always remember they have to define the scheduler after they have prepared the training dataloader (since its size changes). This way is easier.
• the scheduler step should be synchronized with the optimizer step, so when the step is skipped for the optimizer (in mixed precision), it should also be in the scheduler. With this PR, this is now the case.
• this will unlock more cases where the optimizer step is skipped: having Accelerate automatically handle gradient accumulation

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[muellerzr]

Great work on this! Left a few suggestions for documentation consistency and a copy/paste victim 😉

[muellerzr]

Clever! I had wondered how you'd do this 😄

[LysandreJik]

As said offline, LGTM!