TorchTPU in practice: running PyTorch on TPUs and evaluating migration

Moving a PyTorch workload from GPUs to TPUs requires more than changing a device string. The migration must preserve model quality, justify engineering effort, and demonstrate that the platform can provision, observe, and recover the workload.

On October 20, 2026, Google announced the open-source availability of TorchTPU, developed in collaboration with Meta's PyTorch team. Its goal is to run PyTorch on Tensor Processing Units, or TPUs, through a more familiar framework experience. That does not establish compatibility or a financial advantage for every workload. (cloud.google.com)

This guide follows an evaluation sequence: compatibility and correctness first, performance second, operational readiness and total cost last. The examples are educational—not measured benchmarks or a complete production platform.

1. What TorchTPU is—and what it does not replace

TorchTPU is a separately distributed, native PyTorch backend that registers the tpu device through PrivateUse1. It provides eager execution and torch.compile integration. It is neither a CUDA emulator nor an automatic converter for custom CUDA extensions. (pytorch.org)

Two different migrations need separate plans:

• PyTorch/CUDA → TorchTPU: review device handling, libraries, operators, and GPU-specific kernels.
• PyTorch/XLA → TorchTPU: also revisit execution semantics and integrations tied to torch_xla.

Google's documentation positions TorchTPU as the replacement for PyTorch/XLA, recommending it for new workloads and migration of existing ones. However, replacing torch_xla does not remove XLA: the compiler remains part of the execution stack. (docs.cloud.google.com)

Execution architecture

A simplified compiled execution path is:

PyTorch model → TorchDynamo/FX → TorchTPU backend → StableHLO → XLA compiler → runtime → TPU

ATen is PyTorch's operator layer; StableHLO is an intermediate operation representation; XLA compiles those operations for the accelerator. The TPU backend for torch.compile uses XLA rather than TorchInductor. (torchtpu.readthedocs.io)

Infrastructure follows a separate path:

Pipeline → Kubernetes Job → scheduling onto a TPU node → Python process → TPU runtime

For the architecture proposed here, keep datasets, checkpoints, and outputs outside the Pod lifecycle. Treat the cluster as execution capacity, not the only durable home for artifacts.

2. What changes in application behavior and performance

Eager execution still involves compilation

Default eager execution preserves imperative behavior, but device work remains asynchronous and can involve operation compilation. torch.compile exposes larger optimization regions. Reading device tensors on the CPU or making Python decisions from their values can force synchronization and interrupt optimization opportunities. (torchtpu.readthedocs.io)

The practical implication is to measure first execution, warmup, and steady state separately. A Python call returning is not necessarily evidence that the accelerator has finished its work.

Shapes become an operational contract

Changing dimensions can trigger recompilation. TorchTPU documents bounded dynamism with upper limits, not unrestricted shape changes at zero cost. (torchtpu.readthedocs.io)

For the pilot, define an explicit input catalog: batch size, sequence length, resolution, and dtype. Compare length bucketing, padding, and representative variable inputs. Include edge cases rather than testing only the average request.

Numerical precision needs an explicit policy

A tensor stored as float32 does not guarantee an identical arithmetic path across backends. TorchTPU documents BF16-based matrix multiplication with FP32 accumulation and higher-precision modes that require additional work. (torchtpu.readthedocs.io)

Set the comparison policy before testing. Check outputs, gradients, loss stability, and final task quality. Do not relax tolerances merely to pass a test, or require bitwise equality without a domain-specific reason.

API coverage and production maturity are different questions

The initial feature matrix distinguishes validated capabilities from experimental ones. For example, FSDP2 is listed as experimental with performance caveats, while SimpleFSDP appears among validated paths. Check the status of the release you adopt before approving a distributed architecture. (torchtpu.readthedocs.io)

3. A reproducible smoke test

Start on supported TPU hardware with Python and dependencies matched to the selected release. Installation guidance recommends the CPU build of PyTorch alongside TorchTPU; running on TPU does not require installing the CUDA build. (github.com)

