[T170073014] Rewrite distributed examples for Tensor Parallel, Sequence Parallel, 2D (FSDP + TP)

[lessw2020]

This PR updates the three distributed examples for Tensor Parallel, Sequence Parallel and 2D with the following main changes:
(note - internal reference - task [T170073014] Rewrite TensorParalell/SequenceParallel Examples using our new UX)

1 - move to torchrun launching (see run_.sh files) and relevant world topology introspection in the setup instead of mp.spawn.
2 - move device mesh creation to new api, init_device_mesh
3 - use custom parallelization plans (ColwiseParallel and RowwiseParallel) rather than the previous prebuilt PairwiseParallel() and SequenceParallel()
4 - For the 2D example - used a more relevant swiglu MLP model to showcase applying 2D to a more sophisticated/llama style situation.
5 - Adds more interactive UI for the user (start, per iter, and completion feedback).

[netlify]

✅ Deploy Preview for pytorch-examples-preview canceled.

Name	Link
🔨 Latest commit	5f4a5d3
🔍 Latest deploy log	https://app.netlify.com/sites/pytorch-examples-preview/deploys/655e7ed39bbef400093f38f9

[lessw2020]

• test failures are related to being unable to import init_device_mesh (?), no gpu's available (?), and lastly need to modify these tests to launch via the .sh files associated with each example (to run torchscript):

Traceback (most recent call last):
  File "tensor_parallelism/tensor_parallel_example.py", line 5, in <module>
    from torch.distributed._tensor.device_mesh import init_device_mesh
ImportError: cannot import name 'init_device_mesh' from 'torch.distributed._tensor.device_mesh' (/opt/hostedtoolcache/Python/3.8.18/x64/lib/python3.8/site-packages/torch/distributed/_tensor/device_mesh.py)
tensor parallel example failed
Traceback (most recent call last):
  File "tensor_parallelism/sequence_parallel_example.py", line 5, in <module>
    from torch.distributed._tensor.device_mesh import init_device_mesh
ImportError: cannot import name 'init_device_mesh' from 'torch.distributed._tensor.device_mesh' (/opt/hostedtoolcache/Python/3.8.18/x64/lib/python3.8/site-packages/torch/distributed/_tensor/device_mesh.py)
sequence parallel example failed
Traceback (most recent call last):
  File "tensor_parallelism/two_d_parallel_example.py", line 18, in <module>
    from torch.distributed._tensor.device_mesh import init_device_mesh
ImportError: cannot import name 'init_device_mesh' from 'torch.distributed._tensor.device_mesh' (/opt/hostedtoolcache/Python/3.8.18/x64/lib/python3.8/site-packages/torch/distributed/_tensor/device_mesh.py)
2D parallel example failed
Requires at least 8 GPUs to run, but got 0.

[wanchaol]

First pass, I think we can do things a lot simpler. please see inline comments

[wanchaol]

I see, this needs to be consolidate to a device mesh API, cc @wz337

[wanchaol]

Looks much better! wondering what's the reason to keep original.py and also some inline comments about imports, etc.

[wanchaol]

This looks great! Only have some nits for logging, thanks for addressing the comments!