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

NxD Core for Training

NeuronX Distributed Core (NxD Core) is a package for supporting different distributed training mechanisms for Neuron devices. It provides XLA-friendly implementations of some of the more popular distributed training 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.

About NeuronX-Distributed (NxD) for Training

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.

The NeuronX Distributed Training (NxD Training) library is a collection of open-source tools and libraries designed to empower customers to train PyTorch models on AWS Trainium instances. It combines both ease-of-use and access to features built on top of NxD Core library. Except for a few Trainium-specific features, NxD Training is compatible with training platforms like NVIDIA’s NeMo.

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

• Developer guide for Tensor Parallelism

• Developer guide for Pipeline Parallelism

• Developer guide for Activation Memory reduction

• Developer guide for save/load checkpoint

• Developer guide for Neuron-PT-Lightning

• Developer guide for model and optimizer wrapper

Tutorials

• Training Tutorials

• Inference Tutorials

Misc

• NxD Core Release Notes (neuronx-distributed)

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