[Ray SGD] use_local flag + Worker group abstraction (#10539)

Co-authored-by: Richard Liaw <rliaw@berkeley.edu>

python/ray/tune/integration/torch.py

@@ -75,7 +75,6 @@ class _TorchTrainable(tune.Trainable):
         remote_trainable = remote_trainable.options(
             **self.get_remote_worker_options())
 
-        address = setup_address()
         self.workers = [
             remote_trainable.remote(
                 config=config,
@@ -83,6 +82,10 @@ class _TorchTrainable(tune.Trainable):
             for rank in range(num_workers)
         ]
 
+        # Address has to be IP of rank 0 worker's node.
+        address = ray.get(
+            self.workers[0].execute.remote(lambda _: setup_address()))
+
         pgroup_params = self.default_process_group_parameters()
         from functools import partial
         setup_on_worker = partial(