PyTorch is the framework behind most AI models, used by an estimated 85% of AI applications. It was built with Nvidia GPUs in mind, which is one reason Nvidia's grip on AI computing has been so hard to loosen. This week the two biggest cloud providers moved at the same time. On 19 October Amazon released TorchNeuron for its Trainium chips, and on 20 October, at PyTorch Conference 2026, Google open-sourced TorchTPU for its Tensor Processing Units, built together with Meta's PyTorch team.
What changes in your code
Both projects plug into PyTorch through its official interface for new hardware, so tensors are ordinary PyTorch tensors. For most models, moving to a TPU is about choosing the device.
# pip install torch_tpu (on a Cloud TPU VM)
import torch
import torch_tpu # registers the "tpu" device
device = torch.device("tpu")
model = MyModel().to(device)
batch = batch.to(device)
loss = model(batch).loss
loss.backward()
# Optional: compile for extra speed
model = torch.compile(model, backend="tpu")import torch
device = torch.device("cuda")
model = MyModel().to(device)
batch = batch.to(device)
loss = model(batch).loss
loss.backward()
model = torch.compile(model)How TorchTPU works
- Strict mode (default)
- Runs each operation as normal PyTorch, asynchronously. Easiest to debug.
- Debug mode
- Synchronous, with bounds checks, for tracking down errors.
- DeferAndFuse
- Groups operations together to reduce memory traffic; Google reports 15% faster in its fine-tuning test, and 50% or more in some cases.
- torch.compile(backend="tpu")
- Captures the whole graph and compiles it through OpenXLA into TPU machine code, cached across restarts.
- Distributed training
- DDP, FSDP2 and DTensor through a TPU process group, for both identical and divergent workloads.
- Custom kernels
- Pallas, Tokamax and Helion kernels for when you need hand-tuned speed.
Google-reported TorchTitan example. Seconds per training step.
TorchTPU, PyTorch/XLA and TorchNeuron compared
TorchTPU
Google TPUs- Status
- open source, v0.2.0, Apache 2.0
- Chips
- TPU v5e and newer, up to Ironwood
- Feels like normal PyTorch, eager by default
- vLLM, SGLang and Transformers (device_map="tpu") support
Google's recommended path for all new TPU work.
PyTorch/XLA
Google TPUs, older- Status
- being replaced by TorchTPU
- Records a graph lazily, so you have to mark what to compile
- Only supports identical work on every chip
Plan a migration if you use it today.
TorchNeuron
Amazon Trainium- Status
- public beta (Neuron 2.33)
- Chips
- Trainium 2 and Trainium 3
- Eager plus torch.compile through the Neuron compiler
- FSDP, DTensor, DDP and tensor parallelism; serving with vLLM Neuron
Amazon is also upstreaming code to make PyTorch less tied to Nvidia.
Should you move?
- 1
Check what you actually run
Fine-tuning a Hugging Face model or serving it with vLLM is the easiest case. Custom CUDA kernels are the hardest.
- 2
Price the same job on each chip
Compare cost per training step or per million tokens served, not price per hour.
- 3
Start with a small, real workload
Port one fine-tuning or inference job and compare speed and results.
- 4
Keep code portable
Use device-agnostic code and torch.compile, so switching chips stays a config change.
Check your understanding
1/3How do you point PyTorch code at a TPU with TorchTPU?
TorchTPU plugs into PyTorch's official interface for new hardware.
2/3What happens to PyTorch/XLA?
Google recommends TorchTPU for all new TPU workloads.
3/3Which chips does TorchNeuron target?
It is in public beta as part of Neuron 2.33.
You got 0 out of 3.
More for Python and ML engineers: what's new in scikit-learn 1.9, Python 3.15 and Apache Spark 4.2.
Sources: MarkTechPost on TorchTPU (20 October 2026); TechCrunch, "Google has joined Amazon in attacking Nvidia's stronghold with AI app developers" (20 October 2026); the TorchTPU README. Performance figures are Google's own; check the official documentation before installing.
Alessandro Ecclesie AgazziWritten by Alessandro Ecclesie Agazzi, analytics engineer in London. Updated 21 October 2026. About me