[DeepSpeed] strange learning rate schedule in linear_schedule_with_warmup
## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->

- `transformers` version: 4.3.2
- Platform: Linux
- Python version: 3.7.3
- PyTorch version (GPU?): 1.7 (yes)
- Tensorflow version (GPU?): N/A
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: yes (DeepSpeed)

### Who can help
@stas00

<!-- Your issue will be replied to more quickly if you can figure out the right person to tag with @
 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.


Models:

- albert, bert, xlm: @LysandreJik
- blenderbot, bart, marian, pegasus, encoderdecoder,  t5: @patrickvonplaten, @patil-suraj
- longformer, reformer, transfoxl, xlnet: @patrickvonplaten
- fsmt: @stas00
- funnel: @sgugger
- gpt2: @patrickvonplaten, @LysandreJik
- rag: @patrickvonplaten, @lhoestq
- tensorflow: @jplu

Library:

- benchmarks: @patrickvonplaten
- deepspeed: @stas00
- ray/raytune: @richardliaw, @amogkam
- text generation: @patrickvonplaten
- tokenizers: @n1t0, @LysandreJik
- trainer: @sgugger
- pipelines: @LysandreJik

Documentation: @sgugger

HF projects:

- nlp datasets: [different repo](https://github.com/huggingface/nlp)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Examples:

- maintained examples (not research project or legacy): @sgugger, @patil-suraj
- research_projects/bert-loses-patience: @JetRunner
- research_projects/distillation: @VictorSanh

 -->

## Information

Model I am using (Bert, XLNet ...): GPT-2

The problem arises when using:
* [x] the official example scripts: (give details below)
* [ ] my own modified scripts: (give details below)

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [ ] my own task or dataset: (give details below)

## To reproduce

I am trying using deepspeed for run_clm.py to train GPT-2 (from scratch).
I want to use the same scheduler (`linear_schedule_with_warmup`) and `optimizer` as ones used in run_clm.py.
So, the `scheduler` and `optimizer` sections are removed in `examples/tests/deepspeed/ds_config.json`,
and the original ones are used.

My `ds_config.json` is as follows:

```
{
    "fp16": {
        "enabled": true,
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "hysteresis": 2,
        "min_loss_scale": 1
    },

   "zero_optimization": {
       "stage": 2,
       "allgather_partitions": true,
       "allgather_bucket_size": 2e8,
       "overlap_comm": true,
       "reduce_scatter": true,
       "reduce_bucket_size": 2e8,
       "contiguous_gradients": true,
       "cpu_offload": true
   },

   "zero_allow_untested_optimizer": true,

    "steps_per_print": 2000,
    "wall_clock_breakdown": false
}
```

I ran the following command (using 4GPUs in one node):

$ cd examples/language-modeling/
$ deepspeed run_clm.py \
   --output_dir=/somewhere \
   --model_type=gpt2 \
   --do_train \
   --dataset_name wikitext \
   --dataset_config_name wikitext-2-raw-v1 \
   --tokenizer_name gpt2 \
   --block_size=512 \
   --num_train_epochs=5 \
   --warmup_steps=100 \
   --learning_rate=2e-5 \
   --per_device_train_batch_size=32 \
   --per_device_eval_batch_size=32 \
   --save_steps=10000 \
   --save_total_limit=5 \
   --dataloader_drop_last \
   --deepspeed ds_config.json \
   --logging_steps=10

The learning rate schedule was strange. The following is a screenshot of tensorboard. 
 
![image](https://user-images.githubusercontent.com/16042472/108629197-2fc69700-74a2-11eb-9d05-c8efd476c489.png)

The initial learning rate was 1e-5, which should be 0. The learning rate went up to 2e-5 (it was OK), and went down to 0 around the middle (before the end), which was strange.

I tested a `WarmupDecayLR` scheduler in `deepspeed` (without `transformers`), and it seemed OK.
So, I think the utilization of this scheduler in `transformers` is strange.

<!-- If you have code snippets, error messages, stack traces please provide them here as well.
     Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting
     Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.-->

## Expected behavior

The learning rate schedule through `deepspeed` should be the same as the original one used in `run_clm.py`.

<!-- A clear and concise description of what you would expect to happen. -->

Incidentally a bug fix was just merged as part of: https://github.com/huggingface/transformers/pull/10310
- the scheduler step was getting run twice.

Could you please re-test with `transformers` master?

Thank you!
<|||||>Thank you for your response.

I have tested the latest version, but the following error occurred. There is something wrong in the lr initialization.
```
..
  File "run_clm.py", line 376, in main
    train_result = trainer.train(resume_from_checkpoint=checkpoint)
  File "/home/.../anaconda3/envs/transformers-4.3.2/lib/python3.7/site-packages/transformers/trainer.py", line 1054, in train
    train_result = trainer.train(resume_from_checkpoint=checkpoint)
.. 
  File "/home/.../anaconda3/envs/transformers-4.3.2/lib/python3.7/site-packages/deepspeed/runtime/lr_schedules.py", line 728, in get_last_lr
    assert getattr(self, '_last_lr', None) is not None, "need to call step() first"
AssertionError: need to call step() first
``` <|||||>Thank you for testing with the master version.

Please always post the full backtrace and the full command line you used - otherwise it's impossible to reproduce the problem and know how to fix it.

Thank you.<|||||>Sorry. I just ran the same command shown above, and the full error is as follows:

```
  File "run_clm.py", line 417, in <module>
    main()
  File "run_clm.py", line 376, in main
    train_result = trainer.train(resume_from_checkpoint=checkpoint)
  File "/home/.../anaconda3/envs/transformers-4.3.2/lib/python3.7/site-packages/transformers/trainer.py", line 1054, in train
Traceback (most recent call last):
  File "run_clm.py", line 417, in <module>
Traceback (most recent call last):
  File "run_clm.py", line 417, in <module>
    self._maybe_log_save_evaluate(tr_loss, model, trial, epoch)
  File "/home/.../anaconda3/envs/transformers-4.3.2/lib/python3.7/site-packages/transformers/trainer.py", line 1135, in _maybe_log_save_evaluate
Traceback (most recent call last):
  File "run_clm.py", line 417, in <module>
    main()
  File "run_clm.py", line 376, in main
    if version.parse(torch.__version__) >= version.parse("1.4")
  File "/home/.../anaconda3/envs/transformers-4.3.2/lib/python3.7/site-packages/deepspeed/runtime/lr_schedules.py", line 728, in get_last_lr
    main()
  File "run_clm.py", line 376, in main
        assert getattr(self, '_last_lr', None) is not None, "need to call step() first"
AssertionErrortrain_result = trainer.train(resume_from_checkpoint=checkpoint): need to call step() first

  File "/home/.../anaconda3/envs/transformers-4.3.2/lib/python3.7/site-packages/transformers/trainer.py", line 1054, in train
    main()
  File "run_clm.py", line 376, in main
    train_result = trainer.train(resume_from_checkpoint=checkpoint)
  File "/home/.../anaconda3/envs/transformers-4.3.2/lib/python3.7/site-packages/transformers/trainer.py", line 1054, in train
    train_result = trainer.train(resume_from_checkpoint=checkpoint)
  File "/home/.../anaconda3/envs/transformers-4.3.2/lib/python3.7/site-packages/transformers/trainer.py", line 1054, in train
    self._maybe_log_save_evaluate(tr_loss, model, trial, epoch)
  File "/home/.../anaconda3/envs/transformers-4.3.2/lib/python3.7/site-packages/transformers/trainer.py", line 1135, in _maybe_log_save_evaluate
    self._maybe_log_save_evaluate(tr_loss, model, trial, epoch)
  File "/home/.../anaconda3/envs/transformers-4.3.2/lib/python3.7/site-packages/transformers/trainer.py", line 1135, in _maybe_log_save_evaluate
    self._maybe_log_save_evaluate(tr_loss, model, trial, epoch)
  File "/home/.../anaconda3/envs/transformers-4.3.2/lib/python3.7/site-packages/transformers/trainer.py", line 1135, in _maybe_log_save_evaluate
    if version.parse(torch.__version__) >= version.parse("1.4")
  File "/home/.../anaconda3/envs/transformers-4.3.2/lib/python3.7/site-packages/deepspeed/runtime/lr_schedules.py", line 728, in get_last_lr
    if version.parse(torch.__version__) >= version.parse("1.4")
  File "/home/.../anaconda3/envs/transformers-4.3.2/lib/python3.7/site-packages/deepspeed/runtime/lr_schedules.py", line 728, in get_last_lr
    if version.parse(torch.__version__) >= version.parse("1.4")
  File "/home/.../anaconda3/envs/transformers-4.3.2/lib/python3.7/site-packages/deepspeed/runtime/lr_schedules.py", line 728, in get_last_lr
    assert getattr(self, '_last_lr', None) is not None, "need to call step() first"
AssertionError: need to call step() first
    assert getattr(self, '_last_lr', None) is not None, "need to call step() first"
AssertionError: need to call step() first
    assert getattr(self, '_last_lr', None) is not None, "need to call step() first"
AssertionError: need to call step() first
  1%|▋                                                                             | 10/1095 [00:08<16:13,  1.11it/s]
```

Thanks.<|||||>Great, thank you, I'm able to reproduce this problem. Let me investigate and I will get back to you with a solution. <|||||>I understand the problem.

The optimizer doesn't kick in until a much later step, so `lr_scheduler` doesn't get its first `step()` yet and `._maybe_log_save_evaluate` fails to retrieve `get_last_lr` since there wasn't any yet.

A quick workaround is to add `"initial_scale_power": 1,`, which will force the optimizer to churn from step one.
```
    "fp16": {
        "enabled": true,
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 1,
        "hysteresis": 2,
        "min_loss_scale": 1
    },
```
but it might not be an optimal solution. https://www.deepspeed.ai/docs/config-json/#fp16-training-options

I will think of how to resolve this correctly, but meanwhile please let me know if that resolves the scheduler issue.

to explain - when you use deepspeed's fp16 it skips the optimizer/scheduler calls until the OVERFLOW is no more. And you'd see the following in the log:

```
OVERFLOW! Rank 0 Skipping step. Attempted loss scale: 4294967296, reducing to 4294967296
OVERFLOW! Rank 0 Skipping step. Attempted loss scale: 4294967296, reducing to 2147483648.0
OVERFLOW! Rank 0 Skipping step. Attempted loss scale: 2147483648.0, reducing to 1073741824.0
```

This is also probably why you see an odd behavior as you reported originally (besides the double step bug I fixed)<|||||>OK, this PR should work too if you would like to try it instead: https://github.com/huggingface/transformers/pull/10362
<|||||>Thank you for your work.

I set `"initial_scale_power": 1`, re-ran the command, and the training finished without errors.
 
![image](https://user-images.githubusercontent.com/16042472/108927816-622fea00-7684-11eb-8071-c3f101d21c9c.png)

After the PR #10362 is merged into master, I will try it. Thanks.<|||||>FYI,  it has been merged.<|||||>Thanks.

I have tested the latest version (without setting `"initial_scale_power": 1`), and the learning rate behavior is as expected! 

![image](https://user-images.githubusercontent.com/16042472/108942404-47687000-769a-11eb-82af-a68bf1815345.png)

Thanks for your work. It is very useful to use deepspeed in transformers.<|||||>Thank you for your feedback and supporting this problem fixing process, @tomohideshibata 