85% of AI Code Is PyTorch. Google and Amazon Want It Off Nvidia

Google's TorchTPU and Amazon's TorchNeuron let PyTorch apps move to their own AI chips with a one-line change. Here is what launched, why Nvidia should care, and what it means for Indian developers.

• 7 min read

Last updated

Share:
85% of AI Code Is PyTorch. Google and Amazon Want It Off Nvidia

Amazon says 85% of AI researchers and developers write their code in PyTorch. For years, most of that code has quietly assumed one thing: that it will run on an Nvidia GPU.

This week, two of the biggest cloud companies tried to break that assumption on the same stage. Amazon opened up TorchNeuron on Monday. On Tuesday, Google followed with TorchTPU.

Quick answer

  • Google's TorchTPU and Amazon's TorchNeuron are open-source PyTorch backends that run existing AI code on Google TPUs and AWS Trainium chips.
  • The pitch is a configuration change, not a rewrite: in many cases .to("cuda") becomes .to("tpu") or .to("neuron").
  • Nvidia's real moat has been software habits, and both companies are going after exactly that.
  • For Indian developers and students, PyTorch skills now travel across more hardware, so learning it is a safer bet than ever.

Why PyTorch is the real battlefield

PyTorch is the open-source toolkit that developers use to build, train and run AI models. TechCrunch notes that, by some accounts, it is used by 85% of AI apps.

Meta (then Facebook) built it in 2016. Back then, Nvidia GPUs were already the standard hardware for AI research, so many PyTorch features were designed around them.

In 2022, Meta handed PyTorch to the Linux Foundation. But the Nvidia tilt stayed baked into how most AI apps are written by default. That is the advantage Google and Amazon are now attacking.

What Google launched: TorchTPU

Google announced TorchTPU at PyTorch Conference 2026 and released it as open source on GitHub. It lets PyTorch apps run on Google's Tensor Processing Units, or TPUs, and makes it easier to move them to Google Cloud.

It was co-developed with the PyTorch team at Meta. Google says it started the effort in mid 2025.

According to the PyTorch blog post introducing it, TorchTPU so far:

  • supports both eager mode (for debugging) and compiled mode through torch.compile
  • covers about 450 core ATen operations, the building blocks of most deep learning models
  • has been validated out of the box on the 5,000 most downloaded Hugging Face models, including Qwen, DeepSeek, Gemma, Llama and Whisper
  • works with popular tools such as vLLM, SGLang, TorchTitan and Ray

This is not Google's first attempt. In 2020, it released PyTorch/XLA, but the PyTorch blog admits that version needed intrusive code changes and had slow compile times. TorchTPU was built to fix that.

Google also launched a TPU agent suite in preview to help developers move their apps and fix errors. And it updated LiteRT, its tool for running AI on phones, laptops and embedded devices, to work better with PyTorch apps.

What Amazon launched: TorchNeuron

Amazon's version, TorchNeuron, does the same job for its own Trainium chips. Amazon says existing PyTorch code runs on Trainium with a single configuration change and no rewrite.

GeekWire reports that more than 30 companies and internal Amazon teams used it during a private beta. Amazon tests it every night against more than 100,000 models from Hugging Face.

One early user, the speech AI startup Deepgram, says its benchmarks showed Trainium processing about 60% more audio per dollar than its existing GPU setup.

Amazon is also going a step further. It is submitting code to the main PyTorch project that removes ties to specific chips, so PyTorch works with more hardware by default.

Amazon chip business run rate of more than $20 billion vs an estimated $50 billion as a standalone business

Source: GeekWire, citing Andy Jassy's shareholder letter

That chart shows why Amazon cares. In his shareholder letter in April, CEO Andy Jassy wrote that the chip business, including Trainium and Graviton, had an annual revenue run rate of more than $20 billion. He estimated it would be about $50 billion as a standalone chip business.

Google vs Amazon: who did more?

Both tools aim at the same problem, but TechCrunch suggests Google may be one-upping Amazon in a few ways.

Google says it worked directly with Meta, Hugging Face, Inferact (maker of vLLM) and RadixArk (maker of SGLang) so TorchTPU behaves like normal PyTorch inside those tools. In Hugging Face Transformers v5, deploying a model to TPUs is as simple as passing device_map="tpu".

Amazon's edge is its upstream work on PyTorch itself and its list of production users. Anthropic uses Trainium to train and run Claude, and GeekWire reports OpenAI agreed in February to use 2 gigawatts of Trainium capacity.

Google TorchTPUAmazon TorchNeuron
ChipTensor Processing Unit (TPU)Trainium
CloudGoogle CloudAWS
Code change.to("tpu").to("neuron")
Extra toolsTPU agent suite, LiteRT updateUpstream PyTorch changes, open-source kernels
Testing5,000 most downloaded Hugging Face models100,000+ Hugging Face models nightly

Does this end Nvidia's lead?

Not overnight. TechCrunch is clear that developers will not ditch Nvidia hardware en masse just because switching got easier.

Even the PyTorch team is honest about limits. Changing model.to("cuda") to model.to("tpu") should make a model run, but it will not deliver peak TPU performance, because the model's shapes and operations may still be tuned for GPUs.

Still, the pressure is building from several sides. GeekWire notes that Microsoft said last week it plans to run most of its AI services on its own Maia chips. Amazon is also reportedly in talks to buy Decart, a startup that helps move AI models off Nvidia chips, in a deal that could value it at about $7 billion.

What it means for Indian developers

India's AI startups and students mostly work in PyTorch, often on rented cloud GPUs. Here is what changes for them:

  1. Your PyTorch skills get more portable. The same code can now target GPUs, TPUs and Trainium, so time spent learning PyTorch is not tied to one vendor.
  2. Cost comparisons become real. If a workload can move with a config change, you can test it on more than one cloud and compare price per result, the way Deepgram did.
  3. Try before you migrate. TorchTPU installs with pip install torch_tpu, and TorchNeuron with pip install torch-neuronx, so a small test is cheap in engineering time.
  4. Learn eager vs compiled mode. Both tools lean on torch.compile for production speed. Understanding it is useful on any hardware.

Mistakes to avoid

  • Assuming "no code change" means "same speed." Expect to profile and tune before a TPU or Trainium run matches a tuned GPU setup.
  • Ignoring beta labels. TorchNeuron launched as a public beta, and Google's agent suite is in early preview.
  • Switching for hype alone. Move a workload only when your own benchmark shows a cost or speed gain.

FAQs

What is TorchTPU? It is Google's open-source, PyTorch-native backend for its TPU chips. It lets existing PyTorch models run on TPUs with minimal code changes.

What is TorchNeuron? It is Amazon's open-source PyTorch backend for its Trainium chips on AWS. Amazon says existing PyTorch code runs on Trainium with a single configuration change.

Do I need to stop using Nvidia GPUs? No. These tools add options. Your PyTorch code can still run on Nvidia GPUs, and you can test the alternatives side by side.

Are these tools free? The software is open source. Running it still requires paying for TPU or Trainium time on Google Cloud or AWS.

Sources

Tagged Topics

#PyTorch#Google TPU#AWS Trainium#Nvidia#AI Chips#Developers#India

Thanks for reading Kali Zubaan

Explore more articles in Tech & AI or return to the main publication feed.

← Back to All Stories