NxD Core for Inference

Note

Maintenance status

The inference-context NxD Core documentation in this section is in maintenance status — it continues to receive critical fixes and remains the supported path for Trn1/Inf2 inference and for custom NxD Inference API serving (for example Flux/diffusion). This content is not end-of-support. For new inference work, the go-forward path is the vLLM Neuron plugin; the task-oriented Inference on Neuron page shows where each path fits.

NeuronX Distributed Core (NxD Core) is a package for supporting different distributed inference mechanisms for Neuron devices. It provides XLA-friendly implementations of some of the more popular distributed inference techniques. As the size of the model scales, fitting these models on a single device becomes impossible and hence we have to make use of model sharding techniques to partition the model across multiple devices.

As part of this library, we enable support for Tensor Parallelism sharding technique with other distributed library supported to be added in future.

About NeuronX-Distributed (NxD) Inference

NeuronX Distributed (NxD Core) provides fundamental building blocks that enable you to run advanced inference workloads on AWS Inferentia and Trainium instances. These building blocks include parallel linear layers that enable distributed inference, a model builder that compiles PyTorch modules into Neuron models, and more.

As part of NxD Core, Neuron offers NxD Inference, which is a library that provides optimized model and module implementations that build on top of NxD Core. For more information about NxD Inference, see NxD Inference Overview.

For examples of how to build directly on NxD Core, see the following:

• Llama 3.2 1B inference sample

• T5 3B inference tutorial [html] [notebook]

NxD Core for Inference Documentation

Setup

Install PyTorch Neuron on Trn1 to create a pytorch environment. It is recommended to work out of a Python virtual environment (such as venv) so as to avoid package installation issues.

You can install the neuronx-distributed package using the following command:

python -m pip install neuronx_distributed --extra-index-url https://pip.repos.neuron.amazonaws.com

App Notes

• Tensor Parallelism Overview

• Pipeline Parallelism Overview

• Activation Memory Reduction

• Context Parallelism Overview

API Reference Guide

• Distributed Strategies APIs

• Training APIs

• Inference APIs

• ModelBuilderV2 API Reference

Developer Guide

• About NeuronX-Distributed (NxD) Inference

Tutorials

• Training Tutorials

• Inference Tutorials

Misc

• Component Release Notes for NxD Core