GKE with TPU, in depth: slice node pools, labels and chip requests, Indexed Jobs, Multislice with JobSet, capacity options and failure modes

How GKE models a TPU slice

A Cloud TPU VM hosts up to eight TPU chips; on the multi-host machine types for v4, v5e, v5p and Trillium it hosts four. A slice is a group of chips connected by ICI, the TPU interconnect, in a 2D or 3D mesh or torus topology written as a shape such as 2x4, 4x4 or 2x2x2. The number of chips is the product of the shape. GKE represents each TPU VM as one Kubernetes node and gives you two kinds of node pool.

• Single-host TPU slice node pools. Each VM is its own small slice; chips on different nodes are not connected by ICI. Nodes are independent, so the pool can scale node by node, much like a GPU pool. Good for inference and small fine-tunes.
• Multi-host TPU slice node pools. The nodes together form one slice. GKE treats them as a single unit: creation is atomic, so if one VM cannot be created none are, and the set is immutable, so you cannot add a node to an existing slice. When any VM needs repair, GKE shuts down every VM in the slice and all Pods on it are evicted.

That last property shapes everything else. A multi-host slice fails, restarts and is preempted as a whole, so the job running on it must be designed to restart as a whole from a checkpoint. It is also why the scheduling unit you plan around is the slice, not the node.

Labels, chip requests and node counts

Pods reach TPUs through two node labels and one extended resource. The label cloud.google.com/gke-tpu-accelerator names the TPU generation, and cloud.google.com/gke-tpu-topology names the slice shape. Chips are requested as google.com/tpu. GKE cannot hand out part of a node's TPUs: a Pod on a multi-host slice node requests all of the chips on that node, which on the four-chip machine types means google.com/tpu: 4.

Newer generations, including Ironwood, have their own values and minimum GKE versions; take them from the current GKE TPU documentation or from kubectl get nodes -L cloud.google.com/gke-tpu-accelerator on a live cluster rather than from a blog post. The node count of a multi-host pool follows from the shape: chips in the topology divided by chips per VM. The documentation's example is a v4 topology of 2x4x4, which is 32 chips, so 32 / 4 = 8 nodes. A v5e 4x4 is 16 chips on 4 nodes. Get this wrong and the node pool request is rejected.

from math import prod

def slice_nodes(topology: str, chips_per_vm: int = 4) -> int:
    chips = prod(int(x) for x in topology.split("x"))
    if chips % chips_per_vm:
        raise ValueError(f"{topology} is {chips} chips, not a multiple of {chips_per_vm}")
    return chips // chips_per_vm

assert slice_nodes("2x4x4") == 8      # TPU v4 example from the GKE docs
assert slice_nodes("4x4") == 4        # TPU v5e, 16 chips

Creating a slice node pool

In a Standard cluster you create the slice as a node pool. The flags that matter are the machine type, the topology, the node count and the capacity choice. Autopilot clusters provision the equivalent from the Pod's selectors, but the same shape rules apply.

gcloud container node-pools create v5e-4x4 \
  --cluster=train-cluster \
  --location=REGION \
  --node-locations=ZONE \
  --machine-type=ct5lp-hightpu-4t \
  --tpu-topology=4x4 \
  --num-nodes=4 \
  --spot
# alternatives to --spot: --reservation-affinity=specific --reservation=NAME,
# or --flex-start; with --enable-autoscaling, start at --num-nodes=0

Pick the region and zone from where TPU capacity and quota actually exist; TPU quota is tracked per generation and per region, with separate quotas for on-demand and preemptible chips. On GKE versions before 1.28, Pods had to run privileged to reach the TPU devices; the documentation states that from GKE 1.28 nodes no longer need privileged mode, so drop privileged: true on current clusters.

Running a multi-host job

A multi-host slice runs one Pod per node, and those Pods must find each other. The documented pattern is an Indexed Job, whose Pods get stable ordinal indexes, plus a headless Service, which gives each Pod a DNS name. The TPU runtime and JAX use those to form the slice. Ports 8471 (TPU VM communication) and 8431 (runtime metrics) are exposed in Google's examples.

apiVersion: v1
kind: Service
metadata:
  name: headless-svc
spec:
  clusterIP: None
  selector:
    job-name: tpu-train
---
apiVersion: batch/v1
kind: Job
metadata:
  name: tpu-train
spec:
  completions: 4            # one Pod per node in the 4x4 slice
  parallelism: 4
  completionMode: Indexed
  backoffLimit: 0           # a slice restarts as a unit; let the controller above retry
  template:
    spec:
      subdomain: headless-svc
      restartPolicy: Never
      nodeSelector:
        cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
        cloud.google.com/gke-tpu-topology: 4x4
      containers:
      - name: train
        image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
        ports:
        - containerPort: 8471
        - containerPort: 8431
        command: ["python", "-u", "train.py"]
        resources:
          limits:
            google.com/tpu: 4

Inside the container, a JAX program on TPU initialises without arguments: the runtime discovers the other hosts of the slice. The first thing every training script should do is check that it sees the slice it asked for, and fail fast if not.

import jax

jax.distributed.initialize()          # on Cloud TPU, peers are discovered automatically
expected_chips = 16                   # 4x4 slice
print(f"process {jax.process_index()} of {jax.process_count()}, "
      f"{jax.local_device_count()} local, {jax.device_count()} global")
