[Config] Make the checkpoint step configurable.

[casper-hansen]

load() takes an argument step. Make this configurable in the configuration, so we can choose from which step to continue training.

torchtitan/train.py

Line 208 in 193ce98

checkpoint_loaded = checkpoint.load()

[awgu]

Related: #571

[gnadathur]

cc: @saumishr

[casper-hansen]

I just did this for now in my own train.py script, but this seems like a low-hanging fruit that someone could raise a PR to make it an argument in the parser.

    job_config.parser.add_argument("--custom.from_checkpoint_step", type=int, default=0)