AWS NxD libraries

Accelerate your machine learning workloads with Neuron’s distributed libraries. Our libraries provide high-level abstractions and optimized implementations for distributed training and inference on AWS Trainium and Inferentia.

What are NeuronX Distributed libraries?

NeuronX Distributed (NxD) libraries are a comprehensive suite of PyTorch-based libraries designed to enable scalable machine learning on AWS Neuron hardware. The NxD ecosystem provides a layered architecture where foundational distributed primitives support higher-level training and inference workflows.

The NxD Stack:

• NxD Core: The foundational layer providing distributed primitives, model sharding techniques, and XLA-optimized implementations

• NxD Training: High-level training library built on NxD Core, offering turnkey distributed training workflows with NeMo compatibility

• NxD Inference: Production-ready inference library with advanced features like continuous batching, speculative decoding, and vLLM integration

Together, these libraries enable developers to scale from prototype to production while leveraging the full performance potential of AWS Trainium and Inferentia instances.

NxD Core

Core distributed training and inference mechanisms for Neuron devices with XLA-friendly implementations.

• Tensor Parallel (TP) sharding

• Pipeline Parallel (PP) support

• Model partitioning across devices

• XLA-optimized distributed operations

• Foundation for other NxD libraries

NxD Inference

PyTorch-based inference library for deploying large models on Inferentia and Trainium.

• Large Language Model (LLM) inference

• Disaggregated inference architecture

• vLLM integration and compatibility

• Model sharding and parallelism

• Performance optimization tools

NxD Training

PyTorch library for distributed training with NeMo-compatible YAML interface.

• Large-scale model training

• NeMo YAML configuration support

• HuggingFace and Megatron-LM models

• Experiment management

• Advanced parallelism strategies

Transformers NeuronX

Legacy Library (Archived)

Original transformer inference library - now superseded by NxD Inference.

• Status: Support ended 9/16/2025

• Migration path to NxD Inference

• Archived documentation available

• Legacy workload support

• Recommended: Migrate to NxD Inference

Hardware compatibility

Library	Inf1	Inf2	Trn1/Trn1n	Trn2	Inference	Training
NxD Core	❌	✅	✅	✅	✅	✅
NxD Inference	❌	✅	✅	✅	✅	❌
NxD Training	❌	❌	✅	✅	❌	✅

Third-Party partner libraries

🤗 Hugging Face Optimum Neuron

Standard Hugging Face APIs for Trainium and Inferentia with SageMaker support.

⚡ PyTorch Lightning

Professional AI framework with maximal flexibility and NxD integration.

🔬 AXLearn

JAX-based library for distributed training with AWS Trainium integration.