Amazon Unveils TorchNeuron to Ease PyTorch Use on Trainium Chips

Amazon has introduced TorchNeuron, a PyTorch backend that lets developers move existing AI models onto its Trainium chips with fewer code changes, easing a hurdle in competing with Nvidia's hardware.

TorchNeuron adds Trainium2 and Trainium3 as device options within PyTorch, a framework used to build AI models. Developers can direct a model to the chips with a setting such as .to("neuron"). Programs that already use PyTorch's hardware detection interface may need no device setting change, while compiled models must select the Neuron backend.

The public code repository identifies TorchNeuron as beta software and carries an Apache 2.0 license. The beta requires fixed input shapes for compiled models and does not support some PyTorch features, including sparse tensors. Developers will need to check whether their existing models run as intended.

The software makes it easier to test Amazon's chips alongside Nvidia-based systems. It does not, on its own, establish a performance or cost advantage for a particular workload. Amazon plans to present TorchNeuron at the PyTorch Conference in San Jose, California, on Oct. 20.