[BE] Lr schduler flatten

[mori360]

Currently, lr_scheduler is stored differently as optimizer, model and data_loader, with keys to be "lr_scheduler_0", "lr_scheduler_1", ... stored in the state
This PR aims to flatten lr_shceduler so that all the schedulers would be stored as a list of state_dict under self.state['lr_scheduler'], which is consistent with optimizer

Here we have the assumption that all the optimziers have the same lr_scheduler, thus only to save a single lr_scheduler's state_dict and load it to all the schedulers works here.

The lr_scheduler has the state_dict like:
{'base_lrs': [0.0003], 'last_epoch': 1, 'verbose': False, '_step_count': 2, '_get_lr_called_within_step': False, '_last_lr': [2.985074626865671e-06], 'lr_lambdas': [{}]}

The PR is tested by 2 parts:

1. test lr_scheduler value before and after checkpoint, resharding with degree changes on tp and pp.
   [dp=2, tp=4, pp=1] -> [dp=2, tp=1, pp=4]
   [dp=2, tp=1, pp=4] -> [dp=2, tp=4, pp=1]
   date_loader does not support resharding right now.

logs:
[dp=2, tp=4, pp=1]
step 5 before saving to checkpoint:
[{'lr': 8.955223880597014e-06, ...}]

step 10 after loading from checkpoint and reshard to [dp=2, tp=2, pp=2]:
[{'lr': 1.6417910447761194e-05, ...}, {'lr': 1.6417910447761194e-05, ...}]

[dp=8, tp=1, pp=1]
step 5 without checkpoint:
[{'lr': 8.955223880597014e-06, ...}]
step 10 without checkpoint:
[{'lr': 1.6417910447761194e-05, ...}]

2. Memory trace:
   Before the flatten, rerun llama3_8b.toml from step 5 to step 10:

After the flatten, rerun llama3_8b.toml from step 5 to step 10:

[tianyu-l]

I think it won't be this simple. Both OptimizersContainer and ModelWrapper define state_dict and load_state_dict to handle flattening and unflattening. Since we don't have things like get_model_state_dict and set_model_state_dict for lr scheduler in torch.distributed.checkpoint.state_dict, we likely will need to manually write something for the LambdaLR we are using. See #738 (comment)

Let's work with @fegin on this.

[mori360]

Compared lr_schedulers before and after flattening, with/without checkpoint
lr_scheduler values are consistent with changes here

[tianyu-l]

does it support DCP resharding? e.g. PP degree from 2 to 4 across two jobs

[fegin]

I think this PR doesn't address the resharding issue, hence the [BE] prefix. Supporting lr resharding deserve a separate PR.

[tianyu-l]

We may need to explicitly copy state_dict before loading. Otherwise there could be silent errors. See details of behavior here https://github.com/pytorch/pytorch/blob/v2.6.0/torch/optim/lr_scheduler.py#L359

Please add more detailed comment/NOTE here in the code.
Please add verified experiment results in the PR summary.

We should consider add unit test under the test folder to guard the behavior, but feel free to do this in a later PR.

[tianyu-l]

I checked the LambdaLR scheduler code. It seems the only thing matters in the state is the current step which is an int, so load_state_dict will automatically make copies. See last_epoch in https://github.com/pytorch/pytorch/blob/v2.6.0/torch/optim/lr_scheduler.py#L122

Therefore, the behavior should be correct, as long as we don't modify training.steps and training.warmup_steps when resuming from a checkpoint.

But for safety, let's still explicitly call .copy() on the state_dict, as the overhead is small.

Let's document this in the code here.

[mori360]

Thanks for the comments. Add more detailed comments here with our discussion results.
Update summary with multi-optimizer results, with the same lr values after checkpoint and resharding

[tianyu-l]

Looks awesome. Thanks for the effort!