Create a MIG with multi-host TPU slice

This document describes how to create a managed instance group (MIG) that forms a multi-host TPU slice.

Before you begin

  • Review the limitations for creating MIGs with TPU instances.
  • If you haven't already, set up authentication. Authentication verifies your identity for access to Google Cloud services and APIs. To run code or samples from a local development environment, you can authenticate to Compute Engine by selecting one of the following options:

    1. Install the Google Cloud CLI. After installation, initialize the Google Cloud CLI by running the following command: 

      gcloud init

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

    2. Set a default region and zone.

Prerequisites

Before you proceed to create a multi-host TPU slice, you must do the following:

  1. Choose the TPU version: select the TPU version that is suitable for your workload. For a list of TPU versions by workload type, see Recommended TPU versions by workload types.

  2. Validate TPU availability in your preferred location: TPUs are available in specific Google Cloud regions. To use a TPU version, ensure its availability in your preferred region. For a list of TPU locations, see TPU availability.

  3. Ensure that your project has sufficient TPU quota: if you are creating a multi-host TPU slice with on-demand or Spot VMs, you must have sufficient TPU quota available in the region that you want to use. Creating a multi-host TPU slice that consumes a TPU reservation doesn't require any TPU quota as quota is used when the reservation was created. For a list of TPU quota names, see TPU quota and for instructions on how to view the quota, see View and manage quotas

  4. Choose a TPU consumption option: select a consumption option that best fits your workload, its duration, and your cost needs. For a list of consumption option availability by TPU versions, see TPU consumption options.

  5. Choose a topology: select a topology supported for the selected TPU version. For a list of topologies available for each TPU version, see TPU topology.

Create a MIG with multi-host TPU slices

  1. Create an instance template.
  2. Create a workload policy.
  3. Create the MIG.

Create an instance template

The command to create an instance template depends on the consumption option you use: on-demand, Spot, reservation-bound, or flex-start. For more information about consumption options, see About VM provisioning models.

Create an instance template for an on-demand TPU VM

The following command creates an instance template using the on-demand consumption option:

gcloud compute instance-templates create INSTANCE_TEMPLATE_NAME \
  --machine-type=MACHINE_TYPE \
  --maintenance-policy=TERMINATE \
  --image-family=IMAGE_FAMILY \
  --image-project=IMAGE_PROJECT

Replace the following placeholders:

  • INSTANCE_TEMPLATE_NAME: The name of your instance template.
  • MACHINE_TYPE: The machine type for the TPU VM, for example, ct6e-standard-8t.
  • IMAGE_FAMILY: The OS image family for the TPU VM. If you want to install a specific OS version, use the --image flag. For more information about OS images, see OS images.
  • IMAGE_PROJECT: The project that contains the OS image. For TPU images, this is ubuntu-os-accelerator-images.

Create an instance template for a TPU Spot VM

The following command creates an instance template using the Spot consumption option:

gcloud compute instance-templates create INSTANCE_TEMPLATE_NAME \
  --machine-type=MACHINE_TYPE \
  --maintenance-policy=TERMINATE \
  --instance-termination-action=STOP \
  --provisioning-model=SPOT \
  --image-family=IMAGE_FAMILY \
  --image-project=IMAGE_PROJECT

Replace the following placeholders:

  • INSTANCE_TEMPLATE_NAME: The name of your instance template.
  • MACHINE_TYPE: The machine type for the TPU VM, for example, ct6e-standard-8t.
  • IMAGE_FAMILY: The OS image family for the TPU VM. If you want to install a specific OS version, use the --image flag. For more information about OS images, see OS images.
  • IMAGE_PROJECT: The project that contains the OS image. For TPU images, this is ubuntu-os-accelerator-images.

Create an instance template for a TPU reservation-bound VM

The following command creates an instance template using the reservation-bound consumption option:

gcloud compute instance-templates create INSTANCE_TEMPLATE_NAME \
  --machine-type=MACHINE_TYPE \
  --maintenance-policy=TERMINATE \
  --instance-termination-action=DELETE \
  --reservation-affinity=specific \
  --provisioning-model=reservation-bound \
  --reservation=RESERVATION_NAME \
  --image-family=IMAGE_FAMILY \
  --image-project=IMAGE_PROJECT

Replace the following placeholders:

  • INSTANCE_TEMPLATE_NAME: The name of your instance template.
  • MACHINE_TYPE: The machine type for the TPU VM, for example, ct6e-standard-8t.
  • RESERVATION_NAME: The name of the specific reservation to consume.
  • IMAGE_FAMILY: The OS image family for the TPU VM. If you want to install a specific OS version, use the --image flag. For more information about OS images, see OS images.
  • IMAGE_PROJECT: The project that contains the OS image. For TPU images, this is ubuntu-os-accelerator-images.

Create an instance template for a TPU Flex-start VM

The following command creates an instance template using the flex-start consumption option:

