Repository: microsoft/DeepSpeed

Issue:
[BUG] Deepspeed elastic launcher is incompatible with torch 2.0
**Describe the bug**

If I run `deepspeed` with `--elastic_training`, I got the following error:

```
Traceback (most recent call last):
  File "./lib/python3.9/site-packages/torch/distributed/elastic/metrics/api.py", line 129, in wrapper
    result = f(*args, **kwargs)
  File "./lib/python3.9/site-packages/torch/distributed/elastic/agent/server/api.py", line 723, in run
    result = self._invoke_run(role)
  File "./lib/python3.9/site-packages/deepspeed/elasticity/elastic_agent.py", line 126, in _invoke_run
    self._initialize_workers(self._worker_group)
  File "./lib/python3.9/site-packages/torch/distributed/elastic/metrics/api.py", line 129, in wrapper
    result = f(*args, **kwargs)
  File "./lib/python3.9/site-packages/torch/distributed/elastic/agent/server/api.py", line 692, in _initialize_workers
    self._rendezvous(worker_group)
  File "./lib/python3.9/site-packages/torch/distributed/elastic/metrics/api.py", line 129, in wrapper
    result = f(*args, **kwargs)
  File "./lib/python3.9/site-packages/torch/distributed/elastic/agent/server/api.py", line 556, in _rendezvous
    self._set_master_addr_port(
TypeError: _set_master_addr_port() takes 3 positional arguments but 4 were given
```

The torch repository's pull request [#88922](https://github.com/pytorch/pytorch/pull/88922) introduces an additional parameter called `local_addr`. You can find the specific code changes related to this parameter in [this section](https://github.com/pytorch/pytorch/commit/5d37890b8eb952560f22b1b8ff661a024bf4b6d9#diff-0332dff4db45fd26518da1eddd4a93eb965a006de743edddba2316b64e2b1f6cR507) of the commit. However, it appears that the deepspeed project is still using the outdated method and has not incorporated this new parameter yet.

**To Reproduce**

Run any script with `--elastic_training`

**System info (please complete the following information):**
 - Python 3.9
 - torch 2.0.1

**Launcher context**
`deepspeed` launcher




Discussion:
I want to learn about elastic training in DeepSpeed. I see the folder 'deepspeed/elasticity', but I can't find an introduction to this folder in the official tutorial. Can someone help? Thank you very much! @VeryLazyBoy @loadams 
Hi @Truring, as a fellow newcomer in this area, I suggest checking out the following resources for a quick understanding:

Code comments on elastic config: [GitHub link](https://github.com/microsoft/DeepSpeed/blob/a603a2130f63207c00b626c062b868ee90145994/deepspeed/elasticity/elasticity.py#L234)
Documentation on elastic config: [DeepSpeed docs](https://www.deepspeed.ai/docs/config-json/#elastic-training-config-v01-and-v02)

For DeepSpeed's elastic launcher, it relies on `torch.distributed.elastic`. You can refer to the documentation for torch at: [PyTorch docs](https://pytorch.org/docs/stable/distributed.elastic.html)
@VeryLazyBoy Thank you very much for sharing the resources! These materials are very useful to me.
I got the same error in my environment.

**Environment**

pytorch: 2.0.0+cu118
deepspeed: 0.13.0

"ds_config.json":
```
{
   "steps_per_print": 2000,
   "checkpoint": {
     "use_node_local_storage": true
   },
   "elasticity": {
     "enabled": true,
     "micro_batch_sizes": [256,512],
     "max_train_batch_size": 1024
   }
}
```
launching command:
```
deepspeed \
  --master_addr=rogpt1 \
  --elastic_training \
  --min_elastic_nodes=1 \
  --max_elastic_nodes=2 \
  --hostfile=hostfile \
  train.py \
  --deepspeed_config ds_config.json
```
** Error **
```
rogpt1:   File "/opt/conda/lib/python3.10/site-packages/torch/distributed/elastic/agent/server/api.py", line 556, in _rendezvous
rogpt1:     self._set_master_addr_port(
rogpt1: TypeError: DSElasticAgent._set_master_addr_port() takes 3 positional arguments but 4 were given
```

Resolving patch:
diff --git a/deepspeed/elasticity/elastic_agent.py b/deepspeed/elasticity/elastic_agent.py
--- a/deepspeed/elasticity/elastic_agent.py
+++ b/deepspeed/elasticity/elastic_agent.py
@@ -39,7 +39,10 @@ def __init__(
         self.ds_env = env
 
     @staticmethod
-    def _set_master_addr_port(store: Store, master_addr: Optional[str], master_port: Optional[int]):
+    def _set_master_addr_port(store: Store,
+                              master_addr: Optional[str],
+                              master_port: Optional[int],
+                              local_addr: Optional[str] = None):
         if master_port is None:
             sock = _get_socket_with_port()
             with closing(sock):


Test patch: