Repository: google/jax

Issue:
Support multinode training on GPU
I don't have a node with 8 gpus. I have two nodes each with 4 gpus. So is it possible to train a model on multiple nodes?



Discussion:
This *is* actually something that does work right now but it's still experimental. There's also no real public-facing API for it yet; you have to type in some obscure and fairly magical things to set it all up correctly.

We should polish it off and document it!
Can you say a bit more about your model, though? Would gradient all-reductions across multiple nodes suffice?
@hawkinsp Technically, I'm training a reformer model using Trax library.
And I assume you're just looking for data parallelism, i.e., partitioning a minibatch across GPUs, not partitioning in any other way (e.g., model parallelism)?
@hawkinsp yeah my concern is data parallelism 

@hawkinsp Can you please share your notes on this (don't need a stable api) ? We are trying some hybrid data/model/pipeline parallelism  so it is a little different from @py4 but would love to get started with data parallelism
Data parallelism would of value to other projects that use XLA as well (eg https://www.tensorflow.org/swift).  Exposing this functionality in a standardized way would help drive progress in the broader ecosystem!
> I don't have a node with 8 gpus. I have two nodes each with 4 gpus. So is it possible to train a model on multiple nodes?

Hello py4, I am meeting the same problem, have you found some solutions? 
> This _is_ actually something that does work right now but it's still experimental. There's also no real public-facing API for it yet; you have to type in some obscure and fairly magical things to set it all up correctly.
> 
> We should polish it off and document it!

Hello hawkinsp, Could you please provide more details about how to run data parallel with multi node GPUs?
@hawkinsp We are also interested in running JAX code on multiple nodes. Anything (hacky or not) that you can share would be appreciated. Thanks!
I really enjoyed Jax during my DM internship and wanted to use it on my university SLURM cluster, but the lack of a clear (official) data parallel (multi-node) solution is a huge blocker to increasing Jax adoption outside of Google where you cant just grab a TPU pod and `pmap` across the pod. A single 8 (GPU) replica setup can barely train a Resnet50 imagenet classifier. Training SimCLR or any other large SOTA model is currently impossible without multi-node data parallelism. 
I would love this feature! I enjoy Jax, but I've been largely using DeepSpeed due to its ability to distribute across clusters.
Any progress on this issue ? Using JAX to train a model on multi-node, multi-GPU is becoming a very important features for us.
@hawkinsp This is a significant bottleneck for scaling on multi-node GPU clusters. Is there any update on this issue? 
Also, there was a recent `pjit` [tutorial](https://jax.readthedocs.io/en/latest/jax-101/08-pjit.html) that explains multi-node TPU scaling but doesn't mention about GPUs. Is that planned to be updated in the future? 
@sudhakarsingh27 I constantly monitoring the jax releases, and there is something WIP that you might be interested in https://github.com/google/jax/pull/8364
See also: #9582
Yes indeed. We haven't advertised it that much yet, but (a) you need to initialize the cluster using that API, and (b) you need to follow the same rules of multi-host programming that also apply on TPU, documented here: https://jax.readthedocs.io/en/latest/multi_process.html

I suspect we can consider this issue closed when we've documented (a) in the document (b).
@hawkinsp @zhangqiaorjc 
Multinode (or multiprocess) doesn't seem to work with the following jax(lib) versions:
```
jax                           0.3.13                                                                                                                                                                       
jaxlib                        0.3.10+cuda11.cudnn82
```
Ran [the attached](https://github.com/google/jax/files/8744863/jax_multi_node_experiment.zip) minimal code on single node with 8 V100 GPUs as follows (2 processes with 4 GPUs each):

```
CUDA_VISIBLE_DEVICES="0,1,2,3" python jax_multi_node_experiment.py 0 &
CUDA_VISIBLE_DEVICES="4,5,6,7" python jax_multi_node_experiment.py 1
```
---
I could check that multi process(host/node) first fails with `jax[cuda]==0.3.12` installed with following command

```
pip install jax[cuda]==0.3.12 -f https://storage.googleapis.com/jax-releases/jax_releases.html
```
I get the following error when I run the multi-process jax commands above:
```
127.0.0.1:65432 2 1
I0525 00:05:16.228919 139978761119552 distributed.py:59] Connecting to JAX distributed service on 127.0.0.1:65432
I0525 00:05:16.245648 139978761119552 xla_bridge.py:330] Unable to initialize backend 'tpu_driver': NOT_FOUND: Unable to find driver in registry given worker: 
I0525 00:05:16.246569 139742444975936 xla_bridge.py:330] Unable to initialize backend 'tpu_driver': NOT_FOUND: Unable to find driver in registry given worker: 
I0525 00:05:18.227763 139978761119552 xla_bridge.py:330] Unable to initialize backend 'tpu': INVALID_ARGUMENT: TpuPlatform is not available.
I0525 00:05:18.228022 139978761119552 xla_bridge.py:330] Unable to initialize backend 'cuda': make_gpu_client() got an unexpected keyword argument 'platform_name'
I0525 00:05:18.228085 139978761119552 xla_bridge.py:330] Unable to initialize backend 'rocm': make_gpu_client() got an unexpected keyword argument 'platform_name'
global devices= [GpuDevice(id=0, process_index=0), GpuDevice(id=1, process_index=0), GpuDevice(id=2, process_index=0), GpuDevice(id=3, process_index=0)]
local devices= [GpuDevice(id=0, process_index=0), GpuDevice(id=1, process_index=0), GpuDevice(id=2, process_index=0), GpuDevice(id=3, process_index=0)]
I0525 00:05:18.246024 139742444975936 xla_bridge.py:330] Unable to initialize backend 'tpu': INVALID_ARGUMENT: TpuPlatform is not available.
I0525 00:05:18.246273 139742444975936 xla_bridge.py:330] Unable to initialize backend 'cuda': make_gpu_client() got an unexpected keyword argument 'platform_name'
I0525 00:05:18.246334 139742444975936 xla_bridge.py:330] Unable to initialize backend 'rocm': make_gpu_client() got an unexpected keyword argument 'platform_name'
global devices= [GpuDevice(id=0, process_index=0), GpuDevice(id=1, process_index=0), GpuDevice(id=2, process_index=0), GpuDevice(id=3, process_index=0)]
local devices= [GpuDevice(id=0, process_index=0), GpuDevice(id=1, process_index=0), GpuDevice(id=2, process_index=0), GpuDevice(id=3, process_index=0)]

```

For reference, here's the ouput from `jax[cuda]==0.3.10` where multi-process seems to be working okay:
```
127.0.0.1:65432 2 1
I0525 00:09:03.394093 140366043674432 distributed.py:59] Connecting to JAX distributed service on 127.0.0.1:65432
I0525 00:09:03.410755 140366043674432 xla_bridge.py:263] Unable to initialize backend 'tpu_driver': NOT_FOUND: Unable to find driver in registry given worker: 
I0525 00:09:03.410994 140588577851200 xla_bridge.py:263] Unable to initialize backend 'tpu_driver': NOT_FOUND: Unable to find driver in registry given worker: 
I0525 00:09:05.517608 140366043674432 xla_bridge.py:263] Unable to initialize backend 'tpu': INVALID_ARGUMENT: TpuPlatform is not available.
global devices= [GpuDevice(id=0, process_index=0), GpuDevice(id=1, process_index=0), GpuDevice(id=2, process_index=0), GpuDevice(id=3, process_index=0), GpuDevice(id=4, process_index=1), GpuDevice(id=5, process_index=1), GpuDevice(id=6, process_index=1), GpuDevice(id=7, process_index=1)]
I0525 00:09:05.517817 140588577851200 xla_bridge.py:263] Unable to initialize backend 'tpu': INVALID_ARGUMENT: TpuPlatform is not available.
local devices= [GpuDevice(id=4, process_index=1), GpuDevice(id=5, process_index=1), GpuDevice(id=6, process_index=1), GpuDevice(id=7, process_index=1)]
global devices= [GpuDevice(id=0, process_index=0), GpuDevice(id=1, process_index=0), GpuDevice(id=2, process_index=0), GpuDevice(id=3, process_index=0), GpuDevice(id=4, process_index=1), GpuDevice(id=5, process_index=1), GpuDevice(id=6, process_index=1), GpuDevice(id=7, process_index=1)]
local devices= [GpuDevice(id=0, process_index=0), GpuDevice(id=1, process_index=0), GpuDevice(id=2, process_index=0), GpuDevice(id=3, process_index=0)]

```

Resolving patch:
diff --git a/jax/_src/distributed.py b/jax/_src/distributed.py
--- a/jax/_src/distributed.py
+++ b/jax/_src/distributed.py
@@ -102,45 +102,53 @@ def initialize_preemption_sync_manager(self):
 def initialize(coordinator_address: Optional[str] = None,
                num_processes: Optional[int] = None,
                process_id: Optional[int] = None):
-  """Initialize distributed system for topology discovery.
+  """Initializes the JAX distributed system.
 
-  Currently, calling ``initialize`` sets up the multi-host GPU backend and Cloud
-  TPU backend.
+  Calling :func:`~jax.distributed.initialize` prepares JAX for execution on
+  multi-host GPU and Cloud TPU. :func:`~jax.distributed.initialize` must be
+  called before performing any JAX computations.
 
-  If you are on GPU platform, you will have to provide the coordinator_address
-  and other args to the `initialize` API.
+  The JAX distributed system serves a number of roles:
 
-  If you are on TPU platform, the coordinator_address and other args will be
-  auto detected but you have the option to provide it too.
+    * it allows JAX processes to discover each other and share topology information,
+    * it performs health checking, ensuring that all processes shut down if any process dies, and
+    * it is used for distributed checkpointing.
+
+  If you are using GPU, you must provide the ``coordinator_address``,
+  ``num_processes``, and ``process_id`` arguments to :func:`~jax.distributed.initialize`.
+
+  If you are using TPU, all arguments are optional: if omitted, they
+  will be chosen automatically from the Cloud TPU metadata.
 
   Args:
-    coordinator_address: IP address and port of the coordinator. The choice of
+    coordinator_address: the IP address of process `0` and a port on which that
+      process should launch a coordinator service. The choice of
       port does not matter, so long as the port is available on the coordinator
       and all processes agree on the port.
-      Can be None only for TPU platform. If coordinator_address is None on TPU,
-      then it will be auto detected.
-    num_processes: Number of processes. Can be None only for TPU platform and
-      if None will be determined from the TPU slice metadata.
-    process_id: Id of the current process. Can be None only for TPU platform and
-      if None will default to the current TPU worker id determined via the TPU
-      slice metadata.
+      May be ``None`` only on TPU, in which case it will be chosen automatically.
+    num_processes: Number of processes. May be ``None`` only on TPU, in
+      which case it will be chosen automatically based on the TPU slice.
+    process_id: The ID number of the current process. The ``process_id`` values across
+      the cluster must be a dense range ``0``, ``1``, ..., ``num_processes - 1``.
+      May be ``None`` only on TPU; if ``None`` it will be chosen from the TPU slice
+      metadata.
 
   Raises:
-    RuntimeError: If `distributed.initialize` is called more than once.
+    RuntimeError: If :func:`~jax.distributed.initialize` is called more than once.
 
   Example:
 
-  Suppose there are two GPU hosts, and host 0 is the designated coordinator
+  Suppose there are two GPU processs, and process 0 is the designated coordinator
   with address ``10.0.0.1:1234``. To initialize the GPU cluster, run the
   following commands before anything else.
 
-  On host 0:
+  On process 0:
 
-  >>> jax.distributed.initialize('10.0.0.1:1234', 2, 0)  # doctest: +SKIP
+  >>> jax.distributed.initialize(coordinator_address='10.0.0.1:1234', num_processes=2, process_id=0)  # doctest: +SKIP
 
-  On host 1:
+  On process 1:
 
-  >>> jax.distributed.initialize('10.0.0.1:1234', 2, 1)  # doctest: +SKIP
+  >>> jax.distributed.initialize(coordinator_address='10.0.0.1:1234', num_processes=2, process_id=1)  # doctest: +SKIP
   """
   global_state.initialize(coordinator_address, num_processes, process_id)
   atexit.register(shutdown)


Test patch: