TorchTPU and TorchNeuron go after Nvidia’s real moat: PyTorch’s GPU defaults

Google has released TorchTPU as an open-source project, letting PyTorch apps run on its Tensor Processing Units and on Google Cloud through a configuration change rather than a rewrite, TechCrunch reports. A day earlier, Amazon released TorchNeuron for its Trainium chips and said it is submitting code to the main PyTorch project to decouple features from Nvidia hardware (Amazon’s TorchNeuron). PyTorch, developed by Meta, then Facebook, in 2016 and handed to the Linux Foundation in 2022, is the toolkit most AI apps are built on, by some accounts 85% of them, and it grew up around Nvidia GPUs. Hugging Face detailed the TorchTPU integration separately.

Google says it worked directly with Meta, Hugging Face, Inferact, the maker of vLLM, and RadixArk, the maker of SGLang, so TorchTPU behaves as developers expect with tools they already use. It also shipped a TPU agent suite to help port apps and fix errors, and updated LiteRT for on-device PyTorch apps.

Nvidia’s lock-in has always been as much about code as silicon: PyTorch features were built for its GPUs because that is what researchers had in 2016. Making a move to TPUs a configuration change goes directly at switching costs, and the compatibility work with inference engines is what tends to move developers. For AI startups, easier portability adds negotiating leverage on compute prices. Nvidia, meanwhile, has been strengthening its own developer tooling (Nsight Python) and financing custom-silicon efforts that still use its infrastructure (Nvidia and MediaTek).

Amazon’s move may matter more over time. Upstream code that decouples PyTorch from Nvidia would benefit every hardware maker, not only AWS; a vendor-specific layer helps one cloud, while a more neutral core helps the field. Whether that code is accepted, and how the Linux Foundation-governed project handles it, will be telling.

TechCrunch flags the caution: easy portability does not mean developers will leave Nvidia en masse. Performance, tooling depth and habit are sticky, and a port that works is not necessarily fast or cheap. Independent benchmarks of popular PyTorch workloads on TorchTPU versus Nvidia GPUs will matter. There is also a trade-off: escaping one chip vendor by moving to one cloud’s proprietary hardware creates a different dependency rather than freedom. The best outcome for developers is PyTorch that runs well anywhere. More in our Google coverage, Nvidia coverage and AI tools coverage.