Repository: Lightning-AI/pytorch-lightning
Pull request: Flexible and easy to use HSDP setting
## What does this PR do?

<!--
Please include a summary of the change and which issue is fixed.
Please also include relevant motivation and context.
List any dependencies that are required for this change.

If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

The following links the related issue to the PR (https://docs.github.com/en/free-pro-team@latest/github/managing-your-work-on-github/linking-a-pull-request-to-an-issue#linking-a-pull-request-to-an-issue-using-a-keyword)
-->

Fixes #19502

It allows users to provide `device_mesh` as a tuple to the `FSDPStrategy`.

<!-- Does your PR introduce any breaking changes? If yes, please list them. -->

<del>The test case for hybrid fsdp strategy would be broken, because `_init_sharding_strategy` is moved out of the `__init__` function.

<details>
  <summary><b>Before submitting</b></summary>

- Was this **discussed/agreed** via a GitHub issue? (not for typos and docs)
- [ ] Did you read the [contributor guideline](https://github.com/Lightning-AI/lightning/blob/master/.github/CONTRIBUTING.md), **Pull Request** section?
- [ ] Did you make sure your **PR does only one thing**, instead of bundling different changes together?
- Did you make sure to **update the documentation** with your changes? (if necessary)
- Did you write any **new necessary tests**? (not for typos and docs)
- [ ] Did you verify new and **existing tests pass** locally with your changes?
- Did you list all the **breaking changes** introduced by this pull request?
- Did you **update the CHANGELOG**? (not for typos, docs, test updates, or minor internal changes/refactors)

<!-- In the CHANGELOG, separate each item in the unreleased section by a blank line to reduce collisions -->

</details>

## PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing, make sure you have read the [review guidelines](https://github.com/Lightning-AI/lightning/wiki/Review-guidelines). In short, see the following bullet-list:

<details>
  <summary>Reviewer checklist</summary>

- [ ] Is this pull request ready for review? (if not, please submit in draft mode)
- [ ] Check that all items from **Before submitting** are resolved
- [ ] Make sure the title is self-explanatory and the description concisely explains the PR
- [ ] Add labels and milestones (and optionally projects) to the PR so it can be classified

</details>

<!--

Did you have fun?

Make sure you had fun coding 🙃

-->


<!-- readthedocs-preview pytorch-lightning start -->
----
📚 Documentation preview 📚: https://pytorch-lightning--19504.org.readthedocs.build/en/19504/

<!-- readthedocs-preview pytorch-lightning end -->

Problem statement:
Allow flexible and easy to configure HSDP
### Description & Motivation

The `FSDPStrategy` can use hybrid sharding strategy to shard across smaller sets of ranks in the global dist group. However, it is not flexible enough to let user easily specify the sharding scale.


### Pitch

The `FSDPStrategy` can use hybrid sharding strategy to shard across smaller sets of ranks in the global dist group. Currently there are two path to use it in Lightning:
1. Specify `sharding_strategy` as one of the hybrid sharding strategies. This will shard within one node, and replicate across nodes.
2. Specify `sharding_strategy` as one of the hybrid sharding strategies, and provide `process_group` as `kwards` to `FSDPStrategy`. This let user specify how large the sharding scale is. However, it is not easy for user to insert torch dist groups creation code and prepare the `process_group` ahead of time, because Lightning handles torch dist init_process_group automatically in trainer, or the fabric launcher.

So I'm looking forward to a easier way to use HSDP within Lightning, like:
```FSDPStrategy(..., sharding_strategy="HYBRID_SHARD", fsdp_size=16)```
to easily shard at specified scale, and let Lightning handle `process_group` preparation for PyTorch FSPD wrapper.

### Alternatives

_No response_

### Additional context

_No response_

cc @borda @awaelchli @carmocca

Hints from issue discussion:
I agree we need a better way to specify this. PyTorch 2.2 introduced the [device mesh](https://pytorch.org/tutorials/recipes/distributed_device_mesh.html), 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 

```py
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
> 
> ```python
> 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 code`process_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: https://github.com/Lightning-AI/pytorch-lightning/blob/a6273d17870336886432d7f79900a07fe7f872a6/src/lightning/fabric/strategies/fsdp.py#L297

> 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: https://github.com/Lightning-AI/pytorch-lightning/blob/a6273d17870336886432d7f79900a07fe7f872a6/src/lightning/fabric/strategies/fsdp.py#L297
> 
> > 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 https://github.com/Lightning-AI/pytorch-lightning/pull/19504 as suggested.

#QuantumDominator

Commit message: add simple docs

Patch to review:
diff --git a/src/lightning/fabric/strategies/fsdp.py b/src/lightning/fabric/strategies/fsdp.py
index 5679155380bea..fb4f033d61398 100644
--- a/src/lightning/fabric/strategies/fsdp.py
+++ b/src/lightning/fabric/strategies/fsdp.py
@@ -125,10 +125,14 @@ class FSDPStrategy(ParallelStrategy, _Sharded):
             - ``"SHARD_GRAD_OP"``: Shards gradients and optimizer states only. Model parameters get replicated.
             - ``"NO_SHARD"``: No sharding (identical to regular DDP).
             - ``"HYBRID_SHARD"``: Shards model parameters, gradients, and optimizer states within a single machine, but
-              replicates across machines.
+              replicates across machines. See also the `device_mesh` parameter below.
 
             Also accepts a :class:`torch.distributed.fsdp.ShardingStrategy` enum value.
 
+        device_mesh: A tuple `(sharding size, replication size)` that defines over how many devices to shard and
+            replicate the model. The product of the two numbers must equal the world size. Only valid in combination
+            with the `HYBRID_SHARD` sharding strategy.
+
         state_dict_type: The format in which the state of the model and optimizers gets saved into the checkpoint.
 
             - ``"full"``: The full weights and optimizer states get assembled on rank 0 and saved to a single file.
@@ -253,6 +257,12 @@ def setup_environment(self) -> None:
         super().setup_environment()
         self._setup_distributed()
 
+        # if 'device_mesh' in the `_fsdp_kwargs` is provided as a tuple, update it into the `DeviceMesh` object here
+        if isinstance(self._fsdp_kwargs.get("device_mesh"), tuple):
+            from torch.distributed.device_mesh import init_device_mesh
+
+            self._fsdp_kwargs["device_mesh"] = init_device_mesh("cuda", self._fsdp_kwargs["device_mesh"])
+
     @override
     def setup_module_and_optimizers(
         self, module: Module, optimizers: List[Optimizer]
diff --git a/src/lightning/pytorch/strategies/fsdp.py b/src/lightning/pytorch/strategies/fsdp.py
index 9887e89df9723..55f8eb5afc011 100644
--- a/src/lightning/pytorch/strategies/fsdp.py
+++ b/src/lightning/pytorch/strategies/fsdp.py
@@ -125,10 +125,14 @@ class FSDPStrategy(ParallelStrategy):
             - ``"SHARD_GRAD_OP"``: Shards gradients and optimizer states only. Model parameters get replicated.
             - ``"NO_SHARD"``: No sharding (identical to regular DDP).
             - ``"HYBRID_SHARD"``: Shards model parameters, gradients, and optimizer states within a single machine, but
-              replicates across machines.
+              replicates across machines. See also the `device_mesh` parameter below.
 
             Also accepts a :class:`torch.distributed.fsdp.ShardingStrategy` enum value.
 
+        device_mesh: A tuple `(sharding size, replication size)` that defines over how many devices to shard and
+            replicate the model. The product of the two numbers must equal the world size. Only valid in combination
+            with the `HYBRID_SHARD` sharding strategy.
+
         state_dict_type: The format in which the state of the model and optimizers gets saved into the checkpoint.
 
             - ``"full"``: The full weights and optimizer states get assembled on rank 0 and saved to a single file.
@@ -260,6 +264,12 @@ def setup_environment(self) -> None:
         assert self.cluster_environment is not None
         _init_dist_connection(self.cluster_environment, self._process_group_backend, timeout=self._timeout)
 
+        # if 'device_mesh' in the `kwargs` is provided as a tuple, update it into the `DeviceMesh` object here
+        if isinstance(self.kwargs.get("device_mesh"), tuple):
+            from torch.distributed.device_mesh import init_device_mesh
+
+            self.kwargs["device_mesh"] = init_device_mesh("cuda", self.kwargs["device_mesh"])
+
     def _get_process_group_backend(self) -> str:
         return self._process_group_backend or _get_default_process_group_backend_for_device(self.root_device)
 


Code review comments:
diff_hunk: @@ -125,10 +125,14 @@ class FSDPStrategy(ParallelStrategy, _Sharded):
             - ``"SHARD_GRAD_OP"``: Shards gradients and optimizer states only. Model parameters get replicated.
             - ``"NO_SHARD"``: No sharding (identical to regular DDP).
             - ``"HYBRID_SHARD"``: Shards model parameters, gradients, and optimizer states within a single machine, but
-              replicates across machines.
+              replicates across machines. See also the `device_mesh` parameter below.
 
             Also accepts a :class:`torch.distributed.fsdp.ShardingStrategy` enum value.
 
+        device_mesh: A tuple `(sharding size, replication size)` that defines over how many devices to shard and
path: src/lightning/fabric/strategies/fsdp.py
text: @user1:
@author Want to double check that the tuple ordering is correct. Looking at https://github.com/pytorch/pytorch/blob/main/torch/distributed/fsdp/_init_utils.py#L162-L163, it looks like the first dimension is for the inter pg and the second dimension is for the intra pg. By this logic, should it be "A tuple `(replication_size, sharding_size)`..." (i.e, swap the ordering of the tuple elements?

@author:
Right. I have it corrected in my code, but haven't pushed it.