Run PyTorch on GKE TPUs with Polyaxon

TorchTPU gives PyTorch a native TPU backend: tensors use device="tpu", eager execution works through ordinary PyTorch operations, and compiled workloads use torch.compile. Google and Meta announced the open-source release on October 20, 2026.

You can schedule that workload through Polyaxon using GKE's TPU resource, then keep its configuration, metrics and model weights in the same run. This walkthrough trains a small MNIST classifier on one TPU7x host, with a simplified Polyaxonfile and eight data-parallel workers.

Four chips are eight devices

For the TPU7x configuration used here, one host has four chips, each with two TensorCores. TorchTPU exposes each TensorCore as a PyTorch device. Kubernetes therefore requests four chips, while torchrun launches eight processes. This mapping is specific to this TPU generation and configuration. Upstream GKE guide.

Conceptual layout for one four-chip TPU7x host. The drawing does not assign physical device IDs or rank order.

A single .to("tpu") process would not automatically train across all eight devices. In the program below, DistributedSampler partitions the dataset across workers; DistributedDataParallel synchronizes their gradients.

Use an existing TPU pool

You need Polyaxon CLI and server 2.18 or later for the simplified workload syntax, an existing project, and a Polyaxon Agent/queue connected to a GKE Standard cluster with a four-chip TPU7x pool. The cluster must have the TPU capacity, quota and placement policy described in GKE's deployment guide.

This recipe retains the upstream TorchTPU example's host networking and privileged container settings. Use a pool and namespace where those settings are permitted. The container also needs access to the package index, MNIST downloads, the Polyaxon API and its artifact store.

Inspect the intended pool before filling in the configuration:

kubectl --context YOUR_CLUSTER get nodes \
  -l cloud.google.com/gke-nodepool=YOUR_TPU_POOL \
  -L cloud.google.com/gke-tpu-accelerator,cloud.google.com/gke-tpu-topology,cloud.google.com/placement-policy-name

The example targets that pool and its 2x2x1 topology. Set YOUR_PLACEMENT_POLICY to the pool's Ironwood workload policy. GKE's current accelerator labels differ from the older label in TorchTPU's reference manifest, so pool selection avoids assuming that string applies to every cluster.

The deployment below was reviewed against upstream source and Polyaxon's configuration handling; it has not been executed on TPU hardware for this article.

Define the Polyaxon job

Create a folder containing torchtpu.yaml and the train_tpu.py script in the next section. Save this as torchtpu.yaml, replacing the two pool/policy placeholders:

name: torchtpu-mnist

container:
  image: us-docker.pkg.dev/ml-oss-artifacts-published/torch-tpu-public-release-container/torch-tpu@sha256:e7509a5699e70f51c52d1c8378cfe30f07a7182c9c73304f02fb5d59feb8524f
  workingDir: "{{ globals.run_artifacts_path }}/uploads"
  securityContext:
    privileged: true

environment:
  hostNetwork: true
  dnsPolicy: ClusterFirstWithHostNet
  annotations:
    cloud.google.com/gke-tpu-slice-topology: 2x2x1
  nodeSelector:
    cloud.google.com/gke-nodepool: YOUR_TPU_POOL
    cloud.google.com/gke-tpu-topology: 2x2x1
    cloud.google.com/placement-policy-name: YOUR_PLACEMENT_POLICY
  tolerations:
    - key: google.com/tpu
      operator: Exists

resources:
  cpu: 8
  memory: 32Gi
  google.com/tpu: 4

env:
  TORCH_TPU_TOPOLOGY: "2,2,1,2"

cmd: |
  python -m pip install --no-cache-dir polyaxon==2.18.0
  python train_tpu.py --download
  eval "$(python -m torch.tpu.distributed)"
  export TORCH_TPU_TOPOLOGY TORCH_TPU_SLICEBUILDER_ADDRESSES
  python -m torch.distributed.run \
    --nnodes=1 --node_rank=0 \
    --master_addr=localhost --master_port=29500 \
    --nproc_per_node=8 train_tpu.py

The image digest was resolved from TorchTPU's public release image on October 24. It includes TorchTPU 0.2.0, PyTorch 2.14.0 CPU and torchvision 0.29.0 CPU; only the Polyaxon SDK is installed at startup. You do not need a CUDA image or a custom Dockerfile for this example.

The TPU resource sets equal requests and limits: four chips for the container, not four per Python process. The scalar CPU and memory entries set requests, using starting values for this small example. Polyaxon already provides a shared-memory mount, so it needs no duplicate volume definition here.

The topology helper runs before the workers initialize the TPU runtime. It sets the local coordinator addresses used to construct the TPU mesh. Within each worker, tpu_dist initializes communication and torch.device("tpu") refers to its assigned device. Do not replace it with a CUDA local-rank device index. TorchTPU distributed initialization.

Train and record the result

Save the following as train_tpu.py. It uses ordinary PyTorch modules, a distributed sampler and SGD. Dataset downloads happen once before the workers start; each worker then loads its own portion with num_workers=0, following the upstream GKE example's guidance about TPU runtime initialization and forked loader processes.

import argparse
from importlib.metadata import version

import torch
import torch.distributed as dist
from torch import nn
from torch.nn.parallel import DistributedDataParallel
from torch.utils.data import DataLoader, DistributedSampler
from torchvision import datasets, transforms
from polyaxon.tracking import Run

