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

About the AWS Neuron SDK#

AWS Neuron is the software development kit (SDK) for machine learning on AWS Trainium and AWS Inferentia, the machine learning accelerators that AWS designs. Use Neuron to train, serve, profile, and monitor deep learning and generative AI workloads on Trainium and Inferentia instances, with the frameworks and libraries that you already use.

The AWS Neuron SDK includes:

  • TorchNeuron (native PyTorch) - Runs PyTorch on Trainium as a native device, with eager mode, torch.compile, and standard distributed APIs such as FSDP, DDP, and DTensor

  • TorchTitan Neuron - Pre-trains large models on Trainium with TorchTitan and TorchNeuron

  • JAX support - Runs JAX workloads on Neuron through a PJRT plugin

  • vLLM Neuron - Serves large language models (LLMs) with the standard vllm serve command and an OpenAI-compatible API

  • vLLM Omni Neuron - Runs diffusion models for image and video generation with vLLM Omni

  • Neuron Kernel Interface (NKI) - Writes custom kernels with direct access to the NeuronCore instruction set, plus pre-optimized kernels in the NKI Library

  • Neuron Compiler - Compiles models from ML frameworks into optimized code for Neuron devices

  • Neuron Runtime - Loads and runs compiled models on NeuronCores, and provides collective communication for distributed workloads

  • Neuron Explorer - Profiles workloads at the framework, model, and NeuronCore level to find and fix performance bottlenecks

  • Neuron Observer - Collects device health, hardware fault, and performance metrics from Trainium clusters

  • Neuron system tools - Command-line tools for device monitoring and administration

  • Deployment support - Deep Learning AMIs (DLAMIs), Deep Learning Containers (DLCs), and integration with Amazon EC2, Amazon EKS, Amazon ECS, AWS Batch, and Amazon SageMaker

For an overview of each component, see What is AWS Neuron?. To run your first workload, see the Neuron quickstarts.

About the “NeuronX” name#

You see “NeuronX” (or “neuronx”) in the names of the Neuron packages for Inf2, Trn1, Trn2, and Trn3 instances, such as the neuronx-cc compiler, the torch-neuronx package for TorchNeuron, and the aws-neuronx-* runtime, driver, and tools packages. Packages without the “x”, such as neuron-cc and torch-neuron, are for first-generation Inf1 instances. Inf1 is no longer supported, and its documentation is in the archive.

Catch up on the latest Neuron news#

What’s New in Neuron

Read about the latest releases and features of the Neuron SDK

Learn about AWS Neuron#

What is AWS Neuron?

Short overview of the AWS Neuron SDK and its components

Neuron architecture

Understand the Neuron hardware and software architecture

Neuron features

Overviews of model development features provided by Neuron

TorchNeuron (native PyTorch)

Run PyTorch on Trainium with eager mode, torch.compile, and standard distributed APIs

JAX on Neuron

Run JAX workloads on Neuron with the PJRT plugin

Training on Neuron

Train models with TorchNeuron and pre-train at scale with TorchTitan

Inference on Neuron

Choose an inference path: vLLM Neuron, vLLM Omni Neuron, or native PyTorch

vLLM Neuron

Serve LLMs on Trainium with vLLM and an OpenAI-compatible API

vLLM Omni Neuron

Generate images and video with diffusion models on Trainium

Neuron Kernel Interface (NKI)

Write custom kernels with direct access to the NeuronCore, or use the NKI Library

Neuron Compiler

The Neuron compiler optimizes models for Neuron hardware

Neuron Runtime

Runtime for executing compiled models on Neuron devices

Neuron developer tools

Profile with Neuron Explorer, monitor cluster health with Neuron Observer, and use system tools

Deep Learning AMIs and Containers

Start from preconfigured Neuron DLAMIs and DLCs

Deploy on AWS

Deploy Neuron workloads on Amazon EC2, Amazon EKS, Amazon ECS, AWS Batch, and more

Resources#

Support#

Contact us#

For support, submit a request with AWS Neuron GitHub issues.

If you want to request a feature or report a critical issue, you can contact us directly at aws-neuron-support@amazon.com.

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