[Distributed] Support index + multi-bin loading

[kwen2501]

Some HF format uses index + multiple binaries instead of index + multiple safetensors.
Distributed workflows already support the later, this PR adds support for the former too.
(The reason is that quantized checkpoint cannot save as safetensor format yet.)

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 9ef2171 with merge base 766bee9 ():
💚 Looks good so far! There are no failures yet. 💚

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

[mikekg]

Does this compose with #1255 ?

[kwen2501]

Does this compose with #1255 ?

Thanks for the review. I hope yes.
Distributed already have support for "index + multi-safetensor" format (thanks to @lessw2020 ).
This PR only adds support for "index + multi-bin" format. (The reason is that quantized checkpoint cannot save as safetensor format yet.)

[mikekg]

Does this compose with #1255 ?

Thanks for the review. I hope yes. Distributed already have support for "index + multi-safetensor" format (thanks to @lessw2020 ). This PR only adds support for "index + multi-bin" format. (The reason is that quantized checkpoint cannot save as safetensor format yet.)

How much of this can be shared with local-only weight loading? Should distributed weight loading just be a general case of local loading (or vice versa, local just a special case of distributed?) Not suggesting to abandon this PR at all, but similar to what you've done to merging model.py and model_dist.py (yay! Awesome work!!!!) could be a real value here.

[lessw2020]

lgtm!
Thanks for adding this!