jax distributed.initialize support tpu?

[JiahaoYao]

recently, i found https://github.com/google/jax/blob/main/jax/_src/distributed.py#L35-L36 adds the tpu in jax distributed.initialize. Does it currently support multiple tpu hosts? And I also wonder how to install jax nightly in order to keep updated with this features?

Here are the details: I have launched 2 tpu:v2-8 vms.

>>> import jax
>>> jax.devices()
[TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0), TpuDevice(id=1, process_index=0, coords=(0,0,0), core_on_chip=1), TpuDevice(id=2, process_index=0, coords=(1,0,0), core_on_chip=0), TpuDevice(id=3, process_index=0, coords=(1,0,0), core_on_chip=1), TpuDevice(id=4, process_index=0, coords=(0,1,0), core_on_chip=0), TpuDevice(id=5, process_index=0, coords=(0,1,0), core_on_chip=1), TpuDevice(id=6, process_index=0, coords=(1,1,0), core_on_chip=0), TpuDevice(id=7, process_index=0, coords=(1,1,0), core_on_chip=1)]

How can I use distributed.initialize to connected the two hosts?

[JiahaoYao]

Thanks @jakevdp for editting. I build latest tpu version from the source, and run the nvidia_gpu_pjit.py (link thanks @zhangqiaorjc ) for the test collective behavior.

python3 nvidia_gpu_pjit.py --server_addr="10.128.0.25:1456" --num_hosts=2 --host_idx=0
python3 nvidia_gpu_pjit.py --server_addr="10.128.0.25:1456" --num_hosts=2 --host_idx=1

And here are the results

ubuntu@t1v-n-207dbaa4-w-0:~$ python3 nvidia_gpu_pjit.py --server_addr="10.128.0.25:1456" --num_hosts=2 --host_idx=0
['t1v-n-207dbaa4-w-0:2664546204052242694:10.128.0.25']
Traceback (most recent call last):
  File "nvidia_gpu_pjit.py", line 50, in <module>
    app.run(main)
  File "/home/ubuntu/.local/lib/python3.8/site-packages/absl/app.py", line 312, in run
    _run_main(main, args)
  File "/home/ubuntu/.local/lib/python3.8/site-packages/absl/app.py", line 258, in _run_main
    sys.exit(main(argv))
  File "nvidia_gpu_pjit.py", line 25, in main
    jax.distributed.initialize(FLAGS.server_addr, FLAGS.num_hosts, FLAGS.host_idx)
  File "/home/ubuntu/.local/lib/python3.8/site-packages/jax/_src/distributed.py", line 89, in initialize
    raise RuntimeError('Number of workers does not equal the number of '
RuntimeError: Number of workers does not equal the number of processes. Auto detecting process_id is not possible.Please pass process_id manually.

and the other host

ubuntu@t1v-n-78ea0c76-w-0:~$ python3 nvidia_gpu_pjit.py --server_addr="10.128.0.25:1456" --num_hosts=2 --host_idx=1
['t1v-n-78ea0c76-w-0:8021122269074014007:10.128.0.26']
Traceback (most recent call last):
  File "nvidia_gpu_pjit.py", line 50, in <module>
    app.run(main)
  File "/home/ubuntu/.local/lib/python3.8/site-packages/absl/app.py", line 312, in run
    _run_main(main, args)
  File "/home/ubuntu/.local/lib/python3.8/site-packages/absl/app.py", line 258, in _run_main
    sys.exit(main(argv))
  File "nvidia_gpu_pjit.py", line 25, in main
    jax.distributed.initialize(FLAGS.server_addr, FLAGS.num_hosts, FLAGS.host_idx)
  File "/home/ubuntu/.local/lib/python3.8/site-packages/jax/_src/distributed.py", line 89, in initialize
    raise RuntimeError('Number of workers does not equal the number of '
RuntimeError: Number of workers does not equal the number of processes. Auto detecting process_id is not possible.Please pass process_id manually.

[JiahaoYao]

So, is there any way to connect multiple v2-8 tpu instances?

[ayaka14732]

I believe that this is not possible. Multi-host initialisation only works for TPU pods, like TPU v3-16, v4-16, etc.

[ayaka14732]

Update: I believe that this feature is not available because TPU Pods are interconnected with high-speed cables, which is different from the normal net cables.

[hawkinsp]

A relevant new announcement: search for "multislice" in https://cloud.google.com/blog/products/compute/announcing-cloud-tpu-v5e-and-a3-gpus-in-ga