Do we need an N-dim sub-DeviceMesh?

[botbw]

🚀 The feature, motivation and pitch

Hey there! Currently torch.distributed._tensor.DeviceMesh only supports 1-D sub-meshes, will it be possible to manipulate it as an NDArray and generate N-dim sub-meshes?

For example, in 2-D tensor parallelism together with pipeline parallelism, the mesh looks like [pp, tp0, tp1] == [2, 2, 2], and if an all-gather/all-reduce on tp is needed for rank == 0, mesh[tp0], mesh[tp1] only gives [0, 1] and [0, 2] (3 is missing).

import os
import torch
from torch.distributed._tensor import DeviceMesh, mesh_resources
rank = int(os.environ['RANK'])
if __name__ == "__main__":
    mesh = DeviceMesh("cpu", [
        [
            [0, 1],
            [2, 3]
        ], # pp_rank == 0
        [
            [4, 5],
            [6, 7]
        ]  # pp_rank == 1
    ], mesh_dim_names=['pp', 'tp0', 'tp1'])
    if rank == 0: # pp_rank == tp0_rank == tp1_rank == 0
        print(mesh['tp0'])
        print(mesh['tp1'])

A workaround is to represent tp0 and tp1 in a single dim, but there could be scenario in which tp0 is replicated and only tp2 needs the communication.

Alternatively, could you please provide some workaround for this?

Alternatives

No response

Additional context

No response

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @penguinwu @fegin @XilunWu @wanchaol @fduwjj @wz337 @tianyu-l @wconstab @yf225 @chauhang @d4l3k @msaroufim

[botbw]

hey any updates on this?
@awgu @yf225 @wanchaol @wz337

[wz337]

@botbw Thanks for raising the issue. We are actually working on the feature and hoping to have it in nightly asap so we can get this in 2.4 release as well.