[Feature Request] Resume from Checkpoint

[pbontrager]

Resuming from checkpoint is very important for recovering from crashes or for extending training but there are a lot of moving parts in RL pipelines that we need to decide on. At a minimum we need all of the things that are need from regular training:

• Checkpoint
• Optimizer State
• LR Schedulers (with option to extend)
• data step
• seed (if you can use it)

But on top of this there are a lot of other things that could be restored in an RL pipeline

• replay buffer data
• states of any tools or stateful services
• any additional models critics or reward models that get updated

I think we can start with just the first.

[ebsmothers]

cc @DNXie

[Jack-Khuu]

Dup: #223

[DNXie]

Looking into this, it seems like titan already handled save/load optimizer/dataloader/lr_scheduler states, at least it should've.

checkpointer.states has all this information
I did some test run and confirmed that in checkpoint.save and checkpoint.load, the states contains keys such as 'train_state', 'model', 'optimizer', 'dataloader', 'lr_scheduler'.

However, it seems that the load didn't restore all the states back in the objects.

Looking at the dcp.load which is load from torch.distributed, it doesn't return anything. Am I reading it wrong?

My experiment:
I was looking into the optimizers, and confirmed that at dcp_save, the state_dict['optimizer'] contains all optimizer status, which is passed into dcp.save. I suppose they are saved in the checkpoint:

[0] [DEBUG] dcp_save: state_dict['optimizer']: OptimizersContainer (
[0] Parameter Group 0
[0]     betas: (0.9, 0.95)
[0]     eps: 1e-08
[0]     foreach: False
[0]     fused: True
[0]     lr: 1e-05
[0]     weight_decay: 0.1
[0] )

Then I start another run, this time I set eps=1 in the config. I load from the previously saved checkpoint, which should overwrite the parameter and reset eps=1e-8 in theory. But after loading, nothing changed.

[DNXie]

cc @wwwjn @tianyu-l