[Distributed] Add support for torchchat checkpoint format

[kwen2501]

Distributed workflow supports two formats now, controlled by the --chpt-from flag ("hf" or "torchchat").
The difference is whether there is an index file -- impacts loading in pipeline parallel case.

    if chpt_from == "hf":
        # This format stands for: index file + multiple binary files

    elif chpt_from == "torchchat":
        # This format stands for:
        # single binary file, OR
        # multiple binary files without index files.

torchrun --standalone --nproc-per-node 4 dist_run.py llama3 --pp 2 --chpt-from hf
torchrun --standalone --nproc-per-node 4 dist_run.py llama3 --pp 2 --chpt-from torchchat
both run.
But they generate different responses. To be investigated further.

TODO:
we should create real model instead of creating Transformer.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchchat/1268

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 0565e8b with merge base b217158 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[mikekg]

Can you figure out from what's in the location whether it's one or the other format rather than adding another option? (If push comes to shove, you could expect the checkpoint path point to the index file?)

Also, does this compose w/ #1255?

[kwen2501]

@mikekg Good idea. Should be easy to implement the detection.

[lessw2020]

for better clarity, 'chkpt' is a much better abbreviation for checkpoint. chpt is short for 'chapter' which is confusing.

[lessw2020]

doen = done :)

[lessw2020]

need to pass in if this is torchchat or hf to permute or not?

[lessw2020]

need to check what you are loading here.
torchchat won't need permuting.

[lessw2020]

approving to land.
I think your code here doesn't work b/c you are not checking what type of state dict you are updating. Thus it seems you are going to permute a torchchat checkpoint which doesn't need it.
That may be why: "But they generate different responses. To be investigated further."

Also, please switch the abbreviation for checkpoint to "chkpt' and not "chpt". 'Chpt' to me = 'chapter' and reads confusingly :)