[trainer] solve "scheduler before optimizer step" warning
Author: stas00
Pull request: yes

As discussed in https://github.com/huggingface/transformers/issues/11106 fp16 scaler leads to a warning:

> torch/optim/lr_scheduler.py:132: UserWarning: Detected call of lr_scheduler.step() before optimizer.step(). In PyTorch 1.1.0 and later, you should call them in the opposite order: optimizer.step() before lr_scheduler.step(). Failure to do this will result in PyTorch skipping the first value of the learning rate schedule. See more details at https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate
> warnings.warn("Detected call of lr_scheduler.step() before optimizer.step().

because the optimizer may get skipped until the right scale is found, so we shouldn't run `lr_scheduler.step()` when that happens.

@ptrblck provided a workaround here:
https://discuss.pytorch.org/t/model-weights-not-getting-updated-when-using-autocast/117286/10?u=ptrblck

This is also reported at pytoch: https://github.com/pytorch/pytorch/issues/55585

So the solution is we check the scale before and after and if it changed, then the optimizer wasn't run and we skip the scheduler step then.

Fixes: https://github.com/huggingface/transformers/issues/11106

@sgugger 

Comments:
I'm a bit torn about this solution: it solves the problem exposed by the warning but it creates another problem with no warning, which I had flagged in #11106 as 

> If we somehow manage to catch those skipped optimizer steps and delay the scheduler steps, then we won't respect the number of steps in the scheduler, leading to some wrong end learning rates.

I have no idea if it's better to skip the beginning or end values for the learning rate though.

Also the test is wrong I think it should be
```
optimizer_was_run = scale_before <= scale_after
```
as the scale factor can be multiplied by the `growth_factor` (after a long period without decrease) without skipping a step. It's when it decreases that we know the step was skipped.

> I'm a bit torn about this solution: it solves the problem exposed by the warning but it creates another problem with no warning, which I had flagged in #11106 as
> 
> > If we somehow manage to catch those skipped optimizer steps and delay the scheduler steps, then we won't respect the number of steps in the scheduler, leading to some wrong end learning rates.

If the scheduler and the optimizer are now synchronized why would this happen?

I think it's the external step counter that is out of sync, so we are off at the total number of steps the Trainer does and the optimizer/scheduler see - so the end may be cut off as some "promised steps" won't be seen by the scheduler.

deepspeed already runs `scheduler.step()` only if `optimizer.step()` was run so it's in the same boat.

> I have no idea if it's better to skip the beginning or end values for the learning rate though.

I'd say that potentially cutting of the end is safer.

Also the optimizer stepping could be skipped in the middle of the run as well.

> ```
> optimizer_was_run = scale_before <= scale_after
> ```

Fixed, thank you for catching this!

> If the scheduler and the optimizer are now synchronized why would this happen?

The scheduler was built with a certain number of total steps (for instance go linearly from 1e-4 to 0 in 500 steps). So by skipping those initial steps, we won't be seeing the last learning rates.

> The scheduler was built with a certain number of total steps (for instance go linearly from 1e-4 to 0 in 500 steps). So by skipping those initial steps, we won't be seeing the last learning rates.

Yes! My apologies I understood this from your earlier comment. I just meant that until this PR the scheduler was not synced with optimizer.

So best to truncate the last learning rates where it's usually fixed, or doesn't quite matter if it's a huge run and if it is cyclical it doesn't matter for sure if I understand the situation correctly. Let's perhaps ask it differently - in what situations do you think this mismatch/missing few last steps would practically matter? 

Thinking more and I think it's actually better to skip at the end, since the "short" part of a scheduler is often the warmup (for instance if we set `warum_steps=50`). So I revert my previous objection and I'm okay with the PR :-)

That's a very good point!

And this will also synchronize with the behavior one gets under deepspeed.

Thank you for this brainstorming, @sgugger!