Cloud TPU Multislice Overview

Cloud TPU Multislice is a full stack performance-scaling technology that enables a training job to use multiple TPU slices within a single pod or on slices in multiple Pods with standard data parallelism. With TPU v4 chips, this means training jobs can use more than 4096 chips in a single run. For training jobs that require less than 4096 chips, a single slice can offer the best performance. However, multiple smaller slices are more readily available, allowing for a faster startup time when Multislice is used with smaller slices.

When deployed in Multislice configurations, TPU chips in each slice communicate through inter-chip-interconnect (ICI). TPU chips in different slices communicate by transferring data to CPUs (hosts) which in turn transmit the data over the data-center network (DCN). For more information about scaling with Multislice, see How to scale AI training to up to tens of thousands of Cloud TPU chips with Multislice.

Developers don't have to write code to implement inter-slice DCN communication. The XLA compiler generates that code for you and overlaps communication with computation for maximum performance.

Concepts

Accelerator type

The shape of each TPU slice that comprises a Multislice. Each slice in a multislice request is of the same accelerator type. An accelerator type consists of a TPU type (v4 or later) followed by the number of TensorCores. For example, v5litepod-128 specifies a TPU v5e with 128 TensorCores.

Auto-repair

When a slice encounters a maintenance event, preemption or hardware failure, Cloud TPU will create a new slice. If there are insufficient resources to create a new slice, the creation won't complete until hardware becomes available. After the new slice is created all other slices in the Multislice environment will be restarted so training can continue. With a properly configured startup script, the training script can automatically relaunch without user intervention, loading and resuming from the latest checkpoint.

Data Center Networking (DCN)

A higher latency, lower-throughput network (when compared with ICI) that connects TPU slices in a Multislice configuration.

Gang scheduling

When all TPU slices are provisioned together, at the same time, guaranteeing either all or none of the slices are provisioned successfully.

Interchip interconnect (ICI)

High speed, low latency internal links that connect TPUs within a TPU Pod.

Multislice

Two or more TPU chip slices that can communicate over DCN.

Node

In the Multislice context, node refers to a single TPU slice. Each TPU slice in a Multislice is given a node ID.

Startup script

A standard Compute Engine startup script that is run every time a VM is booted or rebooted. For Multislice, it is specified in the QR creation request. For more information about Cloud TPU startup scripts, see Manage TPU resources.

Tensor

A data structure that is used to represent multidimensional data in a machine learning model.

Types of Cloud TPU capacity

TPUs can be created from different types of capacity (see Usage Options in How TPU pricing works):

• Reservation: To consume a reservation, you must have a reservation agreement with Google. Use the --reserved flag when creating your resources.

• Spot: Targets preemptible quota using Spot VMs. Your resources may be preempted to make room for requests for a higher priority job. Use the --spot flag when creating your resources.

• On-demand: Targets on-demand quota, which doesn't need a reservation and won't be preempted. The TPU request will be enqueued to an on-demand quota queue offered by Cloud TPU, the availability of resources is not guaranteed. Selected by default, no flags needed.

Get started

1. Set up your Cloud TPU environment.

2. In the Google Cloud console, activate Cloud Shell.

   Activate Cloud Shell

   At the bottom of the Google Cloud console, a Cloud Shell session starts and displays a command-line prompt. Cloud Shell is a shell environment with the Google Cloud CLI already installed and with values already set for your current project. It can take a few seconds for the session to initialize.

To use Multislice, your TPU resources must be managed as queued resources.

Introductory example

This tutorial uses code from the MaxText GitHub repo. MaxText is a high performance, arbitrarily scalable, open source, and well-tested basic LLM written in Python and Jax. MaxText was designed to train efficiently on Cloud TPU.

The code in shardings.py is designed to help you get started experimenting with different parallelism options. For example, data parallelism, fully sharded data parallelism (FSDP), and tensor parallelism. The code scales from single slice to Multislice environments.

ICI parallelism

ICI refers to the high speed interconnect that connects the TPUs in a single slice. ICI sharding corresponds to sharding within a slice. shardings.py provides three ICI parallelism parameters:

• ici_data_parallelism
• ici_fsdp_parallelism
• ici_tensor_parallelism

The values you specify for these parameters determine the number of shards for each parallelism method.

These inputs must be constrained so that ici_data_parallelism * ici_fsdp_parallelism * ici_tensor_parallelism is equal to the number of chips in the slice.

The following table shows example user inputs for ICI parallelism for the four chips available in v4-8:

	ici_data_parallelism	ici_fsdp_parallelism	ici_tensor_parallelism
4-way FSDP	1	4	1
4-way Tensor parallelism	1	1	4
2-way FSDP + 2-way Tensor parallelism	1	2	2

Note that ici_data_parallelism should be left as 1 in most cases because the ICI network is fast enough to almost always prefer FSDP to data parallelism.

This example assumes you are familiar with running code on a single TPU slice such as in Run a calculation on a Cloud TPU VM using JAX. This example show how to run shardings.py on a single slice.

1. Set up the environment:

   $ gcloud auth login
   $ export QR_ID=your-queued-resource-id
   $ export TPU_NAME=your-tpu-name
   $ export PROJECT=your-project-name
   $ export ZONE=us-central1-a
   $ export NETWORK_NAME=your-network-name
   $ export