Frequently Asked Questions - Cloud TPU

This document contains a list of frequently asked questions about Cloud TPUs. It is broken up into sections:

  1. Framework independent FAQs - questions about using Cloud TPUs regardless of what ML framework you are using.
  2. JAX FAQs - questions about using Cloud TPUs with JAX.
  3. PyTorch FAQs - questions about using Cloud TPUs with PyTorch.

Framework independent FAQs

How do I check which process is using the TPU on a Cloud TPU VM?

Run tpu-info on the Cloud TPU VM to print the process ID and other information about the process using the TPU. See supported metrics for the metrics and their corresponding definitions.

   tpu-info

The output from tpu-info is similar to the following:

   TPU Chips
   ┏━━━━━━━━━━━━━┳━━━━━━━━━━━━━┳━━━━━━━━━┳━━━━━━━━┓
   ┃ Chip        ┃ Type        ┃ Devices ┃ PID    ┃
   ┡━━━━━━━━━━━━━╇━━━━━━━━━━━━━╇━━━━━━━━━╇━━━━━━━━┩
   │ /dev/accel0 │ TPU v4 chip │ 1       │ 130007 │
   │ /dev/accel1 │ TPU v4 chip │ 1       │ 130007 │
   │ /dev/accel2 │ TPU v4 chip │ 1       │ 130007 │
   │ /dev/accel3 │ TPU