[Distributed] Accept model names

[kwen2501]

Change 1

Add argparser to accept model name and other arguments (e.g. PP degree).

$ torchrun dist_run.py -h
usage: dist_run.py [-h] [--pp PP] {llama2-7b-chat,llama3}

positional arguments:
  {llama2-7b-chat,llama3}
                        Name of the model to load

options:
  -h, --help            show this help message and exit
  --pp PP               Pipeline parallel degree

Supported model names:

llama2-7b-chat
llama3

(Following model aliases in models.json)

Change 2

Renamed variable hf_model_name to distribution.

A distribution looks like this:
meta-llama/Llama-2-7b-chat-hf
meta-llama/Meta-Llama-3-8B-Instruct

Change 3

Fix output undefined error when PP = 1.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 980d1ff with merge base 26c1d8b ():
💚 Looks good so far! There are no failures yet. 💚

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

[lessw2020]

nit: already filtering based on ranks in the dist_run.sh launcher rather than having express rank based logging in the code.

[kwen2501]

Yeah, it is a hard choice.. I don't know if regular users use filter (I never remember that option's name) and want to have a minimal command line work.

[lessw2020]

looks good!
I also have toml config supporting code in distributed/config_manager.py that I had integrated earlier, so we should add in both cmd line and toml support ultimately (similar to titan where you can use either).

[Jack-Khuu]

Rebasing on main to pick up CI fix

[kwen2501]

Rebasing on main to pick up CI fix

Thanks @Jack-Khuu !