The following command requires an actual release selected by your team:

python3 -m venv .venv
source .venv/bin/activate
: "${TORCH_TPU_VERSION:?Set a validated TorchTPU release}"
python -m pip install "torch_tpu==${TORCH_TPU_VERSION}" --extra-index-url https://download.pytorch.org/whl/cpu
python -m pip check
python -m pip freeze > requirements-resolved.txt

For production, promote the approved dependency resolution into a hash-verified lockfile and an immutable image. pip freeze records versions but does not replace that process by itself.

Save the following as smoke.py. It measures synthetic inference on one device with inputs already resident on the TPU. Synchronization uses the public torch.accelerator.synchronize API. (github.com)

import json
import os
import time
from importlib.metadata import version

import torch

mode = os.environ.get('EXECUTION_MODE', 'eager')
if mode not in {'eager', 'compiled'}:
    raise ValueError('EXECUTION_MODE must be eager or compiled')

torch.manual_seed(42)
device = torch.device('tpu')
model = torch.nn.Sequential(
    torch.nn.Linear(128, 256),
    torch.nn.ReLU(),
    torch.nn.Linear(256, 32),
).eval().to(device=device, dtype=torch.bfloat16)
x = torch.randn(64, 128).to(device=device, dtype=torch.bfloat16)

if mode == 'compiled':
    model = torch.compile(model, backend='tpu')

warmup = 5
iterations = 50

with torch.inference_mode():
    torch.accelerator.synchronize()
    start = time.perf_counter()
    y = model(x)
    torch.accelerator.synchronize()
    first_call_seconds = time.perf_counter() - start

    for _ in range(warmup):
        y = model(x)
    torch.accelerator.synchronize()

    start = time.perf_counter()
    for _ in range(iterations):
        y = model(x)
    torch.accelerator.synchronize()
    elapsed = time.perf_counter() - start

    finite = bool(torch.isfinite(y.float()).all().cpu().item())

if not finite:
    raise RuntimeError('Non-finite output')

print(json.dumps({
    'torch': torch.__version__,
    'torch_tpu': version('torch_tpu'),
    'mode': mode,
    'first_call_seconds': first_call_seconds,
    'mean_batch_seconds': elapsed / iterations,
    'samples_per_second': iterations * x.shape[0] / elapsed,
    'finite': finite,
}))

Run each mode in a separate process:

EXECUTION_MODE=eager python smoke.py
EXECUTION_MODE=compiled python smoke.py

This program illustrates a measurement protocol, not a performance improvement. The first-call measurement excludes part of process initialization and can be affected by persistent caches. Five warmup iterations are an educational choice, not a universal recommendation. Extend warmup until execution stabilizes and record cache conditions.

The test excludes data loading, backward computation, optimizer updates, distributed communication, and convergence. Add those stages before evaluating training.

4. Kubernetes: a complete pilot Job

Google Kubernetes Engine, or GKE, schedules TPU workloads using accelerator and topology selectors together with the google.com/tpu extended resource. The selected topology must match the hardware and chip allocation. This example uses a single-host v6e 1x1 configuration. (docs.cloud.google.com)

Prerequisites include a compatible GKE Standard cluster, a provisioned TPU node pool, available capacity, and image-pull permissions. Build an image containing the validated environment and /app/smoke.py. The image address below is illustrative and must be replaced.

apiVersion: v1
kind: Namespace
metadata:
  name: ml-evaluation
---
apiVersion: v1
kind: ServiceAccount
metadata:
  name: torchtpu-runner
  namespace: ml-evaluation
---
apiVersion: batch/v1
kind: Job
metadata:
  name: torchtpu-smoke
  namespace: ml-evaluation
