On Oct. 20, Amazon opened TorchNeuron to the public as a beta. It is free, open-source software that lets PyTorch run on Amazon's Trainium chips. Amazon says 85% of AI researchers and developers write their code in PyTorch. Amazon first announced TorchNeuron at re:Invent in December, so the news is the release, not the idea. The pitch is one change: point PyTorch at a device called neuron. AWS's Colin Brace, quoted by TechCrunch, put the goal this way: "The real story is removing the lock-in that kept developers from having a real choice at all."
That one line is real. It is also the easy part. Cartesia, a voice AI company quoted in the announcement, ran its code with minimal changes. Then it put its effort into kernels, the low-level programs that run on the chip itself. It wrote them in NKI, Amazon's language for its own chips, for the layers that run most often and set its response time. Deepgram says Trainium handles about 60% more audio per dollar than its GPU baselines, after a multi-month port with AWS's Generative AI Innovation Center. Both accounts come from Amazon's announcement, not from independent tests. The first release also leaves out dynamic shapes, meaning inputs that change size. A changed size means recompiling, and some operations have no Trainium version yet.
My read: the code moves easily, but the speed comes from kernels written for Trainium, and those would not run on another chip. Lock-in has not ended. It has moved into the kernels. Before you price a move, ask AWS which of your workloads would need hand-tuned kernels and who writes them. Then ask for the cost of moving back, quoted next to the cost of moving in.
Update, Oct. 21: Google released TorchTPU, a similar tool for its own TPU chips. Same question, other cloud.