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

Install AWS Neuron SDK#

Install the AWS Neuron SDK to enable deep learning acceleration on Inferentia and Trainium instances.

Note

New to Neuron? Start with the quickstart guide for a complete end-to-end tutorial.

Quick Start Decision Tree#

Answer these questions to find your installation path:

1. Which framework?

2. Installation method?

  • AWS Deep Learning AMI — fastest setup, pre-configured with all dependencies. Best for getting started and single-user development.

  • Deep Learning Container — Docker-based, portable across EC2, ECS, EKS. Best for production deployments and CI/CD pipelines.

  • Manual installation — full control over packages and versions. Best for custom OS images and shared clusters.

Instance Comparison#

Instance

NeuronCore

Use Case

Status

Trn3

v4

Training and inference (latest generation)

Current

Trn2

v3

Training and inference

Current

Trn1

v2

Training and inference

Current

Inf2

v2

Inference

Current

Inf1

v1

Legacy inference

Legacy

Installation by Framework#

PyTorch

Recommended for most users

  • PyTorch 2.13 or 2.14 with TorchNeuron, the native Neuron backend

  • Eager mode and torch.compile

  • Supports: Trn2, Trn3

Most Popular

JAX

For JAX users

  • JAX 0.7.0 - 0.10.0 with Neuron backend

  • XLA compilation

  • Supports: Inf2, Trn1, Trn2, Trn3

Multi-framework DLAMI#

🚀 Neuron multi-framework DLAMI

Pre-configured AMI with PyTorch, JAX, and vLLM virtual environments ready to use. The fastest way to get started with any framework.

Common Issues#

⚠️ Module not found errors

If you see “No module named ‘torch_neuronx’” or similar:

  1. Verify virtual environment is activated

  2. Check Python version: python --version (3.12 or later for TorchNeuron)

  3. Reinstall: pip install --force-reinstall torch-neuronx

See Installation Troubleshooting for more details.

⚠️ Instance type not recognized

Ensure you’re using a Neuron-supported instance:

  • Check with: aws ec2 describe-instance-types --instance-types <type>

  • Verify Neuron devices: neuron-ls

See Installation Troubleshooting for more details.

⚠️ Version compatibility issues

Check version compatibility:

See Installation Troubleshooting for more details.

Legacy Inf1 Support#

Warning

Inf1 uses the legacy NeuronCore v1 architecture. For new projects, use Trn2 or Trn3.

  • Current-generation NeuronCores (v3 on Trn2, v4 on Trn3)

  • Broader framework support (TorchNeuron, JAX, vLLM Neuron)

  • Active development and feature updates

Inf1 Installation (Legacy, Archived)

Install Neuron SDK for Inferentia 1 instances (archived)

Archived

Additional Resources#

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