gcloud compute instance-templates create INSTANCE_TEMPLATE_NAME \
    --machine-type=MACHINE_TYPE \
    --maintenance-policy=TERMINATE \
    --instance-termination-action=DELETE \
    --provisioning-model=FLEX_START \
    --max-run-duration=DURATION \
    --image-family=IMAGE_FAMILY \
    --image-project=IMAGE_PROJECT

Replace the following placeholders:

  • INSTANCE_TEMPLATE_NAME: The name of your instance template.
  • MACHINE_TYPE: The machine type for the TPU VM, for example, ct6e-standard-8t.
  • DURATION: The maximum duration for which the TPU VM can run.
  • IMAGE_FAMILY: The OS image family for the TPU VM. If you want to install a specific OS version, use the --image flag. For more information about OS images, see OS images.
  • IMAGE_PROJECT: The project that contains the OS image. For TPU images, this is ubuntu-os-accelerator-images.

Create a workload policy

You must create a workload policy with the accelerator-topology parameter (for example, 4x4, 8x8, or 4x4x4). The accelerator topology configures the MIG to treat the instances as a single, interconnected slice.

The following command creates a workload policy:

gcloud compute resource-policies create workload-policy WORKLOAD_POLICY_NAME \
  --type=high-throughput \
  --accelerator-topology=TOPOLOGY \
  --region=REGION

Replace the following placeholders:

  • WORKLOAD_POLICY_NAME: The name of your workload policy.
  • TOPOLOGY: The topology of the TPU VMs, for example, 4x4x8. For more information about topology for each version of TPU, see TPU topology.
  • REGION: The region for your workload policy.

Create a MIG

Create a zonal or a regional MIG by using the gcloud compute instance-groups managed create command as follows:

  • To create a zonal MIG containing a multi-host TPU slice, use the following command:

     gcloud compute instance-groups managed create MIG_NAME \
    --size=MIG_SIZE \
    --target-size-policy-mode=bulk \
    --template=INSTANCE_TEMPLATE_URL \
    --zone=ZONE \
    --default-action-on-vm-failure=do-nothing \
    --workload-policy=WORKLOAD_POLICY_URL

  • To create a regional MIG containing a multi-host TPU slice, use the following command:

     gcloud compute instance-groups managed create MIG_NAME \
    --size=MIG_SIZE \
    --target-size-policy-mode=bulk \
    --template=INSTANCE_TEMPLATE_URL \
    --region=REGION \
    --default-action-on-vm-failure=do-nothing \
    --workload-policy=WORKLOAD_POLICY_URL \
    --target-distribution-shape=any-single-zone \
    --instance-redistribution-type=none

Replace the following placeholders:

  • MIG_NAME: The name of your MIG.
  • MIG_SIZE: The number of VMs in the MIG.
  • INSTANCE_TEMPLATE_URL: the URL of the instance template that you want to use to create instances in the MIG. The URL can contain either the ID or name of the instance template. Specify one of the following values:
    • For a regional instance template: projects/PROJECT_ID/regions/REGION/instanceTemplates/INSTANCE_TEMPLATE_ID
    • For a global instance template: INSTANCE_TEMPLATE_ID
  • ZONE: The zone for your MIG.
  • REGION: The region for your MIG.
  • WORKLOAD_POLICY_URL: The URL of the workload policy that you want to use to create instances in the MIG. For example: projects/PROJECT_ID/regions/WORKLOAD_POLICY_REGION/resourcePolicies/WORKLOAD_POLICY_NAME.

Create VMs with custom names in a MIG

You can create VMs in a MIG by specifying custom names for each VM. This is useful for debugging and ensuring instances are created in a specific order.

MIGs that contain a multi-host TPU slice use the bulk mode of target size policy. When creating VMs with custom names in such a MIG, the following applies:

  • You must first verify that the MIG doesn't have VMs in it. If the MIG has VMs, you must either resize the MIG to target size 0 or create another MIG with target size 0.

  • You can only use the REST API to create VMs with custom names.

Create VMs with custom names by using one of the following REST API methods:

  • For a zonal MIG, use the instanceGroupManagers.createInstances.

     POST https://compute.googleapis.com/compute/v1/projects/PROJECT_ID/zones/ZONE/instanceGroupManagers/MIG_NAME/createInstances
 {
   "instances": [
     {
       "name": "INSTANCE_NAME_1"
     },
     {
       "name": "INSTANCE_NAME_2"
     },
     ...
   ]
 }

  • For a regional MIG, use the regionInstanceGroupManagers.createInstances.

     POST https://compute.googleapis.com/compute/v1/projects/PROJECT_ID/regions/REGION/instanceGroupManagers/MIG_NAME/createInstances
 {
   "instances": [
     {
       "name": "INSTANCE_NAME_1"
     },
     {
       "name": "INSTANCE_NAME_2"
     },
     ...
   ]
 }

Replace the following placeholders:

  • PROJECT_ID: The ID of the project where the MIG exists.
  • ZONE: The zone of the MIG.
  • REGION: The region of the MIG.
  • INSTANCE_NAME_1,2,..: The names of the VMs to add to the specified MIG.

What's next