adding support for LR schedule for full distributed finetune

[tginart]

My understanding is that the full multi-gpu fine-tuning doesn't yet support learning rate schedules.

Would it be possible to add support for this? Even basic ones, such as linear warmup follow by cosine or linear decay?

I can also take a look at doing this myself if I could get a pointer to the code.

Thank you!

[felipemello1]

hey @tginart ,

We have it in our recipes here:

torchtune/recipes/full_finetune_distributed.py

Line 341 in 6923833

self._lr_scheduler = self._setup_lr_scheduler(

It gets called in the training loop here:

torchtune/recipes/full_finetune_distributed.py

Line 798 in 6923833

self._lr_scheduler.step()

The issue seems that we don't have it exposed in some configs. I will fix that. You can add this to your config, and it should work:

lr_scheduler:
  _component_: torchtune.training.lr_schedulers.get_cosine_schedule_with_warmup
  num_warmup_steps: 100

Thanks for flagging that!

[tginart]

Thank you!

[Ankur-singh]

@tginart thanks for pointing it out!

The issue seems that we don't have it exposed in some configs. I will fix that. You can add this to your config, and it should work:

@felipemello1, I apologize if I'm overstepping here, but I think this could be part of a bigger issue. For someone familiar with the codebase (at least the recipes), it's clear how to adapt the config to their needs. They can easily see the different components used in the recipe and the various arguments allowed for each one. However, for others, this process isn’t as obvious since it’s not documented. I believe documenting this "esoteric" knowledge would greatly improve the accessibility of the library for new users, while also offering more flexibility for advanced users.

[felipemello1]

hey @Ankur-singh , i agree! We try to expose as much as possible through our configs and add comments so users feel ok exploring them. We also document most of our performance flags and have them in the readme. But i 100% agree that there is room to improve by a lot. This LR scheduler is an example of something that we added to the recipe and forgot to expose in the configs.

If you feel that there is something else in specific that we should prioritize adding, or some tutorial, please feel free to suggest or make a PR.