cosine_with_restarts schedule not restart
Author: haixpham
Pull request: no

### System Info

transformers 4.37.2

### Who can help?

@muellerzr @pacman100 

### Information

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

### Tasks

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

### Reproduction

- This is the official implementation of `cosine_with_restarts` lambda function in `optimization.py` [link](https://github.com/huggingface/transformers/blob/7eb3ba82241c927053689270a0751f4ff5d33c54/src/transformers/optimization.py#L178)
```
def _get_cosine_with_hard_restarts_schedule_with_warmup_lr_lambda(
    current_step: int, *, num_warmup_steps: int, num_training_steps: int, num_cycles: int
):
    if current_step < num_warmup_steps:
        return float(current_step) / float(max(1, num_warmup_steps))
    progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))
    if progress >= 1.0:
        return 0.0
    return max(0.0, 0.5 * (1.0 + math.cos(math.pi * ((float(num_cycles) * progress) % 1.0))))
```
This function will create a single cosine cycle from [0..pi] -> [1, 0]

I modified it as follows:
```
def _get_cosine_with_hard_restarts_schedule_with_warmup_lr_lambda(
    current_step: int, *, num_warmup_steps: int, num_training_steps: int, num_cycles: int
):
    if current_step < num_warmup_steps:
        return float(current_step) / float(max(1, num_warmup_steps))
    # progress over the entire training length
    if current_step >= num_training_steps:
        return 0.0
    if num_cycles < 1:
        num_cycles = 1
    cycle_length = int(math.ceil((num_training_steps - num_warmup_steps) / num_cycles))
    subprogress = float((current_step - num_warmup_steps) % cycle_length) / float(max(1, cycle_length))
    return max(0.0, 0.5 * (1.0 + math.cos(math.pi * subprogress)))
```
and plot the curve to confirm that it's working (1000 steps, 5 restarts)
![lr_schedule](https://github.com/huggingface/transformers/assets/32718796/98057199-e840-40c3-a764-962961b2c3f1)

As you can see from the above plot - the `cosine_with_restarts` schedule is the same as `cosine` schedule, while my modification (green) indeed restarts the schedule.

- Another issue is with logging the learning rate. With the above fix and provide `num_cycles` to TrainingArguments, I debugged and found the right learning rate was calculated, however, the logger outputs a different learning rate curve similar to `cosine` schedule. I haven't found where the learning rate is logged.

### Expected behavior

- the `cosine_with_restarts` schedule should really restart
- The log output of `cosine_with_restarts` should reflect correct value

Comments:
I'm not sure how I feel about suddenly modifying an implementation that's been this way for 5+ years now (even if it may be right), I feel users would be quite shocked by this change. I'd rather see we add another argument potentially that allows for this more aggressive restart schedule?

I found why my fix didn't work: because I was using Deepspeed. At this call: https://github.com/huggingface/transformers/blob/345b9b1a6a308a1fa6559251eb33ead2211240ac/src/transformers/integrations/deepspeed.py#L319
the `scheduler_specific_kwargs=trainer.args.lr_scheduler_kwargs,` line is missing, hence `num_cycles == 1` is always true.
This made me scratch head for a while. In the latest main branch, this call has been replaced with `Trainer.create_scheduler()` to get around the problem https://github.com/huggingface/transformers/blob/cbe58b4269457a6ca66a556224b23f9ef246f905/src/transformers/integrations/deepspeed.py#L348

@muellerzr but the method was really supposed to work like that, as documented here https://huggingface.co/docs/transformers/main_classes/optimizer_schedules#transformers.get_cosine_with_hard_restarts_schedule_with_warmup

Good digging and thanks for the check :) Would you like to open a PR with your proposed fix then? (Unless `main` already solves this a different way)

I have already maintained a forked copy of transformers, frozen at 4.37.2, so I can't fork another version to sync with the latest MAIN. I also don't have permission to push branch to `huggingface/transformers` repo. Can you please make the change and create a PR instead?

Sure but I'm still not fully with you

> the scheduler_specific_kwargs=trainer.args.lr_scheduler_kwargs, line is missing, hence num_cycles == 1 is always true.
This made me scratch head for a while. In the latest main branch, this call has been replaced with Trainer.create_scheduler() to get around the problem 

So, are you saying without deepspeed the problem still exists? And again, it's still a 5 year old behavior, I'm not keen on changing it if this is the first time we've heard this reported. 

Apologies for the lack of clarity. In the previous transformers version 4.37.2, `get_scheduler()` is called directly in `integrations/deepspeed.py`, where `scheduler_specific_kwargs=trainer.args.lr_scheduler_kwargs` is missing. In the current master, the problem with deepspeed/learning rate scheduler is gone, because the code has been simplified and the troubled part has been replaced with Trainer.create_scheduler().

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.

From what I can see, this *should* be fine because it copies the `Trainer` which should include the original `args.lr_scheduler_kwargs`. Can you verify this can be closed now @haixpham ?

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.