Incorrect learning rate when using 'cosine_with_restarts' scheduler type
Author: tijsg
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.20.1
- Platform: Linux-5.4.0-121-generic-x86_64-with-glibc2.29
- Python version: 3.8.10
- Huggingface_hub version: 0.8.1
- PyTorch version (GPU?): 1.12.0+cu116 (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?: yes
- Using distributed or parallel set-up in script?: no

### Who can help?

@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

Start training a Roberta model for masked LM using the following parameters.
All parameters are shared for the sake of context, but the parameter that fails is "lr_scheduler_type".
The learning rate remains fixed at a default value.

```py
training_args = TrainingArguments(
    output_dir='./roberta',
    overwrite_output_dir=True,
    evaluation_strategy = 'steps',
    num_train_epochs=100,
    lr_scheduler_type='cosine_with_restarts',
    warmup_steps=100,
    weight_decay=0.01,
    per_device_train_batch_size=20,
    per_device_eval_batch_size=20,
    save_steps=2048,
    eval_steps=2048,
    save_total_limit=3,
    report_to="wandb",
    ignore_data_skip=True,
    gradient_accumulation_steps=4, 
    gradient_checkpointing=True,
    fp16=True
)

trainer = Trainer(
    model=model,
    args=training_args,
    data_collator=data_collator,
    train_dataset=train_dataset,
    eval_dataset=val_dataset,
    #prediction_loss_only=True,
)
```

### Expected behavior

The learning rate should behave like configured and like documented for 'cosine_with_restarts', instead of remaining fixed through the epochs.

Comments:
Could you explain to us how you are seeing the learning rate being constant? Just tried your code and adding a print statement for the learning rate at every step, and I see something that changes.

I'm using weight and biases at wandb.ai to keep track of my metrics
![image](https://user-images.githubusercontent.com/5707158/182159636-2c728deb-6f23-43eb-836c-dd6ae5041a7b.png)


training_args = TrainingArguments(
    output_dir='./roberta',
    overwrite_output_dir=True,
    evaluation_strategy = 'steps',
    num_train_epochs=100,
    learning_rate=1e-4,
    weight_decay=0.01,
    per_device_train_batch_size=20,
    per_device_eval_batch_size=20,
    save_steps=2048,
    eval_steps=2048,
    save_total_limit=3,
    report_to="wandb",
    ignore_data_skip=True,
    gradient_accumulation_steps=4, 
    gradient_checkpointing=True,
    fp16=True
)

Results in 
![image](https://user-images.githubusercontent.com/5707158/182160682-3c14ecab-57f2-4526-acb4-a24e1d127e7c.png)


So this only shows the learning rates at steps that are multiple of 500 since that is the default logging step I wouldn't trust a curve with 6 points on it.

could you please share how you add print statements for each step?

By changing the source code of the Trainer. YOu can also log more frequently by changing the `loggin_steps` argument.

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.