[BUG] How to checkpoint optimiser states to resume fine-tuning at a later stage?

It looks like the code (https://github.com/microsoft/DeepSpeedExamples/blob/master/applications/DeepSpeed-Chat/training/step1_supervised_finetuning/main.py) uses custom functions (e.g. save_zero_three_model) which uses NCCL primitives (e.g. All-Gather) to gather the sharded model parameters and store them effectively.

How should I go about similarly for storing optimiser states. We are trying to do long fine-tuning with rich and diverse examples (~3M examples) and thus would require support for checkpointing optimiser and lr_scheduler states and load the model. If these states are sharded, what is the best method? Could you please point to relevant materials or DeepSpeed code snippets to perform the same?

DeepSpeed engine has utilities such as save_checkpoint which seems to transparently do this. However, the DeepSpeed-Chat code doesn't seem to utilise the engine.

Any quick help here would be much appreciated. Thanks.

Start with applications/DeepSpeed-Chat/training/step1_supervised_finetuning/main.py and inspect save_zero_three_model alongside DeepSpeed engine save_checkpoint. Determine how sharded optimiser and lr_scheduler states should be stored and restored for resumed fine-tuning. Done means the training flow can save and load the model, optimiser state, and scheduler state across checkpoints.