How do I connect to TPU inside a GCP AI Platform Notebook?

[sayakpaul]

Hi.

I am trying to run the experiments from the official JAX-based implementation of the Vision Transformer. The provided Colab Notebook is only compatible with the TPUs provided by the Colab runtime as obvious. The following code listing is used to configure JAX's backend to run on TPUs:

import os
if 'google.colab' in str(get_ipython()) and 'COLAB_TPU_ADDR' in os.environ:
  # Make sure the Colab Runtime is set to Accelerator: TPU.
  import requests
  if 'TPU_DRIVER_MODE' not in globals():
    url = 'http://' + os.environ['COLAB_TPU_ADDR'].split(':')[0] + ':8475/requestversion/tpu_driver0.1-dev20191206'
    resp = requests.post(url)
    TPU_DRIVER_MODE = 1

  # The following is required to use TPU Driver as JAX's backend.
  from jax.config import config
  config.FLAGS.jax_xla_backend = "tpu_driver"
  config.FLAGS.jax_backend_target = "grpc://" + os.environ['COLAB_TPU_ADDR']
  print('Registered TPU:', config.FLAGS.jax_backend_target)
else:
  print('No TPU detected. Can be changed under "Runtime/Change runtime type".')

It seems like that it cannot be extended to support TPUs in other runtimes such as in an AI Platform Notebook that is attached to a TPU (v3-8).

So, was wondering if there is a way to achieve this.

[lucasliunju]

Hi I'm trying to run the vision transformer code on colab tpu and I find there is an error when I run the part of Fine Tune.

That is this part:

for step, batch, lr_repl in zip(
tqdm.notebook.trange(1, total_steps + 1),
ds_train.as_numpy_iterator(),
lr_iter
):

opt_repl, loss_repl, update_rngs = update_fn_repl(
opt_repl, lr_repl, batch, update_rngs)

In addition, I also try to run this code on my own TPU (V3-8). It also has an error in this part and the error is segmentation fault.

That problem confused me several days and I'm not sure whether you have this question.

Thank you very much!

Have a nice weekend.

[lucasliunju]

For your question, I searched it and that maybe work:

from jax.config import config
config.FLAGS.jax_xla_backend = "tpu_driver"
config.FLAGS.jax_backend_target = "grpc://[your tpu ip]:8470"

Yong