Google and AWS Release PyTorch Backends for Their AI Chips, But Optimization Remains Complex

At PyTorch Conference North America 2026 this week, Google unveiled TorchTPU, an open-source PyTorch implementation designed for Google TPUs and now hosted on GitHub. Simultaneously, AWS introduced TorchNeuron, currently in public beta, as an open-source PyTorch backend tailored for AWS Trainium chips, also available on GitHub. Both initiatives target the same goal: reducing barriers for developers who wish to explore alternatives to Nvidia's GPU-dominated ecosystem and its CUDA software stack.

The announcements represent meaningful progress in diversifying hardware options for the PyTorch community, yet the path to optimal performance on these new platforms remains challenging. While migration from Nvidia hardware has become less painful, extracting top-tier performance from Google and AWS accelerators still requires substantial additional work.

Google's Push to Simplify TPU Integration

Google initially previewed TorchTPU with select customers in April, following development that began in mid-2025 according to PyTorch. The project emerged from collaboration between Google and Meta's PyTorch team, addressing a long-standing friction point: running PyTorch models on TPUs has historically required workarounds and compromises.

The new backend eliminates JAX dependencies from the runtime, enabling developers to port existing PyTorch workloads with minimal code modifications. Rather than extensive rewrites, teams can achieve migration through what PyTorch describes as "trivial code changes." Developers gain access to eager execution for model debugging and torch.compile for TPU optimization and memory reduction.

Google's broader conference announcements included integrations spanning Hugging Face, vLLM, SGLang, llm-d, TorchTitan, Ray and TorchRec. The company also contributed LiteRT to the PyTorch ecosystem. Google characterizes these additions as "designed to help you run PyTorch effectively on Google Cloud or in your own infrastructure."

AWS Brings Native PyTorch Support to Trainium

AWS's TorchNeuron extends the native PyTorch approach to its own accelerator lineup. The public beta supports Trainium2 and Trainium3 instances, enabling developers to build and experiment with models directly on Trainium hardware, which AWS designed for deep learning training and inference workloads.

Migration can be straightforward: developers switch to .to('neuron') for eager execution. Compiled models require selecting the Neuron backend, either as the default or within a torch.compile call. TorchNeuron also integrates with Neuron Kernel Interface (NKI), a low-level language and compiler for AWS Trainium and Inferentia NeuronDevices, permitting developers to build custom operators and fine-tune Trainium performance. AWS and PyTorch have published detailed guidance on execution modes and performance tradeoffs.

Progress Comes With Caveats

Both TorchTPU and TorchNeuron represent genuine advances in hardware diversity for PyTorch, but neither solves the underlying challenge of cross-platform optimization. While basic functionality requires only modest code adjustments, reaching peak performance demands platform-specific tuning.

For distributed workloads on TPUs, PyTorch notes that TorchTPU's default Skinny Collectives approach incurs an end-to-end performance penalty of roughly 30–40% compared with optimized configurations. This penalty stems from the XLA compiler's inability to overlap computation and communication under the default setup.

PyTorch recommends that teams pursuing peak performance pursue additional tuning and production-grade stacks: "Of course, for production workloads seeking peak performance we would absolutely recommend pursuing an optimal approach based on Pallas kernels and production-grade stacks such as vLLM-TorchTPU, SGLang-TorchTPU, or TorchTitan-TorchTPU, as examples."

AWS similarly positions TorchNeuron as enabling teams to adapt models "in most cases" through configuration changes rather than code or model rewrites, implicitly acknowledging that certain workloads require deeper modifications. A case study involving AWS's Generative AI Innovation Center and AI company i14 illustrates this reality: adapting the Laplace 3B model for Trainium2 necessitated three distinct changes, including native PyTorch implementation for the GLA attention kernel, redesigned Muon communication, and sequential backward pass compilation.

These releases expand the toolkit available to PyTorch teams seeking alternatives to Nvidia hardware. For now, the practical advantage lies in easing initial migration; achieving competitive performance on unfamiliar accelerators remains a separate undertaking requiring specialized expertise.