DAdhruuv.me / blog
← All articles

4 min readdev.to

How PyTorch Became the Socket Every AI Chip Plugs Into

Google and AWS released PyTorch backends a day apart, both through the same small hook. Here's what that hook is, and why getting code to run on a new chip is the easy half.

On 19 October, AWS released TorchNeuron, which lets PyTorch run directly on its Trainium chips. A day later, Google and Meta released TorchTPU, which does the same for Google's TPUs. Both arrived at the PyTorch Conference, both are open source, and both connect to PyTorch through the same small hook.

That last part is the story. Two companies with no reason to coordinate reached the same conclusion: the road onto a working engineer's laptop runs through one extension point in PyTorch.

What happens when you call .to("cuda")

Write a + b on two tensors and PyTorch does not immediately add anything. It looks up a name, aten::add, in a table.

The part doing that lookup is called the dispatcher. It checks where the tensors live, reads a label off them — CPU, CUDA, and so on — and uses it to pick the function that actually does the addition on that hardware. Those hardware-specific functions are called kernels. A backend is a full set of kernels filed under one label.

PyTorch keeps one label deliberately empty: PrivateUse1. Any hardware vendor can claim it and register their kernels against it from their own repository, without needing their code accepted into PyTorch itself. TorchTPU registers there. TorchNeuron uses the same registration mechanism. A new chip becomes a new label, and model.to("tpu") is what that looks like from your code.

aten::add routed by the PyTorch dispatcher to CPU, CUDA or PrivateUse1, with TorchTPU and TorchNeuron both registering against PrivateUse1
A backend is a set of kernels filed under one dispatch label. TorchTPU names PrivateUse1; AWS says only "PyTorch's built-in mechanism for registering new hardware". Sources: PyTorch blog; Amazon; InfoQ.

This replaces something clumsier. PyTorch on TPUs used to go through JAX, Google's other framework. The TorchTPU team is blunt: the old approach "shoehorn[ed] PyTorch into JAX's modus operandi." The new backend has no JAX in it at runtime, and it runs operations one at a time as you write them, which is how most people debug. The chip speaks PyTorch now, instead of being translated first.

The coverage looks real. Google tested 4,917 models from Hugging Face: 98% ran without crashing, and of those it checked closely, 97% produced the same numbers a GPU would. Amazon says it tests TorchNeuron against hundreds of thousands of models on every build.

Running is the easy half

Here is the TorchTPU team's own framing, caveat included:

"Porting a workload to TPU should be as simple as changing model.to("cuda") to model.to("tpu"). (One cannot achieve peak TPU performance with that simple change... however, the model should run out of the box.)"

Read the bit in brackets, because that is where the work actually is.

A full-width bar labelled 'it runs' shrinking to 40% as default network settings, precision, missing operations and the process layout each take a share
Illustrative: the bar shows the direction of each documented caveat, not a measured benchmark. Source: PyTorch blog, TorchTPU.

  • The numbers come out different. On CUDA, PyTorch multiplies in full 32-bit precision by default. TorchTPU defaults to a faster, lower-precision mode. Your model runs. Its output drifts.
  • The default network settings cost a third of your speed. TorchTPU ships with an option called "Skinny Collectives" switched on, which Google says can cost 30–40% end to end. Collectives are how chips pass results to each other while training.
  • Not every operation exists yet. Around 450 of PyTorch's built-in operations are supported. FlexAttention was not, in one of Google's own examples.
  • The process layout changes. TorchTPU runs one process per chip; JAX runs one per machine. Reinforcement-learning setups that share a chip between jobs need rework. torch.compile also cannot handle changing input shapes yet.

So the part that looked hard, getting the code to run, is nearly solved. The part nobody puts in a launch announcement is that your tuned kernels, your checked numbers and your network settings do not move with .to(). That is months of work per platform, and it is where lock-in actually lives now.

What to do

  • If you are evaluating other chips, stop scoring them on "does our model run." That question is becoming free. Score them on how much of your tuned performance survives the move.
  • Set precision explicitly in any comparison across hardware, or you will blame the chip for a default.
  • If you maintain a library that sits close to PyTorch, PrivateUse1 is the contract worth reading now. It is how your code reaches hardware you will never own.

Where this breaks down

Most of TorchTPU's performance numbers come from Google, and most of its examples do not compare against a GPU on equal footing. The 64-chip test against JAX/MaxText (1,828 against 1,937 tokens per second) is one of the few straight comparisons, and it is against JAX, not a GPU. The project calls itself "in active development." AWS does not mention CUDA at all. Putting these two releases in one frame is my argument, not either company's claim.

Takeaway

The moat was never CUDA the language. It was the thousands of hand-tuned kernels underneath it, and PrivateUse1 is the agreed-on door for building a rival set.

If switching chips costs one line of code and three months of tuning, is that portable?

Sources