About the AWS Neuron SDK#

AWS Neuron is a software development kit (SDK) enabling high-performance deep learning acceleration using AWS Inferentia and Trainium, AWS’s custom designed machine learning accelerators. It enables you to develop, profile, and deploy high-performance machine learning workloads on AWS Inferentia and Trainium instances.

The AWS Neuron SDK includes:

  • Neuron Compiler - Compiles high-level, framework-based models for optimal performance on Neuron devices

  • Neuron Kernel Interface (NKI) - Provides direct compiler access to Neuron device capabilities

  • Neuron Runtime - Executes compiled models on Neuron devices

  • ML Framework integration - Deep support for PyTorch and JAX

  • Training and inference libraries - Distributable training and inference libraries for large-scale models

  • Deployment support - Integration with AWS services like SageMaker, EC2, EKS, and ECS

  • Developer tools - Profiling, monitoring, and debugging utilities

For a full list of AWS Neuron features, see What is AWS Neuron?.

Join our Beta program

Get early access to new Neuron features and tools! Fill out this form and apply to join our Beta program.

What is “NeuronX”?#

“NeuronX” refers to the next-generation AWS Neuron SDK, which provides enhanced capabilities for both inference and training on AWS Inferentia and Trainium instances. NeuronX includes:

  • Support for the latest versions of PyTorch and JAX

  • Advanced compiler optimizations for improved performance

  • Enhanced distributed training libraries for large-scale models

  • Improved profiling and debugging tools

  • Ongoing feature development and support for new instance types

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

Supported ML frameworks

Neuron support for popular ML frameworks including PyTorch and JAX

NeuronX distributed (NxD) libraries

NeuronX distributed libraries for training and inference

Neuron Kernel Interface (NKI)

NKI is a low-level interface for custom, bare-metal kernel development

Neuron Compiler

The Neuron compiler optimizes models for Neuron hardware

Neuron Runtime

Runtime for executing compiled models on Neuron devices

Neuron developer tools

Tools for profiling, debugging, and monitoring Neuron applications

Neuron AWS Neuron Deep Learning AMIs

Deploy the Neuron SDK on EC2 instances with pre-installed Amazon Machine Images (AMIs)

Neuron AWS Neuron Deep Learning Containers

Deploy the Neuron SDK using pre-built Docker deep learning containers (DLCs)

Resources#

Support#

Contact us#

For support, submit a request with AWS Neuron Github issues or visit the Neuron AWS forums for an answer.

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