Update lr_schedules.py

[CoinCheung]

add cosine annealing scheduler

this scheduler is widely used in image classification task, and many llm (e.g. llama) use this also.

[CoinCheung]

@microsoft-github-policy-service agree

[tjruwase]

@CoinCheung, thanks for the PR. A few items to address.

To fix formatting issues use this guide.

Please add unit test: example

Inspect failing CI tests.

[CoinCheung]

@tjruwase @wjessup @dfyz @manuelciosici I have no experience with pr for deepspeed, what is the status of this now? Is there any further operation that needs me to work on?

[tjruwase]

@CoinCheung, thanks for making the changes. We will review and merge once the CI passes.

[CoinCheung]

Hi @tjruwase ,

I have made some fixes, would you please help me launch CI test one more time?

[CoinCheung]

@tjruwase Would you please launch CI one more time ?

[CoinCheung]

Hi @jeffra @mrwyattii I think the problem is not with my fix, it is a inference error, but my fix is about training learning rate scheduler. Can this fix be merged ? Or is there other things that need me to commit?

[tjruwase]

@CoinCheung, sorry for the delay. It seems the issue is with our CI system. Please bear with us while we resolve the problem.

[CoinCheung]

Hi @tjruwase , What is the status of this thread?

[tjruwase]

@CoinCheung, I have restarted CI. Let's see how it goes.

[CoinCheung]

Hi @tjruwase ,

Is this associated with my changes?

[tjruwase]

@CoinCheung, no I don't think it is related to your changes.

[kmn1024]

Should WarmupCosineLR inherit from WarmupLR?
https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/runtime/lr_schedules.py#L774

[tjruwase]

Should WarmupCosineLR inherit from WarmupLR? https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/runtime/lr_schedules.py#L774

Yes, you are correct. It should.

@CoinCheung, are you able to refactor your changes? Thanks!

[CoinCheung]

Hi @tjruwase @kmn1024 , I do not think WarmupCosineLR can be interited from this WarmupLR in this case.

SInce they use different methods to determine the learning rates. For WarmupCosineLR, I use "ratio of original lr values", which I think should be more scientific, while WarmupLR uses specific lr values.

For example, when using WarmupLR, by setting warmup_min_lr=1e-5, warmup_num_steps=100, the scheduler will set lr from 1e-5 to max lr within 100 steps.
When using WarmupCosineLR, by setting warmup_min_ratio=0.1, warmup_num_steps=100, and we assume lr=1e-3 when we defined the optimizer, the scheduler would set lr from 0,1 * lr = 1e-4 to max lr within 100 steps. We do not set specific lr values to a scheduler.

The reason why I feel using ratio is better: we do not need to set specific lr values everywhere in both optimizer and scheduler. When we define an optimizer, we need to consider the learning rate. When we define a scheduler, what we only need to do is to determine the shape of the learning rate curves, rather than its specific values. When you want to keep shape of lr curve and only tune peak lr, you only need to change one place. This follows the principle of "each module only does its own work, and their settings are not impacted by each other".

From my experience of tuning models, this method is less likely to cause mistakes that I change optimizer lr but forgot to change scheduler.

Also in some paper, if I recally correctly, they claimed that they use CosineLR to train their model, and the learning rate anneals from max_ lr to 0.1 * max_lr. I think many other people accept this method of tuning learning rates.

[tjruwase]

@CoinCheung, thanks for your response. I agree with the differences that you identify between WarmupLR and WarmupCosineLR, but these differences are to me simply in the implementation and logic. At the high-level they are similar because of they provide two phases of lr changes: (1) initial phase of warmup/increase, and (2) final phase of no change or decay. Looking more closely we observe significant similarity or duplication in many of the methods including step, state_dict, load_state_dict, get_last_lr, _format_param. These similarities suggest to me opportunities for to code refactor and reuse.

[CoinCheung]

@tjruwase Would it be acceptable if we change args (init args used for define the scheduler object) of WarmupLR ? It has only one sub-class WarmupDecayLR, and I think its usage frequency is not very high.