Monkeypatch Device Mesh ND Slicing

[mvpatel2000]

What does this PR do?

Monkeypatch Device Mesh ND Slicing. This allows us to create child meshes of multiple dimensions.

We patch in this PR from Torch pytorch/pytorch#119752 which will not be released until PyTorch 2.4

Toy script that now works which did not work before this PR:

from torch.distributed._tensor.device_mesh import DeviceMesh, init_device_mesh

from composer.utils import dist
from composer.trainer.mosaic_fsdp import patch_pytorch

dist.initialize_dist('gpu', timeout=10.0)

patch_pytorch()

child_mesh_dim_names = ("PP", "DP")
mesh_dim_names = ("PP", "DP", "TP")
mesh = init_device_mesh(
    'cuda', (2, 2, 2), mesh_dim_names=mesh_dim_names
)
child_mesh = mesh[child_mesh_dim_names]
print(child_mesh)

[snarayan21]

could you indicate which lines are different from torch 2.3 that we're patching? Or, since the torch PR seems to have made quite extensive changes, a comment with a link to the PR in torch would be useful. other than that, lgtm

[mvpatel2000]

could you indicate which lines are different from torch 2.3 that we're patching? Or, since the torch PR seems to have made quite extensive changes, a comment with a link to the PR in torch would be useful. other than that, lgtm

The PR is already in description. It is copy pasting the code from PR into torch -- the changes are equal to the linked PR

[snarayan21]

lgtm!