DeepSpeed integration not respecting `--warmup_steps` in multi-gpu setting
Author: cchen-dialpad
Pull request: no

### System Info

- `transformers` version: 4.26.1
- Platform: Linux-5.15.0-1030-gcp-x86_64-with-debian-bullseye-sid
- Python version: 3.7.16
- Huggingface_hub version: 0.13.4
- PyTorch version (GPU?): 1.13.1+cu117 (True)
- Deepspeed version: 0.8.3

### Who can help?

 @stas00 @sgugger

### Information

- [X] The official example scripts
- [ ] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

1. DeepSpeed config file `ds_config_zero2_no_optim.json` (adapted from https://huggingface.co/docs/transformers/main_classes/deepspeed#deepspeed-trainer-integration)
```
{
    "zero_optimization": {
        "stage": 2,
        "allgather_partitions": true,
        "allgather_bucket_size": 2e8,
        "overlap_comm": true,
        "reduce_scatter": true,
        "reduce_bucket_size": "auto",
        "contiguous_gradients": true
    },
    "gradient_accumulation_steps": 1,
    "gradient_clipping": "auto",
    "steps_per_print": 2000,
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "wall_clock_breakdown": false
 }
```
2. DeepSpeed launch the training script with 4 GPUs (the script is copied from https://github.com/huggingface/transformers/blob/main/examples/pytorch/language-modeling/run_clm.py)
```
deepspeed --num_gpus=4 run_clm.py \
    --deepspeed $HOME/projects/ds_config_zero2_no_optim.json \
    --report_to mlflow \
    --evaluation_strategy steps \
    --logging_strategy steps \
    --logging_steps 20 \
    --save_strategy steps \
    --save_steps 10000 \
    --max_steps 2560 \
    --num_train_epochs 2 \
    --learning_rate 0.00006 \
    --warmup_steps 100 \
    --model_name_or_path gpt2-medium \
    --do_train \
    --do_eval \
    --dataset_name wikitext \
    --dataset_config_name wikitext-2-raw-v1 \
    --output_dir ./tmp/gpt2_medium-test-mlflow \
    --overwrite_output_dir \
    --per_device_train_batch_size=1 \
    --per_device_eval_batch_size=1 \
    --max_train_samples=2560 \
    --max_eval_samples=10
```
3. Note that the `--warmup_steps 100` and `--learning_rate 0.00006`, so by default, learning rate should increase linearly to 6e-5 at step 100. But the learning rate curve shows that it took 360 steps, and the slope is not a straight line. 
<img width="1036" alt="image" src="https://user-images.githubusercontent.com/47165889/231864162-12f80df5-2827-4bb9-b706-aac66eae5a47.png">
4. Interestingly, if you deepspeed launch with just a single GPU `--num_gpus=1`, the curve seems correct
<img width="927" alt="image" src="https://user-images.githubusercontent.com/47165889/231865235-3d564e5d-4d60-4c5f-ad75-588e35283789.png">

6. The above model is `gpt2-medium`, but training other models such as `gpt2` (with 2 GPUs) also has a similar behavior. For example, note below that `gpt2` at step-100 has a learning rate about `5.82e-05` :
<img width="945" alt="image" src="https://user-images.githubusercontent.com/47165889/231867661-232a017b-d192-490d-822f-0de8d269ba4d.png">



### Expected behavior

The default scheduler should warmup the learning rate linearly, to the specified max learning rate at the specified `warmup_steps`.

Comments:
Thank you for the detailed report and an easy way to reproduce it, @cchen-dialpad  

I'm very much doubting this has anything to do with deepspeed, since you're not using its LR Scheduler, but the default `get_linear_schedule_with_warmup` in HF Trainer. That is you should see the same behavior w/o using deepspeed.

Now I recommend adding: `--logging_steps 1 --logging_strategy steps` and checking the actual LR reports after each step, rather than the graph which most likely extrapolates.

Can you try with the latest `transformers`? I can't repro your report with it. I lowered the number of warm up steps to 10 and you can see that it gets there in 10 steps:

```
$ deepspeed --num_gpus=4 examples/pytorch/language-modeling/run_clm.py --logging_steps 1 --logging_strategy steps --report_to none --evaluation_strategy steps --save_strategy steps     --save_steps 10000     --max_steps 200     --num_train_epochs 2     --learning_rate 0.00006     --warmup_steps 10     --model_name_or_path gpt2-medium     --do_train     --do_eval     --dataset_name wikitext     --dataset_config_name wikitext-2-raw-v1     --output_dir ./tmp/gpt2_medium-test-mlflow     --overwrite_output_dir     --per_device_train_batch_size=1     --per_device_eval_batch_size=1     --max_train_samples=200     --max_eval_samples=10 --evaluation_strategy no --deepspeed ds.json

{'loss': 3.3841, 'learning_rate': 6e-06, 'epoch': 0.02}                                                                                   
{'loss': 3.1834, 'learning_rate': 1.2e-05, 'epoch': 0.04}                                                                                 
{'loss': 3.0764, 'learning_rate': 1.8e-05, 'epoch': 0.06}                                                                                 
{'loss': 3.1638, 'learning_rate': 2.4e-05, 'epoch': 0.08}                                                                                 
{'loss': 3.0197, 'learning_rate': 3e-05, 'epoch': 0.1}                                                                                    
{'loss': 3.3902, 'learning_rate': 3.6e-05, 'epoch': 0.12}                                                                                 
{'loss': 3.0261, 'learning_rate': 4.2e-05, 'epoch': 0.14}                                                                                 
{'loss': 3.1501, 'learning_rate': 4.8e-05, 'epoch': 0.16}                                                                                 
{'loss': 3.1557, 'learning_rate': 5.4000000000000005e-05, 'epoch': 0.18}                                                                  
{'loss': 3.1952, 'learning_rate': 6e-05, 'epoch': 0.2}  
```

Also fyi if you were to run under fp16, you could get optimizer skipping steps while it's tuning up its mixed precision scale factor, so that could also contribute to taking more steps than requested. But you're not using fp16 so this shouldn't be the cause. Just sharing this as FYI.

@stas00 Ah, it's working properly with transformers `v4.28.0` and `v4.27.0`. Not sure how it was fixed from the release info, but thanks for checking! Closing now.

> Also fyi if you were to run under fp16, you could get optimizer skipping steps while it's tuning up its mixed precision scale factor, so that could also contribute to taking more steps than requested. But you're not using fp16 so this shouldn't be the cause. Just sharing this as FYI.

Thanks for sharing, I believe I will run into this later :)

same issues with version 4.29.dev0, the lr is increasing at the first several steps while configured with a warmup scheduler in deepspeed config.