I agree we need a better way to specify this. PyTorch 2.2 introduced the device mesh, so we should probably use that to specify the size, rather than having the user construct the process group matrix themselves.

Having an argument like you suggest could work, but it might be confusing to have this for anything other than hybrid sharding.

Since passing in a device mesh already works

from torch.distributed.device_mesh import init_device_mesh
mesh = init_device_mesh("cuda", (2, 4))

strategy = FSDPStrategy(..., device_mesh=mesh)

I suggest that we simplify this by allowing the user to set a tuple device_mesh=(2, 4) (in addition to DeviceMesh) and internally we initialize the device mesh for them if it's a tuple. Then we don't need to introduce a new argument.

Since passing in a device mesh already works

from torch.distributed.device_mesh import init_device_mesh
mesh = init_device_mesh("cuda", (2, 4))

strategy = FSDPStrategy(..., device_mesh=mesh)

I suggest that we simplify this by allowing the user to set a tuple device_mesh=(2, 4) (in addition to DeviceMesh) and internally we initialize the device mesh for them if it's a tuple. Then we don't need to introduce a new argument.

This seems reasonable as well, and simpler. In pytorch codeprocess_group and device_mesh end up being handled in the same function: https://github.com/pytorch/pytorch/blob/1d14adfa66e2ae437253eebe223710588648eee7/torch/distributed/fsdp/_init_utils.py#L152C5-L152C47

But pytorch FSDP doc does not document the device_mesh argument very well, and users would need to know what the numbers in device_mesh tuple means (which is the fsdp size and which is ddp size).

Great to hear you like it. Would you be interested to draft it? It would be relatively straightforward:

1. Store the device_mesh argument as attribute in the strategy
2. in Strategy setup initialize it and pass it to the FSDP wrapper:

   pytorch-lightning/src/lightning/fabric/strategies/fsdp.py

   Line 297 in a6273d1

   module = FullyShardedDataParallel(

But pytorch FSDP doc does not document the device_mesh argument very well, and users would need to know what the numbers in device_mesh tuple means (which is the fsdp size and which is ddp size).

Yes agreed. This is typical for PyTorch, their distributed features are always very short on docs. I think we would want to document this well on our side (both the API and the user guide). We have a relatively thorough guide already: https://lightning.ai/docs/fabric/stable/advanced/model_parallel/fsdp.html

There's a device_mesh recipe available at https://pytorch.org/tutorials/recipes/distributed_device_mesh.html

Great to hear you like it. Would you be interested to draft it? It would be relatively straightforward:

1. Store the device_mesh argument as attribute in the strategy
2. in Strategy setup initialize it and pass it to the FSDP wrapper:

   pytorch-lightning/src/lightning/fabric/strategies/fsdp.py

   Line 297 in a6273d1

   module = FullyShardedDataParallel(

But pytorch FSDP doc does not document the device_mesh argument very well, and users would need to know what the numbers in device_mesh tuple means (which is the fsdp size and which is ddp size).

Yes agreed. This is typical for PyTorch, their distributed features are always very short on docs. I think we would want to document this well on our side (both the API and the user guide). We have a relatively thorough guide already: https://lightning.ai/docs/fabric/stable/advanced/model_parallel/fsdp.html

Sure. I will iterate in the draft PR above, when I have some bandwidth.

I updated the PR #19504 as suggested.

#QuantumDominator