AWS Neuron SDK Documentation

Neuron Deep Learning AMIs (DLAMIs) are pre-configured Amazon Machine Images designed for running machine learning workloads on AWS Inferentia and Trainium instances. They come with Neuron drivers, frameworks, and libraries pre-installed to eliminate manual setup.

Supported Instance Types:

• Inf2, Trn1, Trn1n, Trn2, Trn3

Important Compatibility Notes:

• Inf1 Support: Starting with version 2.26.1, DLAMIs no longer support Inf1 instance types. For Inf1 workloads, use DLAMIs released up to SDK version 2.26.
• Ubuntu 22.04: Ubuntu 22.04 has reached end-of-support. New deployments should use Ubuntu 24.04.

NxD Training (NeuronX Distributed Training) is a PyTorch library designed for end-to-end distributed training on AWS Trainium instances. It provides turnkey workflows for model pre-training, supervised fine-tuning (SFT), and parameter-efficient fine-tuning (PEFT) using LoRA.

Key capabilities include:

• Distributed Strategies: Support for Data Parallelism, Tensor Parallelism, Sequence Parallelism, Pipeline Parallelism, and ZeRO-1.
• Integration: Uses PyTorch Lightning for organized training code and supports HuggingFace and Megatron-LM model formats.
• Experiment Management: Integrated checkpointing, logging, and S3 storage support.
• Compatibility: Compatible with training platforms like NVIDIA's NeMo (excluding Trainium-specific features).

NxD Inference (NeuronX Distributed Inference) is an ML inference library included with the AWS Neuron SDK designed to simplify the deployment of deep learning models on AWS Inferentia and Trainium instances.

Key capabilities include:

• High-Performance Inference: Supports continuous batching and speculative decoding.
• Model Distribution: Distributes large models across multiple NeuronCores using Tensor Parallelism and Sequence Parallelism.
• Optimizations: Leverages KV Cache, Flash Attention, and Quantization.
• vLLM Integration: Provides seamless integration with vLLM for production-ready deployments.
• Supported Models: Optimized for models such as Llama-3.1, DBRX, and Mixtral.

Neuron Explorer is the recommended profiling tool for AWS Neuron workloads on AWS Trainium. It provides end-to-end profiling support, allowing ML engineers to trace execution from source code down to hardware operations. It supports both single-node and distributed applications and can be accessed via CLI, a web UI, or through a Visual Studio Code extension.

Note on Migration: Neuron Explorer replaces the older Neuron Profiler and Neuron Profiler 2.0. Support for the older profilers ended in the Neuron 2.29 release. Neuron Explorer is backwards compatible with profile data captured using the older profilers, though existing files must be reprocessed before viewing.

vLLM on Neuron enables high-performance Large Language Model (LLM) inference on AWS Trainium and Inferentia instances. It integrates vLLM's advanced features with AWS Neuron's optimized AI inference capabilities via the NxD Inference (neuronx-distributed-inference) library.

Key Capabilities:

• Continuous batching: Efficiently processes multiple requests simultaneously.
• Prefix caching: Reuses KV cache from common prompts to improve time-to-first-token.
• Speculative decoding: Supports Eagle V1.
• Quantization: Supports INT8 and FP8 for optimized performance.
• Multimodal support: Supports models like Llama 4 Scout and Maverick.

Supported Models include:

• Llama 2, 3.1, 3.3
• Llama 4 Scout, Maverick
• Qwen 2.5, Qwen 3
• Custom models onboarded to NxD Inference.

AWS Neuron Reference for NeMo Megatron is a library providing optimized versions of the open-source NeMo and Apex packages specifically for AWS Neuron and AWS EC2 Trn1 instances.

Key features include:

• Distributed Training Support: Supports Tensor Parallel, Pipeline parallel, and Data Parallel configurations for large language models (e.g., GPT-3 175B).
• Optimized Performance: APIs are optimized for XLA-based computation and high-performance communication on Trainium instances.
• Memory Optimization: Utilizes techniques like sequence parallelism (to reduce activation memory footprint), selective or full activation checkpointing (to fit larger models), and SPMD optimizations (to reduce the number of graphs obtained).

JAX NeuronX provides high-performance deep learning acceleration for AWS Trainium (Trn1, Trn2) and AWS Inferentia (Inf2) instances. It integrates these accelerators into JAX as pluggable devices using the PJRT (Plugin Runtime) mechanism, allowing for native JAX device support with minimal code changes.

Key components include:

• libneuronxla: The core integration into JAX's runtime PJRT via the PJRT C-API plugin mechanism. Installing this enables Trainium and Inferentia to be used natively as JAX devices.
• jax-neuronx: A meta-package that provides a tested combination of jax-neuronx, jax, jaxlib, libneuronxla, and neuronx-cc. It also includes Neuron-specific features like the Neuron NKI JAX interface.

Note: JAX NeuronX is currently in beta. Some JAX functionality may not be fully supported.

The NeuronX Runtime Library provides C APIs for interacting with Neuron hardware. It allows developers to:

• Initialize Neuron hardware.
• Load models and input data.
• Execute iterations on loaded models.
• Retrieve output data.

The core library is provided as a shared object (libnrt.so).

NxD Inference (via neuronx-distributed-inference) allows you to use AWS Inferentia and AWS Trainium AI accelerators with vLLM v0.x.

NxD Inference integrates into vLLM by extending the model execution components responsible for loading and invoking models in vLLM's LLMEngine. This means input processing, scheduling, and output processing follow standard vLLM behavior.

Current Capabilities:

• Continuous batching
• Streaming generation

Note on Versions: This guide is for vLLM v0.x (Legacy). For new deployments, Neuron recommends using vLLM v1.x with the vLLM-Neuron Plugin.

TorchNeuron is an open-source PyTorch backend designed for AWS Trainium. It provides native integration for the PyTorch framework, allowing developers to run models on Trainium hardware using standard PyTorch workflows.

Key capabilities include:

• Eager Mode: Supports rapid iteration and experimentation with optimized ATen operators.
• torch.compile Support: Enables JIT compilation via a custom TorchDynamo backend for improved performance.
• Distributed Support: Integrates with torch.distributed for scaling models across hardware.
• Custom Kernels: Integrates with the Neuron Kernel Interface (NKI) for fine-grained performance optimization.

Note: TorchNeuron is currently only available as part of a closed Beta program. Contact your AWS Neuron support representative to participate.