Add checkpoint load step

[msaroufim]

Fixes #662

followed @fegin advice to test this and indeed things are working https://gist.github.com/msaroufim/2925b3f17b631bf370a49f185b6e169d

[checkpoint]
enable_checkpoint = true
folder = "checkpoint"
interval_type = "steps"
interval = 10
model_weights_only = false
export_dtype = "float32"
async_mode = "disabled" # ["disabled", "async", "async_with_pinned_mem"]
load_step = 10

[tianyu-l]

Thanks! It'd be great to also update docs/checkpoint.md

[tianyu-l]

default should be -1? https://github.com/pytorch/torchtitan/blob/main/torchtitan/checkpoint.py#L452

[fegin]

You can just run a local script and checkpoint every 10 steps and then load again from 20th step. The logging should show the training starts from 20th step. If load_step is not used, it should just load the checkpoint and finish the training.

[tianyu-l]

You can just run a local script and checkpoint every 10 steps and then load again from 20th step. The logging should show the training starts from 20th step. If load_step is not used, it should just load the checkpoint and finish the training.

@fegin I think the original ask in #662 was to resume from a particular step, not necessarily the latest one.

[fegin]

What I meant is that if load_step is not specified, the training won't happen because the trainer will just load the last step checkpoint. But if load_step is implemented appropriately, and we specify some intermediate step, like 20, the logging should start from 20. This can be used to verify the implementation.

[tianyu-l]

oh I was assuming Mark's point was how to write a light-weight unit test for it. I think it is not blocking as unit tests for checkpointing in general are missing. I think it's worth updating the checkpoint.md though.

[tianyu-l]

thanks!
Before merge, can you also help update the how-to in docs/checkpoint.md with one more bullet point?

[msaroufim]

thanks!
Before merge, can you also help update the how-to in docs/checkpoint.md with one more bullet point?

Added it, I can do another PR to better document checkpoint.md since right now it just lists configs