This document is relevant for: Inf2, Trn1, Trn2, Trn3

NxD Core for Inference#

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]