spec:
  backoffLimit: 0
  activeDeadlineSeconds: 1800
  ttlSecondsAfterFinished: 86400
  template:
    metadata:
      labels:
        app: torchtpu-smoke
    spec:
      serviceAccountName: torchtpu-runner
      automountServiceAccountToken: false
      restartPolicy: Never
      nodeSelector:
        cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
        cloud.google.com/gke-tpu-topology: 1x1
      tolerations:
        - key: google.com/tpu
          operator: Exists
          effect: NoSchedule
      containers:
        - name: smoke
          image: us-central1-docker.pkg.dev/example-project/ml/torchtpu-smoke:validated
          imagePullPolicy: IfNotPresent
          command: [python, /app/smoke.py]
          env:
            - name: EXECUTION_MODE
              value: compiled
            - name: PYTHONUNBUFFERED
              value: '1'
          resources:
            requests:
              cpu: '2'
              memory: 8Gi
              ephemeral-storage: 1Gi
              google.com/tpu: 1
            limits:
              cpu: '4'
              memory: 16Gi
              ephemeral-storage: 4Gi
              google.com/tpu: 1

Here, backoffLimit: 0 prevents retries from obscuring a failed experiment; the deadline bounds execution; and the TTL cleans up the completed Job later. Export logs before cleanup. Jobs suit finite tasks and do not inherently require an HTTP readiness endpoint. (kubernetes.io)

CPU and memory values are starting points for this example, not production sizing. The container memory limit applies to host RAM, not the TPU's High Bandwidth Memory, or HBM. Accelerator memory requires separate analysis. (torchtpu.readthedocs.io)

Do not copy privileged: true from older examples without justification. GKE documentation associates that requirement with versions earlier than 1.28. Validate device permissions and image hardening for your chosen stack. (docs.cloud.google.com)

kubectl apply -f torchtpu-smoke.yaml
kubectl get pods -n ml-evaluation
kubectl logs -n ml-evaluation job/torchtpu-smoke
kubectl get events -n ml-evaluation --sort-by=.metadata.creationTimestamp

Use a new Job name for another configuration. Preserve the output from every run as evaluation evidence.

5. Evaluate migration without building a misleading benchmark

I recommend five approval gates:

For GPU-versus-TPU comparisons, hold the model, dataset, precision policy, quality target, and latency requirements constant. Document batch-size and parallelism changes needed to use each platform effectively.

For training, prioritize time and cost to target quality, not merely steps per second. For inference, measure throughput within the Service Level Objective, or SLO, including latency percentiles and failures. This is the recommended decision contract for the pilot, not a backend performance promise.

Include queueing, data preparation, compilation, checkpointing, and recovery in the end-to-end report. Keep cold and warm experiments separate: TorchTPU has multiple compilation-cache levels, and persistent caching should not be assumed to be enabled automatically. (torchtpu.readthedocs.io)

An illustrative enterprise scenario

Problem: a nightly pipeline generates catalog embeddings, and the team wants to evaluate an alternative to GPU capacity.

Solution: leave APIs and databases unchanged while moving only the inference worker into a TPU pilot.

Implementation: freeze a representative dataset, compare retrieval quality, run both platforms against the same output contract, and publish results to an isolated location. Use a run identifier to prevent duplicate promotion of outputs.

Expected outcome: a documented decision about quality, completion time, and cost. No percentage improvement is assumed. If the evidence is inconclusive, retain the existing execution path.

6. Production: availability, scaling, and recovery

Distributed training is not API replication

TorchTPU uses torch.distributed with the tpu_dist backend. Planning must map ranks, devices, and communication to physical topology. Increasing Kubernetes replica counts alone does not configure distributed training. (torchtpu.readthedocs.io)

For the proposed architecture, require coordinated worker-group recovery, complete checkpoints, and an injected worker-failure test. Save model and optimizer state, scheduler state, randomness, and data progress. Define Recovery Point Objective, or RPO, and Recovery Time Objective, or RTO, as acceptance criteria.

Place independent serving replicas in distinct failure domains where capacity permits. Do not confuse a regional control plane with a TPU slice spanning zones; these are different architectural concerns. Plan topology and capacity before finalizing availability requirements. (docs.cloud.google.com)

Serving needs application-aware health checks

