This document is relevant for: Inf2
, Trn1
, Trn2
NxD Core#
NeuronX Distributed Core (NxD Core) is a package for supporting different distributed training/inference mechanism for Neuron devices. It would provide xla friendly implementations of some of the more popular distributed training/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 Parallel sharding technique with other distributed library supported to be added in future.
Setup
Install PyTorch Neuron on Trn1 to create a pytorch environment. It is recommended to work out of python virtual env 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
API Reference Guide
Developer Guide
Tutorials
This document is relevant for: Inf2
, Trn1
, Trn2