multi-host tpu on jax

[lucasliunju]

Hi @mattjj and @skye,

It would be better to open an new disscussion.

I try to connect jax and tpu. It works when I use v3-8 version of tpu. My definition on the python code is :

import requests
import os
url = f'http://[ip address of tpu]:8475/requestversion/tpu_driver_nightly'
requests.post(url)
TPU_DRIVER_MODE = 1

from jax.config import config
# The following is required to use TPU Driver as JAX's backend.
config.FLAGS.jax_xla_backend = "tpu_driver"
config.FLAGS.jax_backend_target = "grpc://[ip address]:8470"

However, when I try to connect jax and multi-host tpu (such as v3-128) I find it cannot work when I try to :

print(jax.devices())

The error is :

RuntimeError: Deadline exceeded: Failed to connect to remote server at address: grpc://10.59.106.169:8470. Error fr
om gRPC: Deadline Exceeded. Details:

Currently, I find many great work (vision transformer, sam) is based on jax and multi-host tpu. Therefore, I would like to reproduce these work and learn to use jax. During this process, I think jax is very intertsting and will become more popular and gain more attention.

Thank you very much!

I am looking forward to your reply!

Best regards,
Yong

[skye]

Let's discuss in #6157