if jax.device_count() != expected_chips:
    raise SystemExit(f"expected {expected_chips} chips, got {jax.device_count()}")

Multislice with JobSet

One slice is bounded by its ICI domain. To train on more chips than a slice holds, GKE supports Multislice: several slices joined over the data-centre network (DCN), which is slower than ICI. The standard layout is data parallelism across slices and model or FSDP parallelism inside each slice, so only gradient reductions cross DCN. Multislice requires multi-host slice node pools; a single-host shape such as v5e 2x2 cannot participate.

The orchestration object is a JobSet: one replicated Job per slice. The documented annotation alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool makes each child Job claim a whole node pool, so two slices never share one, and port 8080 carries the Multislice (MXLA) coordinator traffic.

apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
  name: multislice-train
  annotations:
    alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
spec:
  failurePolicy:
    maxRestarts: 4           # restart the whole JobSet from the last checkpoint
  replicatedJobs:
  - name: slice
    replicas: 2              # two 4x4 slices = 32 chips
    template:
      spec:
        parallelism: 4
        completions: 4
        backoffLimit: 0
        template:
          spec:
            nodeSelector:
              cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
              cloud.google.com/gke-tpu-topology: 4x4
            containers:
            - name: train
              image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
              ports:
              - containerPort: 8471
              - containerPort: 8080
              - containerPort: 8431
              resources:
                limits:
                  google.com/tpu: 4

Getting capacity

Getting chips is usually harder than using them. GKE offers four routes, each with a different failure profile.

• Reservations (--reservation-affinity=specific): capacity you have committed to. Predictable; you pay whether or not it is busy.
• On demand: available if the zone has capacity and you have quota. Large multi-host slices often fail to provision atomically.
• Spot (--spot): cheaper, preemptible at any time, and a preemption of one VM takes the whole slice. Checkpoint often.
• Flex-start (--flex-start): waits for capacity and then runs for a bounded time; pairs naturally with a queue.

When several teams share TPU capacity, put a queue in front of it. Kueue admits whole jobs against quotas, so a JobSet that needs two slices either gets both or waits, instead of holding one slice idle while it waits for the second. The general GKE platform pieces, including storage and serving, are covered in GKE for ML.

Worked example: fine-tuning on a v5e 4x4

A team wants to fine-tune a 7B-parameter model on TPU v5e, which has 16 GB of HBM per chip. Full fine-tuning with Adam in mixed precision needs on the order of 16 bytes per parameter for weights, gradients and optimizer state, about 112 GB before activations, so it needs at least 8 chips of state sharding and in practice 16 to leave room for activations. They choose a v5e 4x4: 16 chips, 256 GB of HBM, four nodes.

The plan follows directly. Create a Spot node pool with topology 4x4 and four nodes; deploy the Indexed Job above with FSDP sharding over all 16 chips; write checkpoints to Cloud Storage every 20 minutes, because a Spot preemption restarts the slice; gate the Job behind a Kueue queue so the next experiment waits rather than collides. When they later need twice the throughput, they do not ask for a bigger slice shape in a zone that may not have one; they move to the JobSet above with replicas: 2 and run data parallel across two slices over DCN.

Failure modes

• Pod Pending with no node match. The topology label, accelerator label or chip request does not match any pool. Compare the Pod's nodeSelector with kubectl get nodes -L cloud.google.com/gke-tpu-topology.
• Partial chip request. Requesting 1 or 2 chips on a four-chip node never schedules on a multi-host slice. Request all of them.
• Node pool creation fails. Atomic provisioning means one missing VM fails the whole slice. Retry in another zone, use a reservation, or use flex-start.
• Whole job evicted. Repair, upgrade or Spot preemption of one VM restarts every Pod. Set maintenance windows and exclusions for training periods and checkpoint on a timer.
• Two jobs on one slice. Without the exclusive-topology annotation or anti-affinity, Pods from different jobs can collide on a slice and both hang. Use the annotation.
• Hang at initialisation. Usually one Pod is missing or cannot reach its peers: check the headless Service selector, the subdomain and that all completions were created.

Trade-offs

GKE is the right home for TPUs when your team already runs Kubernetes, wants one scheduler for TPUs, GPUs and CPU services, and is willing to own the YAML. Vertex AI training hides the slice mechanics at the cost of control. Raw TPU VMs created directly through the Cloud TPU API are the simplest for one person running one job, and the hardest to share. The trade on GKE is the same as anywhere else: you gain queueing, quotas and a common platform, and you inherit a scheduler whose basic unit, the node, is smaller than the TPU's real unit, the slice.

What to do next

1. List the TPU generations and quotas in your region, on demand and preemptible, before designing anything.
2. Choose a topology from the model's memory needs, then compute the node count as chips divided by chips per VM.
3. Create a small multi-host pool, run the device-count check in an Indexed Job, and confirm the global chip count.
4. Add checkpointing to Cloud Storage on a timer, and test a restart by deleting one Pod.
5. Put a Kueue queue in front of TPU pools if more than one person submits jobs.
6. Scale out with JobSet replicas and the exclusive-topology annotation when one slice is no longer enough.
7. Set maintenance windows and exclusions so upgrades do not land mid-run.