Repository: Lightning-AI/pytorch-lightning
Pull request: Add timeout to DeepSpeedStrategy
## 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 #20466. allow users to pass `kwargs` to `DeepSpeedStrategy`. This allows users to set timeout when using `DeepSpeedStrategy`.

<!-- Does your PR introduce any breaking changes? If yes, please list them. -->

<details>
  <summary><b>Before submitting</b></summary>

- Was this **discussed/agreed** via a GitHub issue? (not for typos and docs)
- [x] Did you read the [contributor guideline](https://github.com/Lightning-AI/lightning/blob/master/.github/CONTRIBUTING.md), **Pull Request** section?
- [x] 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>

- [x] Is this pull request ready for review? (if not, please submit in draft mode)
- [x] Check that all items from **Before submitting** are resolved
- [x] 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--20474.org.readthedocs.build/en/20474/

<!-- readthedocs-preview pytorch-lightning end -->

Problem statement:
Allowing setting timeout in DeepSpeedStrategy
### Outline & Motivation

In DDPStrategy / FSDPStrategy, the `timeout=datetime.timedelta(seconds=1800)` flag is exposed and thus allowing user to tune. However, in DeepSpeedStrategy, which is a subclass of DDPStrategy, this flag is not exposed, which makes it hard to change the timeout behavior.

Is there any workaround? Otherwise, I think it might be worth adding `kwargs` to the `__init__()` function of DeepSpeedStrategy, and pass along those parameters to the parent class DDPStrategy.

### Pitch

_No response_

### Additional context

_No response_

cc @borda @awaelchli @justusschock

Hints from issue discussion:
Hey @jedyang97 definitely, do you want to send in a quick PR?

Commit message: [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Patch to review:
diff --git a/src/lightning/fabric/strategies/deepspeed.py b/src/lightning/fabric/strategies/deepspeed.py
index 03d90cd5df057..0bb5be4d8adbb 100644
--- a/src/lightning/fabric/strategies/deepspeed.py
+++ b/src/lightning/fabric/strategies/deepspeed.py
@@ -97,6 +97,7 @@ def __init__(
         load_full_weights: bool = False,
         precision: Optional[Precision] = None,
         process_group_backend: Optional[str] = None,
+        **kwargs: Any,
     ) -> None:
         """Provides capabilities to run training using the DeepSpeed library, with training optimizations for large
         billion parameter models. `For more information: https://pytorch-
@@ -239,6 +240,7 @@ def __init__(
             cluster_environment=cluster_environment,
             precision=precision,
             process_group_backend=process_group_backend,
+            **kwargs,
         )
         self._backward_sync_control = None  # DeepSpeed handles gradient accumulation internally
 
@@ -648,7 +650,9 @@ def _init_deepspeed_distributed(self) -> None:
                 f"MEMBER: {self.global_rank + 1}/{self.world_size}"
             )
         self._process_group_backend = self._get_process_group_backend()
-        deepspeed.init_distributed(self._process_group_backend, distributed_port=self.cluster_environment.main_port)
+        deepspeed.init_distributed(
+            self._process_group_backend, distributed_port=self.cluster_environment.main_port, timeout=self._timeout
+        )
 
     def _set_node_environment_variables(self) -> None:
         assert self.cluster_environment is not None
diff --git a/src/lightning/pytorch/strategies/deepspeed.py b/src/lightning/pytorch/strategies/deepspeed.py
index 4fa771114768d..456d6428cbe87 100644
--- a/src/lightning/pytorch/strategies/deepspeed.py
+++ b/src/lightning/pytorch/strategies/deepspeed.py
@@ -119,6 +119,7 @@ def __init__(
         load_full_weights: bool = False,
         precision_plugin: Optional[Precision] = None,
         process_group_backend: Optional[str] = None,
+        **kwargs: Any,
     ) -> None:
         """Provides capabilities to run training using the DeepSpeed library, with training optimizations for large
         billion parameter models. `For more information: https://pytorch-
@@ -263,6 +264,7 @@ def __init__(
             cluster_environment=cluster_environment,
             precision_plugin=precision_plugin,
             process_group_backend=process_group_backend,
+            **kwargs,
         )
 
         self.config = self._load_config(config)
@@ -364,7 +366,9 @@ def _init_deepspeed_distributed(self) -> None:
                 f"MEMBER: {self.global_rank + 1}/{self.world_size}"
             )
         self._process_group_backend = self._get_process_group_backend()
-        deepspeed.init_distributed(self._process_group_backend, distributed_port=self.cluster_environment.main_port)
+        deepspeed.init_distributed(
+            self._process_group_backend, distributed_port=self.cluster_environment.main_port, timeout=self._timeout
+        )
 
     def _set_node_environment_variables(self) -> None:
         assert self.cluster_environment is not None


Code review comments:
diff_hunk: @@ -97,6 +97,7 @@ def __init__(
         load_full_weights: bool = False,
         precision: Optional[Precision] = None,
         process_group_backend: Optional[str] = None,
+        **kwargs: Any,
path: src/lightning/fabric/strategies/deepspeed.py
text: @user1:
The other strategies have `timeout` as an explicit argument (https://github.com/Lightning-AI/pytorch-lightning/blob/be608fa355b835b9b0727df2f5476f0a1d90bc59/src/lightning/fabric/strategies/fsdp.py#L144).

I strongly suggest we do that here as well, adding `kwargs` hampers discoverability.

Let's keep the pattern exactly the same as FSDP.

@author:
That makes sense. I've updated the PR.

@user1:
looks good thanks!