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. What’s your use case?

  • Training ML models → Use Trn1, Trn2, or Trn3

  • Running inference → Use Inf2, Trn1, Trn2, or Trn3

  • Legacy Inf1 support → See Legacy Inf1 Support

2. Which framework?

3. 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.9+ with Native Neuron support

  • Eager mode and torch.compile

  • Supports: Inf2, Trn1, Trn2, Trn3

Most Popular

JAX

For JAX users

  • JAX 0.7+ 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 (should be 3.10+)

  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 legacy NeuronCore v1 architecture. For new projects, use Inf2, Trn1, Trn2, or Trn3 with NeuronCore v2.

  • Inf2 offers 3x better price-performance than Inf1

  • Broader framework support (PyTorch 2.x, JAX)

  • Active development and feature updates

Inf1 Installation (Legacy)

Install Neuron SDK for Inferentia 1 instances

Legacy Hardware

Additional Resources#

Other Platforms#

Rocky Linux 9

Install PyTorch Neuron on Rocky Linux 9 using the Rocky-9-EC2-Base AMI. Covers driver and tools setup, then follows the Amazon Linux 2023 guide for framework installation.

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