For an HTTP service, use startupProbe for initialization, readinessProbe to control traffic admission, and livenessProbe for failures that a restart can resolve. Aggressive probes during compilation can cause unnecessary restarts. (kubernetes.io)

The proposed production policy is to become ready only after loading the model and completing required warmup. Reserve headroom for failures and upgrades. Evaluate autoscaling against queue depth and latency rather than CPU alone. Do not apply an HPA to distributed training as though every process were an independent worker.

Development versus production

During upgrades, change one dimension at a time: image, PyTorch, TorchTPU, runtime, or TPU generation. Retain the previous image and a compatible checkpoint. For inference, use canary traffic and routing rollback. For training, validate resume behavior before retiring the previous path.

7. Security and observability

Use a dedicated ServiceAccount and Workload Identity Federation for GKE to authorize access to Google Cloud services instead of distributing service-account keys. Creating a Kubernetes ServiceAccount does not itself grant Identity and Access Management, or IAM, permissions. (docs.cloud.google.com)

Apply least-privilege Role-Based Access Control, or RBAC; namespace separation; NetworkPolicy traffic restrictions; TLS; protected Secrets; and controlled artifact access. Do not expose profilers or notebooks publicly. Validate metadata-server access and required destinations before tightening egress rules. (kubernetes.io)

For observability, I recommend three dashboards:

• Business outcomes: successful runs, completion deadlines, quality, and cost per result.
• Application behavior: latency, throughput, errors, queue depth, and time without progress.
• Infrastructure: CPU, RAM, restarts, TPU memory, and capacity wait time.

Emit structured logs containing the run identifier, model revision, dependency versions, dtype, and topology. Alert on stalled progress, overdue checkpoints, and quality degradation—not only endpoint downtime.

For device-level investigation, the documentation provides torch.profiler and XProf integration. Check version requirements and capture short representative windows rather than profiling every production step. (torchtpu.readthedocs.io)

8. Evidence-driven troubleshooting

These hypotheses combine the TorchTPU execution model with Kubernetes resource debugging. Confirm the cause before applying a correction. (torchtpu.readthedocs.io)

kubectl get nodes
kubectl get pods -n ml-evaluation -o wide
kubectl describe pod <pod> -n ml-evaluation
kubectl logs <pod> -n ml-evaluation
kubectl logs <pod> -n ml-evaluation --previous
kubectl get events -A --sort-by=.metadata.creationTimestamp

--previous retrieves logs from an earlier container instance in the same Pod. It does not automatically recover logs from deleted Pods. The sample Job also uses restartPolicy: Never, limiting when that option is useful. (kubernetes.io)

9. Cost, alternatives, and the final decision

Use the following evaluation model:

cost per useful result = total cost attributed to the run / results meeting the quality contract

Include accelerator and host charges, storage, networking, observability, preparation, compilation, idle capacity, and failed attempts. Amortize migration engineering separately. Hourly rates and theoretical operations per second are not sufficient decision criteria.

Spot capacity can fit recoverable tasks, but it can be preempted. Evaluate net savings after lost work and recovery time; do not treat it as guaranteed capacity for strict deadlines. (docs.cloud.google.com)

My recommended screening framework is:

These are proposed decision criteria, not a universal ranking. Avoid migration before addressing data bottlenecks, without a team able to operate the new platform, or without a measurable economic or availability objective.

Practical summary

• Separate compatibility from performance, and performance from economics.
• Validate precision, shapes, and dependencies before scaling.
• Measure cold, warm, and end-to-end execution independently.
• Distinguish distributed training from independent serving replicas.
• Make restore, rollback, and available capacity production gates.
• Do not treat a successful smoke test as automatic migration approval.

Conclusion

TorchTPU belongs in a Platform Engineering strategy as an execution option with a testable contract. DevOps provides reproducible packaging and promotion; SRE defines signals, objectives, and recovery; Cloud Engineering validates capacity, isolation, and cost.

Migration is justified when it preserves the business outcome while improving the combination of cost, completion time, or availability—and when the team can explain, measure, and reverse the change.