Welcome to AWS Neuron#
AWS Neuron is the software development kit (SDK) used to run deep learning and generative AI workloads on AWS Inferentia and AWS Trainium powered Amazon EC2 instances (Amazon EC2 Inf1, Inf2, Trn1 and Trn2 instances). It includes a compiler, runtime, training and inference libraries, and profiling tools. Neuron supports customers in their end-to-end ML development lifecycle including building and deploying deep learning and AI models.
ML Frameworks and Libraries - Neuron integrates with PyTorch and JAX, and offers NxD Training and NxD Inference PyTorch libraries for distributed workflows. It also supports Hugging Face Optimum Neuron, PyTorch Lightning, and AXLearn library for JAX model training.
Developer Tools - Neuron provides health monitoring, observability, and profiling tools for AWS Inferentia and Trainium instances. It tracks hardware utilization, model execution metrics, and device information. The Neuron Profiler identifies performance bottlenecks. Neuron also integrates with third-party monitoring tools like Datadog and Weights and Biases.
Compute Kernels - Neuron Kernel Interface (NKI) provides direct hardware access on AWS Trainium and Inferentia, enabling customer to write optimized kernel. NKI provides a Python-based environment with Triton-like syntax. Neuron supports custom C++ operators, allowing developers to extend functionality and enhance deep learning models.
Workloads Orchestrations and Managed Services - Neuron enables you to use Trainium and Inferentia-based instances with Amazon services such as SageMaker, EKS, ECS, ParallelCluster, and Batch. and third-party solutions like Ray (Anyscale) and Domino Data Lab.
Architecture - To understand the architecture of AWS AI Chips, Trn/Inf instances, and NeuronCores visit Trn and Inf Instances Architecture, AWS AI Chips Architecture and AWS NeuronCore Architecture.
For more information about the latest AWS Neuron release, see Neuron 2.21.0 Beta (12/03/2024) and check the Announcements page.
For list of AWS Neuron model samples and tutorials on Amazon EC2 Inf1, Inf2, Trn1, and Trn2 instances,
see Model samples and tutorials.