Run a calculation on a Cloud TPU VM using TensorFlow

Create a Cloud TPU VM with gcloud

Launch a Compute Engine Cloud TPU using the gcloud command. For more information on the gcloud command, see the gcloud reference.

When launching the TPU VM, you can either use the default TPU software version shown in the command or you can refer to Cloud TPU software versions to specify another supported TPU software version.

  $ gcloud compute tpus tpu-vm create tpu-name
    --zone=europe-west4-a
    --accelerator-type=v3-8
    --version=tpu-vm-tf-2.16.1-pjrt
  

Connect to your Cloud TPU VM

You must explicitly connect to your TPU VM using SSH. If you are not automatically connected, use the following command.

  $ gcloud compute tpus tpu-vm ssh tpu-name
      --zone europe-west4-a
   

Run an example using TensorFlow

Once you are connected to the TPU VM, set the following environment variable.

  (vm)$ export TPU_NAME=local
  

When creating your TPU, if you set the --version parameter to a version ending with -pjrt, set the following environment variables to enable the PJRT runtime:

  (vm)$ export NEXT_PLUGGABLE_DEVICE_USE_C_API=true
  (vm)$ export TF_PLUGGABLE_DEVICE_LIBRARY_PATH=/lib/libtpu.so

Create a file named tpu-test.pyin the current directory and copy and paste the following script into it.

import tensorflow as tf
print("Tensorflow version " + tf.__version__)

@tf.function
  def add_fn(x,y):
  z = x + y
  return z

  cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
  tf.config.experimental_connect_to_cluster(cluster_resolver)
  tf.tpu.experimental.initialize_tpu_system(cluster_resolver)
  strategy = tf.distribute.TPUStrategy(cluster_resolver)

  x = tf.constant(1.)
  y = tf.constant(1.)
  z = strategy.run(add_fn, args=(x,y))
  print(z)

Run this script with the following command:

(vm)$ python3 tpu-test.py

This script performs a computation on a each TensorCore of a TPU. The output will look similar to the following:

PerReplica:{
  0: tf.Tensor(2.0, shape=(), dtype=float32),
  1: tf.Tensor(2.0, shape=(), dtype=float32),
  2: tf.Tensor(2.0, shape=(), dtype=float32),
  3: tf.Tensor(2.0, shape=(), dtype=float32),
  4: tf.Tensor(2.0, shape=(), dtype=float32),
  5: tf.Tensor(2.0, shape=(), dtype=float32),
  6: tf.Tensor(2.0, shape=(), dtype=float32),
  7: tf.Tensor(2.0, shape=(), dtype=float32)
}