AWS TorchNeuron Runs PyTorch Code by Specifying Neuron Instead of CUDA

Details

AWS has released TorchNeuron, a native PyTorch backend for Trainium, as an Apache 2.0 open-source project
By specifying the device as neuron instead of cuda, standard PyTorch code runs on Trainium, with support limited to Trn2 and Trn3
When integrated with Hugging Face libraries, 89% of the top 100 text generation models work without additional work
The neuron device is registered through the PyTorch PrivateUse1 extension point, allowing users to choose between immediate execution or torch.compile mode
torch.compile does not support dynamic shapes, making the dynamic=False setting and static KV cache, and input padding mandatory
Applying NKI kernels to the ViT encoder measured a maximum speedup of 2.45x (based on DINOv3 ViT-L/16)
It supports PyTorch 2.14.0 and Python 3.12~3.14, with standalone wheel installation available on Ubuntu 26.04
Some features, such as sparse tensors, NestedTensor, and JVP, are not supported in the initial release

Key point