Inverse Square Root LR Schedule

[mansheej]

Adds the Inverse Square Root LR Scheduler.

This scheduler is meant to easily enable continual learning. It consists of three components:

1. A linear LR warmup.
2. A component where the LR decays as an inverse square root in the number of steps to a constant value at infinite time.
3. An optional linear cooldown.

The image below show two examples of the LR schedule with a 10 step warmup, with either no cooldown (orange) or a 20 step cooldown (blue) run for 100 steps total.

[dakinggg]

Will review in full as well, but is there any reason this should be in LLM foundry? It seems generically useful and I would probably prefer it is in composer directly.

[mansheej]

Will review in full as well, but is there any reason this should be in LLM foundry? It seems generically useful and I would probably prefer it is in composer directly.

This implementation and hyperparameters are a little different from how people typically do it. I wanted to first get it into LLM foundry, run some more experiments and have other people use it so we can work out some kinks and learn about some best practices, and then eventually upstream it into composer with the best practices documented.

[dakinggg]

Could you please add some unit tests testing this schedule produces what you expect?

[mansheej]

Could you please add some unit tests testing this schedule produces what you expect?

Do you have any suggestions for reasonable unit tests for the LR scheduler?

[dakinggg]

I'd suggest:

1. the "build" function can create it successfully
2. set up a couple simple schedules with known values that you expect and test that your scheduler code produces exactly the schedule you expect.

[b-chu]

Add some unit tests to test functionality

[mansheej]

Added unit tests and tried to address all the comments. Failing the Code Quality Checks, but I'm not sure why from the Error message.

[dakinggg]

@mansheej try running pre-commit run --all-files locally

[dakinggg]

Basically LGTM, left a few comments on simplifying the tests.

[mansheej]

Addressed most of the comments. The rest can perhaps be filed for improvement/addressed when we upstream to Composer?

[dakinggg]

I am ok with the level of test coverage here. Before merging, could you please make the PR description more complete?

In this case, I think an example schedule produced by this code (e.g. a wandb graph) and a one sentence description of the gist of the schedule would suffice. Thanks!

[b-chu]

Two small changes, thanks for the tests!