Create a TPU instance

Create a TPU instance

This document shows you how to create a TPU v6e (Trillium) virtual machine (VM) instance, connect to it, and run a calculation on it. The TPU instance that you create in this quickstart uses on-demand, the default consumption option.

Before you begin

Sign in to your Google Cloud account. If you're new to Google Cloud, create an account to evaluate how our products perform in real-world scenarios. New customers also get $300 in free credits to run, test, and deploy workloads.

Install the Google Cloud CLI.

Note: If you installed the gcloud CLI previously, make sure you have the latest version by running gcloud components update.

If you're using an external identity provider (IdP), you must first sign in to the gcloud CLI with your federated identity.

To initialize the gcloud CLI, run the following command:

gcloud init

Create or select a Google Cloud project.

Roles required to select or create a project

• Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
• Create a project: To create a project, you need the Project Creator role (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

Note: If you don't plan to keep the resources that you create in this procedure, create a project instead of selecting an existing project. After you finish these steps, you can delete the project, removing all resources associated with the project.

• Create a Google Cloud project:

  gcloud projects create PROJECT_ID

  Replace PROJECT_ID with a name for the Google Cloud project you are creating.

• Select the Google Cloud project that you created:

  gcloud config set project PROJECT_ID

  Replace PROJECT_ID with your Google Cloud project name.

If you're using an existing project for this guide, verify that you have the permissions required to complete this guide. If you created a new project, then you already have the required permissions.

Verify that billing is enabled for your Google Cloud project.

Enable the Compute Engine API:

Roles required to enable APIs

To enable APIs, you need the Service Usage Admin IAM role (roles/serviceusage.serviceUsageAdmin), which contains the serviceusage.services.enable permission. Learn how to grant roles.

gcloud services enable compute.googleapis.com

Install the Google Cloud CLI.

Note: If you installed the gcloud CLI previously, make sure you have the latest version by running gcloud components update.

If you're using an external identity provider (IdP), you must first sign in to the gcloud CLI with your federated identity.

To initialize the gcloud CLI, run the following command:

gcloud init

Create or select a Google Cloud project.

Roles required to select or create a project

• Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
• Create a project: To create a project, you need the Project Creator role (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

Note: If you don't plan to keep the resources that you create in this procedure, create a project instead of selecting an existing project. After you finish these steps, you can delete the project, removing all resources associated with the project.

• Create a Google Cloud project:

  gcloud projects create PROJECT_ID

  Replace PROJECT_ID with a name for the Google Cloud project you are creating.

• Select the Google Cloud project that you created:

  gcloud config set project PROJECT_ID

  Replace PROJECT_ID with your Google Cloud project name.

If you're using an existing project for this guide, verify that you have the permissions required to complete this guide. If you created a new project, then you already have the required permissions.

Verify that billing is enabled for your Google Cloud project.

Enable the Compute Engine API:

Roles required to enable APIs

To enable APIs, you need the Service Usage Admin IAM role (roles/serviceusage.serviceUsageAdmin), which contains the serviceusage.services.enable permission. Learn how to grant roles.

gcloud services enable compute.googleapis.com

Required roles

To get the permissions that you need to create a TPU instance, connect to it using SSH, and run commands, ask your administrator to grant you the following IAM roles on your project:

• Compute Instance Admin (v1) (roles/compute.instanceAdmin.v1)
• Service Account User (roles/iam.serviceAccountUser)
• Service Usage Admin (roles/serviceusage.serviceUsageAdmin)

For more information about granting roles, see Manage access to projects, folders, and organizations.

You might also be able to get the required permissions through custom roles or other predefined roles.

Create a TPU v6e (Trillium) instance

Run the following commands in your local terminal or using Cloud Shell.

1. To create a TPU v6e instance, run the following gcloud compute instances create command:

   gcloud compute instances create quickstart-tpu \
       --machine-type=ct6e-standard-4t \
       --zone=europe-west4-a \
       --image-project=ubuntu-os-accelerator-images \
       --image-family=ubuntu-accel-2204-amd64-tpu-v5e-v5p-v6e \
       --maintenance-policy=TERMINATE
   

2. Optionally, verify that the instance is running by using the gcloud compute instances describe command:

   gcloud compute instances describe quickstart-tpu --zone=europe-west4-a
   

   In the output, check the status field to verify if the instance is running or still pending.

Connect to the TPU instance

You can connect to a TPU instance using the Google Cloud console or the gcloud CLI.

To connect to the VM using SSH, run the following gcloud compute ssh command:

gcloud compute ssh quickstart-tpu --zone=europe-west4-a

Install JAX

Install the JAX framework by running the following command on your TPU instance:

pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

Note: If this command fails, you might need to run sudo apt update and sudo apt install python3-pip.

Run a system check

Verify that JAX can access the TPU and can run basic operations by running the following commands on your TPU instance:

1. Start the Python 3 interpreter:

   python3
   

2. Import the JAX module:

   >>> import jax
   

3. Display the number of TPU cores available:

   >>> jax.device_count()
   

   The output shows the number of TPU cores. The number of cores is dependent on the TPU version you are using. For more information, see TPU versions.

Perform a calculation

1. Run the following command to add two numbers together:

   >>> jax.numpy.add(1, 1)
   

   The output should be similar to the following:

   >>> Array(2, dtype=int32, weak_type=True)
   

2. Exit the Python interpreter:

   >>> 
   exit()
   

Clean up

To avoid incurring charges to your Google Cloud account for the resources used on this page, delete the Google Cloud project with the resources.

To avoid incurring charges to your Google Cloud account for the resources used on this page, delete the TPU instance using the gcloud compute instances delete command:

gcloud compute instances delete quickstart-tpu --zone=europe-west4-a

What's next

• Learn about Cloud TPU resources in Compute Engine
• Create a Cloud TPU instance using other consumption options
• Create Cloud TPU instances with MIGs