How to train with cosine annealing?

[ZQ-Dev8]

Is it possible to train a model with axolotl using some form of cosine annealing with warm restarts? A hacky manual way would be to train individual epochs with the default cosine scheduler + warmup steps + dataset shuffle, but I seem to remember this was once supported out-of-the-box?

[NanoCode012]

Can you check if this function get_cosine_schedule_with_min_lr within the code does what you want?

[ZQ-Dev8]

@NanoCode012 unless I'm missing something, it does not appear to without some level of modification. What it needs to do is essentially an entire iteration of a cosine schedule, except across a single epoch. So for each epoch during training, there are warmup steps up to a max_lr, with cosine decay down to a min_lr. I've also seen variations of this where the max_lr is reduced per epoch.