Forget to put the global-step in lr scheduler in train.py
Author: SuperCB
Pull request: no

### System Info

- ` transformers` version: 4.30.2 
- Platform: Linux-3.10.0-1160.83.1.el7.x86_64-x86_64-with-glibc2.35
- Python version: 3.10.6  
- Huggingface_hub version: 0.15.1 
- Safetensors version: 0.3.1  
- PyTorch version (GPU?): 2.0.0 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed 
- JaxLib version: not installed
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

### Who can help?

@sgugger 

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

We fine-tuned the llama model using trainer.py and accelerated the training process using DeepSpeed. We chose the cosine function as the type for the lr scheduler. 

### Expected behavior

In the process of training LLM using `train()` in trainer.py, we observed that after reloading the model from a checkpoint with deepspeed zero-3, the learning rate scheduler did not resume iteration from the previous training's learning rate. Instead, there were gaps in the learning rate progression.We used a lr scheduler of the cosine type, which is not provided by Deepspeed. Instead, we observed that the trainer used the [lr_scheduler](https://github.com/pytorch/pytorch/blob/main/torch/optim/lr_scheduler.py#L124) defined in torch.


Comments:
cc @pacman100 

Hello, a minimal reproducer for this please. The checkpointing test loads back the correct lr: https://github.com/huggingface/accelerate/actions/runs/5651936442/job/15310803049#step:5:2259

> Hello, a minimal reproducer for this please. The checkpointing test loads back the correct lr: https://github.com/huggingface/accelerate/actions/runs/5651936442/job/15310803049#step:5:2259

Hello, the difference between your job and ours lies in the **different lr_scheduler type**. We use the cosine lr_scheculer (with `--lr_scheduler_type cosine`) instead of set `"scheduler"` in `deepspeed_config.json`.

What we see in wandb is, the `train/learning_rate` always restart from zero with warmups when we resume training from checkpoints.
![image](https://github.com/huggingface/transformers/assets/9065640/e928d3df-7194-456e-9652-cfd405b30aa3)


As we go through the code(correct us if we were wrong), the reason  is that DeepSpeed doesn't support consine lr_scheduler, so Transformers use [pytorch native lr_scheduler](https://github.com/pytorch/pytorch/blob/e18d53e2df44bccd7231cdf3dad6ea1255221bd4/torch/optim/lr_scheduler.py#L124), which maintains a self-incrementing variable `self.last_epoch` if we don't pass `epoch` when calling `step()`.

We will provide a minimal reproducer later.


We used [Stanford Alpaca](https://github.com/tatsu-lab/stanford_alpaca) to reproduce this bug and here is our startup script and DeepSpeed configuration.
**`startup script`**
```
torchrun --nproc_per_node=8 --master_port=29600 train.py \
	--model_name_or_path '/nvme/models/llama-13b/' \
	--data_path ./alpaca_data.json \
	--bf16 True \
	--output_dir './chpt' \
	--num_train_epochs 3 \
	--per_device_train_batch_size 4 \
	--per_device_eval_batch_size 4 \
	--gradient_accumulation_steps 8 \
	--evaluation_strategy "no" \
	--save_strategy "steps" \
	--save_steps 2000 \
	--save_total_limit 1 \
	--learning_rate 2e-5 \
	--weight_decay 0. \
	--warmup_ratio 0.03 \
	--logging_steps 1 \
	--learning_rate 2e-4 \
	--warmup_ratio 0.1 \
	--tf32 True \
	--lr_scheduler_type 'cosine' \
	--adam_beta2 0.95 \
	--deepspeed 'deepspeed.zero3.json' \
	--save_strategy 'steps' \
	--save_steps 10
```
`deepspeed config`
```
{
  "bf16": {
    "enabled": "auto"
  },
  "zero_optimization": {
    "stage": 3,
    "overlap_comm": true,
    "contiguous_gradients": true,
    "sub_group_size": 1e9,
    "reduce_bucket_size": "auto",
    "stage3_prefetch_bucket_size": "auto",
    "stage3_param_persistence_threshold": "auto",
    "stage3_max_live_parameters": 1e9,
    "stage3_max_reuse_distance": 1e9,
    "stage3_gather_16bit_weights_on_model_save": true
  },
  "gradient_accumulation_steps": "auto",
  "gradient_clipping": "auto",
  "steps_per_print": 5,
  "train_batch_size": "auto",
  "train_micro_batch_size_per_gpu": "auto",
  "wall_clock_breakdown": false,
  "flops_profiler": {
    "enabled": true,
    "profile_step": 5,
    "module_depth": -1,
    "top_modules": 1,
    "detailed": true,
    "output_file": "logs/deepspeed_flops.log"
  }
}
```
**before restart**
![20230726-125554](https://github.com/huggingface/transformers/assets/82354186/4045a5c4-0473-479c-80ec-a1941817f0b4)

![image](https://github.com/huggingface/transformers/assets/82354186/1353067f-7d6a-47a2-8d4d-d1c4d0c0a700)
**after restart**
![new](https://github.com/huggingface/transformers/assets/82354186/87d599f2-3ea8-4b7c-aa0d-0a6c7e15fd25)

**wandb**
![image](https://github.com/huggingface/transformers/assets/82354186/ce025b1b-2248-42cc-ab78-037ac8d5114b)


We save a checkpoint every ten steps, and when we restart the training from the previous checkpoint, we observe that the learning rate after the restart is different from the learning rate before the restart. This indicates that the LR (learning rate) scheduler's state is not reset during the restart.


@pacman100 Hello, I apologize for bothering you. Could you please let us know if there have been any updates or if there might be any corrections needed in our description? Much thanks.

Hello @Dounm, please see https://github.com/huggingface/transformers/issues/25865#issuecomment-1700722492 and https://github.com/huggingface/transformers/issues/25865#issuecomment-1702239905. It should be resolving this issue.

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.