[BE] Lr schduler flatten #794
Conversation
facebook-github-bot
added
the
CLA Signed
| @@ -183,9 +183,9 @@ def __init__( | |||
| "model": ModelWrapper(model_parts), | |||
| "optimizer": optimizers, | |||
| "dataloader": dataloader, | |||
| "lr_scheduler": lr_schedulers, | |||
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
Compared lr_schedulers before and after flattening, with/without checkpoint
lr_scheduler values are consistent with changes here
There was a problem hiding this comment.
does it support DCP resharding? e.g. PP degree from 2 to 4 across two jobs
There was a problem hiding this comment.
I think this PR doesn't address the resharding issue, hence the [BE] prefix. Supporting lr resharding deserve a separate PR.
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 under self.state['lr_scheduler'], which is consistent with optimizer, model and data_loader
The PR is tested by 2 parts:
before and after this PR, lr_shceduler values are the same
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:
