Add LR Scheduler to full finetune recipes

[joecummings]

Following how we do it in LoRA recipes, we should add the ability to use a LR scheduler to the full finetune recipes

Here's everything that would need to be changed:

• Add appropriate functions to recipes/full_finetune_distributed.py (follow recipes/lora_finetune_distribute.py for example)
• Add appropriate functions to recipes/full_finetune_single_device.py (follow recipes/lora_finetune_single_device.py for example)
• Run an example run with Llama3.1 8B and log to Weights & Biases. Post screenshot confirming this works in your PR

I don't think we need to update all our configs yet, just need to make sure the functionality exists.

[felipemello1]

@joecummings , I was thinking that this could be a function in utils that takes an optimizer as input, and just does the step for it, that way it can be shared for any recipe. Maybe it needs a flag for "is_distributed".

This complicates the task, but maybe we could follow a similar pattern for model checkpointing and early stopping, which are also frequently asked features.

[joecummings]

We CaN aLWayS dO iT IN a FollOw Up REfaCtOR

[joecummings]

Kidding, I know that will never happen. I like the idea.

[user074]

I guess it kinda work for the full_finetune_single_device.py. I will update after I finish the training to see how it goes

[SalmanMohammadi]

Can we close this?

[ebsmothers]

@SalmanMohammadi technically I think we still need to add to the distributed recipe to fully close it out

[user074]

I can have a check on the distributed case but probably after mid November

[felipemello1]

fyi, reference implementation from a fork: https://github.com/homebrewltd/torchtune/blob/1ff4fc217e578a1656dc60c343dcee98e7ac013e/recipes/full_finetune_fsdp2.py#L285