TorchTPU is now public, and PyTorch/XLA is marked deprecated. That does not mean an existing training job is ready for a one-line device swap. The new backend aims to preserve ordinary PyTorch semantics, but your shapes, custom operators, distributed setup, and throughput target still need to survive a real pilot.

Three gates for a TorchTPU migration pilot

The useful change is in the execution model. PyTorch/XLA used lazy tensor graph recording and required explicit steps in many workflows. TorchTPU supports eager execution for debugging and torch.compile for compiled execution. The PyTorch team describes a first port as changing model.to("cuda") to model.to("tpu"), then immediately adds the caveat that this is not a promise of peak performance. That caveat is the whole migration story.

A migration pilot you can actually run

Start on a TPU VM with the official TorchTPU installation instructions, not your existing production image. The project README currently lists Python 3.11+ and a pinned torch_tpu==0.2.0 install command using PyTorch's CPU wheel index. Pin the environment. Do not let a general-purpose install silently choose an incompatible Torch build.

Move a small representative batch through the model in eager mode first. The Google Cloud documentation's smoke test creates tensors on torch.device("tpu"), runs a matrix multiply, copies the result back to CPU, and prints its shape and sum. Adapt that idea to your actual model: feed one fixed input, synchronize where your code requires it, and compare output tensors with the same input on your current backend. Use tolerances appropriate to your dtype. A successful import only proves that Python found the package.

Next, enumerate the things that tend to hide inside a training loop: custom C++ or Triton ops, assumptions about CUDA streams, device-specific memory calls, dynamic sequence lengths, checkpoint load and save, optimizer state, and callbacks that touch tensors on rank zero only. The PyTorch post says TorchTPU supports ~450 ATen ops and validated a set of popular models. That is useful evidence, not proof for your model's exact operator path. Search the issue tracker and run your own end-to-end forward and backward steps.

Keep the first pass boring. Run one worker, a short fixed batch, and a checkpoint round trip. Then test your actual distributed mode. TorchTPU exposes native distributed paths including DDP, FSDP2 and DTensor, but a workload that depended on XLA-specific behavior still needs its own migration. Confirm that every rank starts, gradients agree within tolerance, and a resumed job produces the expected next step. Save logs and a known-good checkpoint before changing several variables at once.

Only then compile. Use eager mode to find correctness problems; use torch.compile after the basic run is sound. Measure compilation separately from steady-state steps. If input shapes vary, record how many distinct shapes appear and whether the workload recompiles. A first compiled request can include compile time, so comparing it with a warmed-up GPU run produces a meaningless result. Track throughput, peak device memory, host overhead, and failure rate over a long enough window to include checkpointing and your normal input variation.

What the published numbers do not guarantee

Google reports up to 52% MFU (Model FLOPs Utilization) for Qwen 3 8B with TorchTitan, with scaling across as many as 8,192 Ironwood TPUs. That is a specific vendor-reported training result on a named workload and hardware. It does not tell you what a smaller model, different batch shape, custom op, or inference service will achieve. Treat it as evidence that the stack can run at scale, not as a forecast for your bill or tokens per second.

There is a real upside to testing now: the old PyTorch/XLA repository says it is deprecated and scheduled for archival on November 15, 2026. But a deadline is not a reason to skip a parity test. Checkpoint portability, numerical drift, supported libraries, failure recovery, and the cost of your chosen TPU slice belong in the pilot report. A device swap can get code executing; it cannot decide whether that execution is operationally acceptable.

Keep a small acceptance sheet beside the code. Record the exact package versions, TPU generation and slice size, model revision, dtype, batch shape, warm-up count, compile duration, steady-state throughput, and peak memory. Attach output-difference checks and the checkpoint-resume result. If the change touches a custom operator, test both its common and boundary inputs. This gives the team something more useful than a screenshot of a successful import, and it makes a later regression easier to trace.

A pilot should also have a stop rule. If output parity fails, the job recompiles repeatedly on normal inputs, or sustained throughput misses the target after warm-up, do not expand the slice and hope the numbers improve. Isolate the case, check the backend's current support matrix, and rerun with one variable changed. Keep the old environment available until the replacement has passed the same data, recovery, and cost checks as the workload it is meant to replace.

Sources