Integrate distributed inference with chat/server

[mreso]

This PR integrates distributed inference with the CLI options chat and server and implements Option 2b of #1376.

To test run on machine with 4 gpus:

python  torchchat.py chat llama3.1  --max-new-tokens 10  --distributed --tp 2 --pp 2
python  torchchat.py server  llama3.1  --max-new-tokens 10  --distributed --tp 2 --pp 2

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❗ 2 Active SEVs

There are 2 currently active SEVs. If your PR is affected, please view them below:

• [DomainsOnly] Jobs fail with GLIBC version not found
• Failure to deploy Self-Hosted runners - High Queue Times and canceled jobs

✅ No Failures

As of commit 76895cc with merge base 5da240a ():
💚 Looks good so far! There are no failures yet. 💚

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

[mikekgfb]

What's the gap to using the models described in model_config/models.json. (as implied by TODO comment)

Definitely should not be part of the present PR, but I think as a north star, it would be desirable to grab the same models (and download and mgmt infra etc) for non-distributed and distributed.

[mreso]

Thanks for the comment @mikekgfb. The gap wasn't that big, the args.model just wasn't accessible at that point and I wanted to take a deeper look to fix it right. Removed the cruel hack and I now save the distribution_path when creating the builder_arg. Still not sure if this is the intended way of locating the checkpoint though.

[mikekgfb]

Not sure what you're trying to do. If it's just that the dict/json data structure describing the mapping isn't in scope, maybe what you want is some methods that give you the relevant info?

Also, these seem to be mapping of short names to HF network paths - should we not have a way to [ick them from the local filesystem (since the torchchat cli already manages download and all that). Oh, and if the answer is "we have bigger fish to fry, hooking this up is not highest priority" I will wholeheartedly agree. This is more about understanding the context of this PR.

Where I'm lacking the context is how you go from all the weights being available locally on a node to reading those weights on another node? And maybe that's why you prefer to straight up pick the files from HF? (Although local distribution from an already downloaded set of weights probs has higher bandwidth?) Again, there's much bigger fish to fry, and I think this PR is a good step in the direction of frying those fish ;)