parser = argparse.ArgumentParser()
parser.add_argument("--download", action="store_true")
parser.add_argument("--compile", action="store_true")
args = parser.parse_args()

data_dir = "/tmp/mnist"
transform = transforms.ToTensor()
if args.download:
    datasets.MNIST(data_dir, train=True, download=True)
    datasets.MNIST(data_dir, train=False, download=True)
    raise SystemExit(0)

dist.init_process_group(backend="tpu_dist")
rank, world_size = dist.get_rank(), dist.get_world_size()
device = torch.device("tpu")
torch.manual_seed(42)

train_data = datasets.MNIST(data_dir, train=True, transform=transform)
test_data = datasets.MNIST(data_dir, train=False, transform=transform)
sampler = DistributedSampler(train_data, shuffle=True, seed=42)
train_loader = DataLoader(
    train_data, batch_size=128, sampler=sampler,
    num_workers=0, drop_last=True,
)
test_loader = DataLoader(
    test_data, batch_size=250,
    sampler=DistributedSampler(test_data, shuffle=False), num_workers=0,
)

network = nn.Sequential(
    nn.Flatten(), nn.Linear(784, 256), nn.ReLU(), nn.Linear(256, 10),
).to(device)
model = DistributedDataParallel(network)
if args.compile:
    model = torch.compile(model, backend="tpu")
optimizer = torch.optim.SGD(model.parameters(), lr=0.05)
criterion = nn.CrossEntropyLoss()

run = Run() if rank == 0 else None
if run is not None:
    run.log_inputs(
        epochs=3, batch_size_per_worker=128, world_size=world_size,
        learning_rate=0.05, seed=42, compiled=args.compile,
        torch_version=str(torch.__version__), torchtpu_version=version("torch_tpu"),
    )
print(f"Rank {rank}/{world_size} initialized on {device}", flush=True)

for epoch in range(1, 4):
    sampler.set_epoch(epoch)
    model.train()
    totals = torch.zeros(4, device=device)
    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)
        optimizer.zero_grad(set_to_none=True)
        loss = criterion(model(images), labels)
        loss.backward()
        optimizer.step()
        totals[0] += loss.detach() * labels.numel()
        totals[1] += labels.numel()

    network.eval()
    with torch.no_grad():
        for images, labels in test_loader:
            images, labels = images.to(device), labels.to(device)
            totals[2] += (network(images).argmax(dim=1) == labels).sum()
            totals[3] += labels.numel()

    dist.all_reduce(totals, op=dist.ReduceOp.SUM)
    loss_sum, train_count, correct, test_count = totals.cpu().tolist()
    if run is not None:
        run.log_metrics(
            train_loss=loss_sum / train_count,
            test_accuracy=correct / test_count, step=epoch,
        )

if run is not None:
    checkpoint = run.get_outputs_path("mnist.pt")
    torch.save(
        {name: value.detach().cpu() for name, value in network.state_dict().items()},
        checkpoint,
    )
    run.log_model_ref(path=checkpoint, framework="pytorch")
    run.end()

dist.barrier()
dist.destroy_process_group()

All ranks contribute to the loss and accuracy totals; rank zero writes one set of metrics and the model weights. MNIST's 10,000 test examples divide evenly across these eight workers, so the distributed test sampler adds no repeated examples in this configuration. The weights file supports loading the trained model; it is not a full optimizer checkpoint for exact training recovery.

This small network makes the execution path easy to inspect. It is not a TPU utilization benchmark or a recommended workload for comparing accelerator cost.

Submit, inspect, then compare compilation

From the folder containing the two files, submit to your project's existing TPU queue:

polyaxon run -p YOUR_PROJECT --queue YOUR_AGENT/YOUR_TPU_QUEUE \
  -f torchtpu.yaml -u --logs

-u uploads the folder and makes the files available in the configured working directory. The TPU selectors choose the GKE pool; the Polyaxon queue routes the operation through the appropriate agent. They solve different parts of placement.

In the run, check for initialization from eight ranks, then inspect train_loss, test_accuracy and the mnist.pt output. Polyaxon retains the submitted configuration and tracked metrics and artifacts. A job waiting for TPU capacity has not reached framework initialization; a runtime or collective error needs the worker logs.

For a second run, append --compile to the final train_tpu.py invocation in cmd and submit again. Keep the image, batch size and other settings fixed. The script compiles the DDP model with backend="tpu", following PyTorch's DDP compilation order. Compare the two runs in Polyaxon, including their resulting accuracy. Measure compilation startup separately from steady-state steps before drawing a performance conclusion.

TorchTPU presents native PyTorch APIs while still using OpenXLA and the TPU runtime underneath. That distinction matters when porting CUDA kernels, choosing numerical precision or tuning collective communication. The initial feature matrix also marks some paths, including FSDPv2, experimental. Treat this single-host job as the starting point for your model, rather than assuming every GPU configuration transfers unchanged.

The jobs exit after training. To stop an unfinished operation, use polyaxon ops stop -p YOUR_PROJECT -uid RUN_UUID. Moving beyond one host requires coordinated placement and a launcher for the full TPU slice; increasing this job's process count alone does not create that deployment.