[Scheduler] Add support for cosine and wsd scheduler

[yzhangcs]

What does this PR do?

This PR introduces support for cosine and WSD schedulers. The cosine scheduler is widely used in LLM training (e.g., Pythia, OLMo, Llama), while the WSD scheduler, introduced by MiniCPM, features a three-stage learning rate schedule: warmup, stable, and decay. The stable stage keeps the learning rate constant, beneficial for continual pretraining and flexible training budget adjustments.

Why this PR is necessary

The cosine scheduler is a standard for LLM training, and the WSD scheduler addresses specific needs like stable learning rates during training. This PR also supports three decay types: linear (MiniCPM's original approach), cosine (used in hf transformers), and 1-sqrt (found optimal in this paper).
These additions provide flexibility and improved performance for diverse training scenarios.

[yzhangcs]

Here is running examples for three schedulers.

bash run_train.sh --optimizer.scheduler=linear \
  --training.steps=50 \
  --training.warmup_steps=5 \
  --optimizer.min_lr_ratio=0.1

bash run_train.sh --optimizer.scheduler=cosine \
  --training.steps=50 \
  --training.warmup_steps=5 \
  --optimizer.min_lr_ratio=0.1

bash run_train.sh --optimizer.scheduler=wsd \
  --training.steps=50 \
  --training.warmup_steps=5 \
  --optimizer.min_lr_ratio=0.1

[tianyu-l]

In fact, I think the three can be unified.
We only need to define warmup_ratio, decay_ratio, lr_decay_type (among linear sqrt cosine), and maybe lr_min (calling it ratio again would sound confusing, we can explain in the helper message it's a ratio), to achieve everything.

We can explain in the helper message, or in a doc, how to use them to achieve various combinations, including the three you explicitly wrote today.

[lessw2020]

Hi @yzhangcs - thanks for the PR.
I wanted to comment that we purposefully did not include cosine lr as recent work has shown the linear rate is more effective:
Please see: https://arxiv.org/abs/2310.07831
"We perform the most comprehensive evaluation of learning rate schedules to date, evaluating across 10 diverse deep learning problems, a series of LLMs, and a suite of logistic regression problems. We validate that overall, the linear-decay schedule outperforms all commonly used default schedules including cosine annealing. "

That said, I believe the WSD may be useful in some cases (MSFT had a paper on this a long time ago, and I used it back in AI competitions), and the 1 sqrt seems to be newer work than the above.
But I generically wouldn't recommend anyone use cosine lr just on the basis that others use it...the above work pretty much closed that discussion.
(I'm not against enabling it if people want to do their own comparisons, though a link to the above would probably change their minds..but just clarifying why we did not add it when initially setting the default lr schedule).

Also, we used to display the lr but it was removed b/c most people didn't need it in the display...thus would recommend that this display is optional/configurable.

[yzhangcs]

@tianyu-l Hi, just updated the stuffs you mentioned.

[yzhangcs]

@lessw2020 Thank you for your comments, I'm wondering if it's ok to add some hints and paper link https://arxiv.org/abs/2310.07831 in config_manager?

Also, we used to display the lr but it was removed b/c most people didn't need it in the display...thus would recommend that this display is optional/configurable.

Agreed. I just add the display to make sure the decay is right.
It's ok for me to add a configurable option or directly revert the changes for displaying lr (seems that the number of colors is not enough to distinguish loll).
Which one do you prefer? @lessw2020 @tianyu-l

[tianyu-l]

I left some inline comments. Below are some general comments.

Regarding file length/complexity:
I think we should separate things into a new file called lr_scheduler.py. This can be in a separate PR.

Regarding warmpup steps:
Since warmup behavior is closer to lr_scheduler, should we move training.warmup_steps also to the optimizer section (or even consider creating an lr_scheduler section)?

Regarding logging lr rates:
I think we should still log lr to TensorBoard / WandB, maybe after the #945
It can be called from a new get_lr function of LRSchedulersContainer, so that people can modify/inherit it to adapt to desired behaviors.

[yzhangcs]

@tianyu-l Hello, I've just updated this PR based on your suggestions. Could you please review it again to see if there's anything I might have missed?

[yzhangcs]

Checks for scheduler.decay_type:

bash run_train.sh --scheduler.warmup_steps=4  --scheduler.decay_ratio=0.9  --scheduler.decay_type=linear  --training.steps=40

bash run_train.sh --scheduler.warmup_steps=4  --scheduler.decay_ratio=0.9  --scheduler.decay_type=cosine  --training.steps=40

bash run_train.sh --scheduler.warmup_steps=4  --scheduler.decay_ratio=0.9  --scheduler.decay_type=sqrt  --training.steps=40

Checks for scheduler.decay_ratio:

bash run_train.sh --scheduler.warmup_steps=4  --scheduler.decay_ratio=0.2  --scheduler.decay_type=linear  --training.steps=40

[tianyu-l]

Looks good to me! I suggest naming scheduler to lr_scheduler. Also please do a final rebase.

[yzhangcs]

@tianyu-l Hi, thank you for the feedbacks. Just fixed them.

[tianyu-l]

Thank you!

Could you rebase to latest torchtitan main? I'm seeing changes in merged #945 as part of this PR (e.g. train.py)

[yzhangcs]

@tianyu-l Hey, just wanted to check if everything looks good on your end. I’m still getting familiar with rebase, so I want to make sure I did it correctly.
Let me know if there’s still anything I need to adjust!

[yzhangcs]

On the newest commits:

bash run_train.sh --lr_scheduler.warmup_steps=4  --lr_scheduler.decay_ratio=0.2  --lr_scheduler.decay_type=linear  --training.steps=40

[tianyu-l]

@yzhangcs
Sorry I'm seeing a lot of warnings in CI tests.

[rank0]:[titan] 2025-03-10 22:32:03,960 - root - WARNING - The warmup steps should be less than or equal to the warmup-stable steps (1.9999999999999996). Consider reducing either the decay ratio (0.8) or the warmup steps (2)

Could you help resolve it, e.g. setting decay_ratio to None?

Also the warning is called every single LR scheduler step. Let's move the check outside of this function and do it only once when setting up the scheduler function.
https://github.com/pytorch/torchtitan/pull/938/files#diff-fb95a3c6c32a395a9dafc1625830d088ded3dd406583d5127f7244bd29d2f759R427

[yzhangcs]

@tianyu-l sure, I will create a new pr to fix it.