Support constant lr with cooldown

[LoserCheems]

What does this PR do?

Fixes #35449
Added the 'warmup_stable_cooldown' learning rate scheduler method.
This method allows three phases, linear warmup, stable, and cooldown, where cooldown can be done using the linear, cosine and 1-sqrt methods.

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

@muellerzr and @SunMarc

[SunMarc]

Thanks for the PR @LoserCheems, we already have warmup_stable_decay scheduler with get_wsd_schedule that does practically the same thing as the new scheduler you proposed. I think what you be good is to maybe add more options to the decay stage of the wsd_scheduler as you did in this PR (e.g linear/ 1-sqrt) WDYT ?

[LoserCheems]

Thanks for the PR @LoserCheems, we already have warmup_stable_decay scheduler with get_wsd_schedule that does practically the same thing as the new scheduler you proposed. I think what you be good is to maybe add more options to the decay stage of the wsd_scheduler as you did in this PR (e.g linear/ 1-sqrt) WDYT ?

oh, thank for your suggestion, I intend to integrate different cooldown methods and minimum learning rate into get_wsd_schedule.
In addition, is it possible to rename get_wsd_schedule to get_wsc_schedule, which is more in line with the three phases of warmup, stable, cooldown.

[SunMarc]

Thanks ! We can't really rename it wsc, wsd is a real term that appears in the MiniCPM paper.

[LoserCheems]

Thanks ! We can't really rename it wsc, wsd is a real term that appears in the MiniCPM paper.

Thank you, the renaming of the function is done.

[SunMarc]

Thanks for the PR ! Left a comment

[SunMarc]

Nice ! Just a few nits and we should be able to merge this soon ! Thanks for being reactive

[LoserCheems]

Nice ! Just a few nits and we should be able to merge this soon ! Thanks for being reactive

Everything is ready!

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[SunMarc]

LGTM ! Thanks for iterating ! Just a few nits

[LoserCheems]

LGTM ! Thanks for iterating ! Just a few nits

It's done!

[SunMarc]

LGTM!

[LoserCheems]

@SunMarc Thank you, this PR is ready for merge.