TorchNeuron: Amazon runs PyTorch code on its Trainium chips by changing just one line of configuration

With TorchNeuron, AWS aims to neutralise Nvidia’s software advantage: a PyTorch application can switch to Trainium without any rewriting, and Amazon is contributing upstream to make PyTorch hardware-agnostic.

On 19 October, AWS launched TorchNeuron, a tool that enables applications written in PyTorch to run on Trainium’s in-house chips with just a simple configuration change, without altering the code. The stakes are high: according to Amazon, around 85 per cent of AI applications are based on PyTorch, and it is largely Nvidia’s software ecosystem that currently keeps developers on its GPUs.

Breaking the software lock-in

At the same time, Amazon is contributing code back to the PyTorch project itself to remove dependencies specific to any particular chip manufacturer. Colin Brace, vice-president at Annapurna Labs – the subsidiary that designs AWS’s chips – highlights Trainium’s price-performance ratio and the fact that a single platform can be used for both training and inference.

Amazon cites the example of Deepgram, a specialist in speech recognition, which reportedly processes 70 per cent more audio per dollar spent since migrating to Trainium, whilst maintaining the same level of accuracy in production.