lr scheduler before optimizer step warning
Author: stas00
Pull request: no

for a really long time now we have this warning being spewed up on startup:

> 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()`. 

Not sure if it's all the time, but sure thing a lot and it'd be nice to solve it.

I patched lr_scheduler's `step` to print the counters:
``` # torch/optim/lr_scheduler.py:117 
    def step(self, epoch=None):
        print(f"self._step_count={self._step_count}, self.optimizer._step_count={self.optimizer._step_count}")
```

With standard way we get:

```
net = torch.nn.Linear(1,1)
optimizer = torch.optim.SGD(net.parameters(), lr=1)
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=[lambda epoch: epoch // 30])

optimizer.step()
scheduler.step()

optimizer.step()
scheduler.step()

optimizer.step()
scheduler.step()
```
output:
```
self._step_count=0, self.optimizer._step_count=0
self._step_count=1, self.optimizer._step_count=1
self._step_count=2, self.optimizer._step_count=2
self._step_count=3, self.optimizer._step_count=3
```

But when we run some example with the HF Trainer, e.g. with:
```
export BS=1; rm -rf /tmp/test-clm; PYTHONPATH=src USE_TF=0 CUDA_VISIBLE_DEVICES=0 python \
examples/language-modeling/run_clm.py --model_name_or_path distilgpt2 --dataset_name wikitext \
--dataset_config_name wikitext-2-raw-v1 --do_train --max_train_samples 1 \
--per_device_train_batch_size $BS --output_dir /tmp/test-clm --block_size 128 --logging_steps 1 \
--fp16 --block_size=128 
```

we get:

```
self._step_count=0, self.optimizer._step_count=0
self._step_count=1, self.optimizer._step_count=0
self._step_count=2, self.optimizer._step_count=0
self._step_count=3, self.optimizer._step_count=1
```

as you can see the counter is off by 2 - that's why we get the warning. 

The order of the `step()` calls in the code looks ok, I think our optimizers may get reset after the scheduler was created perhaps?

@sgugger 

Comments:
I think that's because of the dynamic loss scaling in FP16, which skips some optimizer steps at the beginning of training if the gradients overflowed (while it's trying to find the best loss scale). Not sure of the best way to treat this:

- 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 rate
- If we don't change anything, then we skip the first values of the learning rate and get that warning.

Since it's impossible to know beforehand how many optimizer steps will be skipped, I don't see how to deal with that and would just go for 2, with some manual update of ` _step_count` to avoid catching the warning maybe?

You nailed it - it's `--fp16`, removing it resolves the problem.

So the problem is here:

```
                    elif self.use_amp:
                        self.scaler.step(self.optimizer)
                        self.scaler.update()
                    [...]
                    self.lr_scheduler.step()
```

This sounds like a pytorch amp issue then. I stepped through with debugger and indeed - the first time it skips optimizer step because of the overflow. 

I will file an issue with pytorch. https://github.com/pytorch/pytorch/issues/55585

Yes I agree it is more on the PyTorch side to do something for that warning. Thanks for reporting it to them!

Solved at https://github.com/huggingface/transformers/pull/11144