[Distributed] create model on meta device

[kwen2501]

Change 1

Changing from:

    with device:
        model = Transformer(config)

to

    with torch.device("meta"):
        model = Transformer(config)

because when we later load weights, we will swap out the tensors of the model anyway. So saving an on-device init here.

Change 2

Also added a with device context to model.setup_cache() call. So that the caches are directly created on the target device -- saving a model.to(device) call.

[pytorch-bot]

🔗 Helpful Links

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

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

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

[metascroy]

Rebase to fix failing tochao_experimental check

[lessw2020]

lgtm!
Tested with both llama2 and llama3.