Usage of jax.distributed.initialize on HPC clusters

[shyams2]

Hi

I wanted to know if anyone has used jax.distributed.initialize to parallelize across GPUs on several hosts on HPC clusters. I'm not sure about the right way to do things - I'm especially confused about IP that needs to be used - do we just choose one of the machines as the main host? I'm usually used to submitting a script to slurm to run it across several nodes using MPI but not sure what would be done here.

Is there some other way to get things working on these clusters? I'd appreciate any advice/guidance on this front.

Thanks!

[shyams2]

Adding some more details to this, I tried running the provided example from the PR on a single node. However, I'm running into this error:

$ python3 nvidia_gpu_pjit.py --server_addr="172.11.92.39" --num_hosts=1 --host_idx=0
I0215 14:53:07.547920 22468454270784 distributed.py:49] Starting JAX distributed service on 172.11.92.39
E0215 14:53:07.551977037  475947 server_chttp2.cc:40]        {"created":"@1644954787.551954935","description":"No address added out of total 1 resolved","file":"external/com_github_grpc_grpc/src/core/ext/transport/chttp2/server/chttp2_server.cc","file_line":395,"referenced_errors":[{"created":"@1644954787.551953102","description":"Unable to configure socket","fd":7,"file":"external/com_github_grpc_grpc/src/core/lib/iomgr/tcp_server_utils_posix_common.cc","file_line":216,"referenced_errors":[{"created":"@1644954787.551951008","description":"Permission denied","errno":13,"file":"external/com_github_grpc_grpc/src/core/lib/iomgr/tcp_server_utils_posix_common.cc","file_line":189,"os_error":"Permission denied","syscall":"bind"}]}]}
Traceback (most recent call last):
  File "distributed_test.py", line 50, in <module>
    app.run(main)
  File "/home/user/.local/lib/python3.8/site-packages/absl/app.py", line 312, in run
    _run_main(main, args)
  File "/home/user/.local/lib/python3.8/site-packages/absl/app.py", line 258, in _run_main
    sys.exit(main(argv))
  File "distributed_test.py", line 25, in main
    jax.distributed.initialize(FLAGS.server_addr, FLAGS.num_hosts, FLAGS.host_idx)
  File "/home/user/.local/lib/python3.8/site-packages/jax/_src/distributed.py", line 50, in initialize
    _service = xla_extension.get_distributed_runtime_service(coordinator_address,
RuntimeError: UNKNOWN: Failed to start RPC server

Would I need admin rights to run this code?

[zhangqiaorjc]

Try providing the port too?

--server_addr="172.11.92.39:1234"

[shyams2]

Yes! It's working now! Thank you very much! 😄

[shyams2]

@zhangqiaorjc I got this working perfectly fine on one the HPC clusters I have access to. However, I was trying out jax.distributed on another machine, and I couldn't get it to work even on a single node. Here's what I'm seeing when I call jax.distributed.initialize:

2022-02-27 14:47:06.004990: E external/org_tensorflow/tensorflow/compiler/xla/pjrt/distributed/client.cc:177] Connect() failed after 1 retries in 0; most recent failure status: DEADLINE_EXCEEDED: Deadline Exceeded
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-6-9dfe32413fab> in <module>
----> 1 jax.distributed.initialize("10.230.2.199:1235", 1, 0)

~/.local/conda/2021-11-30/lib/python3.8/site-packages/jax/_src/distributed.py in initialize(coordinator_address, num_processes, process_id)
     54                                                         process_id)
     55   logging.info('Connecting to JAX distributed service on %s', coordinator_address)
---> 56   client.connect()
     57
     58   factory = functools.partial(xla_client.make_gpu_client, client, process_id)

RuntimeError: DEADLINE_EXCEEDED: Connect() timed out after 0 with 1 attempts. Most recent failure was: DEADLINE_EXCEEDED: Deadline Exceeded

Is there a reason I'm seeing this - it's pretty much the same python environment AFAIK. How can I go about fixing this? Or do you think this maybe something got to do with permissions on this machine?

Thanks

[shyams2]

Just adding on this in case anyone faces a similar issue - the proxy environment variables on the computer that you're working on may lead jax.distributed to fail. So make sure you unset proxy environment variables before you run your code.

[bhoov]

This was exactly my issue as well. I followed a stack overflow solution here, added the following function to my .bashrc, and run it on each node before I submit a multi-node job:

# Unset Proxy
function unsetproxy() {
    unset {http,https,ftp}_proxy
    unset {HTTP,HTTPS,FTP}_PROXY
}

[szxiangjn]

You save my day!!

[hawkinsp]

Someone should send a PR improving the docs!