The AWS Neuron SDK is the software development kit for deep learning and generative AI on AWS Inferentia and Trainium. It is what you use in place of CUDA on those chips.
What it includes
Per the Neuron documentation, it supports:
- PyTorch (native, through TorchNeuron, labelled Beta) and JAX (Beta integration), for training and inference.
- Custom kernels through the Neuron Kernel Interface (NKI).
- Direct use of the Neuron compiler and runtime.
AWS's Trainium page adds Neuron Explorer and Neuron Monitor for profiling and health metrics.
Why it matters
The software stack decides which models run well on a chip. Check it against your model before you commit, as the PyTorch on Trainium page explains. The two chip lines it serves are compared in Trainium vs Inferentia.
See also NVIDIA's main competitors in AI.
