Troubleshooting JAX - TPU

This guide provides pointers to JAX troubleshooting information to help you identify and resolve problems you might encounter while training JAX models on Cloud TPU.

For a more general guide to getting started with Cloud TPU, see the JAX quickstart.

Note: If you aren't able to resolve your issue using this guide, see Getting Support for further assistance.

General JAX issues

If you run into issues while developing your training model or training with JAX, see the JAX FAQ.

For more general programming errors you might encounter when writing a training application with JAX, see JAX Errors.

Profiling JAX performance

You can understand how your TPU resources are being utilized using the tools described in Profiling JAX performance.

Troubleshooting memory issues

Note: TPU memory (high bandwidth memory is shared across each TPU chip. This enables efficient coordination between the two on-chip TensorCores.

You can monitor how the memory is used with the JAX Device Memory Profiler, but you cannot directly manage how it is used.

The Device Memory Profiler can be used to:

• Figure out which arrays and executables are in TPU memory at a given time, or
• Track down memory leaks.

You cannot specify how TPU memory is allocated for specific operations. For more information on JAX-specific TPU performance issues, see Performance Notes for using TPUs with JAX.

Troubleshooting TPU issues

How can I verify that the TPU is running?

Details

Everything will be run on the TPU as long as JAX doesn't print "No GPU/TPU found, falling back to CPU."

You can verify the TPU is active by either looking at jax.devices(), where you should see several TPU devices displayed, or verify programmatically with: assert jax.devices()[0].platform == 'tpu'.

RuntimeError: Unable to initialize backend 'tpu': UNAVAILABLE: No TPU Platform available.

Details

This runtime error message and/or finding the following in /tmp/tpu_logs/tpu_driver.WARNING on the TPU VM: W1118 17:40:20.985243 23901 tpu_version_flag.cc:57] No hardware is found. Using default TPU version:xxxxxx can indicate that you are running the wrong TPU VM version.

Verify that you are running the current JAX runtime version and retry.