[FSDP] Support scheduled learning rate

[PopSoda2002]

Currently FSDP only supports fixed learning, does not support like warm up, cosine decay like megatron

[ChangyiYang]

I would glad to take it with @aaaandychen :)

[ChangyiYang]

After some investigation, we have two options for LR scheduling:

A. Use PyTorch built-in schedulers

📌 https://github.com/pytorch/pytorch/blob/main/torch/optim/lr_scheduler.py

Pros: easy integration, maintained by PyTorch
Cons: no warmup+decay in one place → not fully compatible with Megatron

B. Implement a custom scheduler based on Megatron

📌 https://github.com/NVIDIA/Megatron-LM/blob/e847643f88758b7f5577b119240d3dfd1b4b4661/megatron/core/optimizer_param_scheduler.py

Pros: supports all Megatron features (warmup, min_lr, cosine/linear/inv-sqrt/WSD)
Cons: needs own maintenance, only support lr, not weight decay

Pytorch currently does not suport dynamic weight decay.

I would prefer option B, any opinion is welcomed!

[PopSoda2002]

Fell free to implement whatever you want! But I am not sure we need decouple FSDP and Megatron

[ChangyiYang]

I will go with option B. Since we are copying the code from megatron, not refer from it, FSDP and Megatron are still decoupled.

[PopSoda2002]

solved by #1040 thanks! @ChangyiYang