This document is relevant for: Inf1, Inf2, Trn1, Trn2

ML framework support on AWS Neuron SDK#

AWS Neuron provides integration with popular machine learning frameworks, enabling you to accelerate your existing models on AWS Inferentia and Trainium with minimal code changes. Choose from our comprehensive framework support to optimize your inference and training workloads.

PyTorch on AWS Neuron

Complete PyTorch integration for both inference and training on all Neuron hardware.

  • PyTorch NeuronX - Inf2, Trn1, Trn2 (inference & training)

  • PyTorch Neuron - Inf1 (inference only)

  • Native PyTorch API compatibility

JAX on AWS Neuron

Beta release

Experimental JAX support with Neuron Kernel Interface (NKI) integration.

  • JAX NeuronX - Neuron hardware support

  • Research and development focus

  • Status: Beta - active

Hardware compatibility matrix#

Framework

Inf1

Inf2

Trn1/Trn1n

Trn2

Inference

Training

torch-neuronx

N/A

✅

✅

✅

✅

✅

torch-neuron

✅

N/A

N/A

N/A

✅

N/A

JAX NeuronX

N/A

✅

✅

N/A

✅

N/A

TensorFlow Neuron

✅

N/A

N/A

N/A

✅

N/A

This document is relevant for: Inf1, Inf2, Trn1, Trn2