AWS Neuron

AWS Neuron is the developer stack for running deep learning and generative AI workloads on AWS Trainium and AWS Inferentia. Built on an open-source foundation, Neuron enables developers to build, deploy and explore natively with PyTorch and JAX frameworks and with ML libraries such as HuggingFace, vLLM, PyTorch Lightning, and others without modifying your code.  It includes a compiler, runtime, training and inference libraries, and developer tools for monitoring, profiling, and debugging. Neuron supports your end-to-end machine learning (ML) development lifecycle from building and deploying deep learning and AI models, optimizing to achieve highest performance and lowest cost, and getting deeper insights into model behavior. 

Neuron enables rapid experimentation, production scale training of frontier models, low level performance optimization through the Neuron Kernel Interface (NKI) for custom kernels, cost optimized inference deployment for agentic AI and reinforcement learning workloads, and comprehensive profiling and debugging with Neuron Explorer. 

Neuron enables rapid AI research by running native PyTorch code unchanged on Trainium. Researchers can try new ideas and iterate quickly with PyTorch Eager mode support. Scaling is easy with PyTorch distributed libraries such as, FSDP, DDP, and DTensor to shard models across chips or scale to multiple nodes. Neuron supports torch.compile, and libraries like TorchTitan and HuggingFace Transformers now work directly on Trainium without modification. Also, JAX developers can utilize Neuron to easily develop, optimize and deploy their models to Inferentia and Trainium.   

Neuron optimizes inference economics for agentic AI and reinforcement learning workloads. Standard vLLM V1 APIs work on Trainium and Inferentia with out-of-the-box high performance with features like Expert Parallelism, disaggregated inference, and speculative decoding, and optimized kernels from the Neuron Kernel Library to maximize token economics at scale.  ML developers can train with HuggingFace Optimum Neuron, PyTorch Lightning, and TorchTitan, then deploy inference with standard vLLM APIs. 

Building AI models requires both rapid innovation and performance optimization. While standard frameworks like PyTorch make it easy to scale experimentation, pushing boundaries of performance requires optimizing the full stack (chip, server and UltraServer). Neuron provides ML performance engineers unparallelled access to our AWS AI chips, through Neuron Kernel Interface(NKI), deeper insights through Neuron Explorer, and our optimized kernel library called Neuron Kernel Library (NKILib). NKI provides APIs for memory allocation, and execution scheduling and direct access to the Trainium ISA enabling control over instruction-level programming. NKI Compiler is open sourced, built on MLIR, and provides developers visibility into the complete compiler pipeline. The open-source Neuron Kernel Library provides optimized implementations with source code, documentation, and benchmarks. Neuron Explorer provides a unified suite of tools that guides developers through their performance optimization and debugging journey. Performance engineers can trace execution from source code down to hardware operations, profile single-node and distributed applications, and receive AI-powered insights and actionable recommendations for kernel optimizations and performance improvements.

AI innovation thrives in open communities where developers can inspect, modify, and contribute. Neuron is committed to the open-source community and fostering innovation. While we move more parts of our stack to opensource, NKI Compiler, Neuron Kernel Driver, Neuron Kernel Library, NxD Inference, Neuron Explorer, and PyTorch, JAX and vLLM integrations are available today fully open-sourced. Open source libraries and tools enables developers to inspect compiler implementations, contribute optimizations, and adapt kernel code without any barriers. Come build with us.