NxD Training (archived)

Warning

This section is archived. NxD Training (NeuronX Distributed Training) and the XLA-based torch-neuronx training flow have reached end-of-support as of Neuron 2.32 and are no longer actively maintained. The content below is preserved for historical reference and may describe features, APIs, or workflows that no longer work with current Neuron releases. For current PyTorch training on Trainium, see PyTorch Support on Neuron.

This page consolidates the archived documentation for NxD Training and the XLA-based torch-neuronx training experience — the library overview, setup and configuration, tutorials, developer guides, migration guides, application notes, and reference material that previously lived under libraries/nxd-training/ and frameworks/torch/torch-neuronx/ (training). Individual archived pages are listed in the sections below and in the navigation tree for this page.

Overview

NxD Training (NeuronX Distributed Training) was a PyTorch library for end-to-end distributed training on AWS Trainium instances, built on top of the NxD Core library. It offered turnkey support for model pre-training, supervised fine-tuning (SFT), and parameter-efficient fine-tuning (PEFT) with LoRA, and was compatible with NVIDIA’s NeMo (except for Trainium-specific features).

Key capabilities included:

• Turnkey workflows for pre-training, SFT, and PEFT (LoRA).

• Distributed strategies [1]: data parallelism, tensor parallelism, sequence parallelism, pipeline parallelism, and ZeRO-1.

• PyTorch Lightning integration for organized training code.

• Ready-to-use model samples in HuggingFace and Megatron-LM formats.

• Experiment management with checkpointing, logging, and S3 storage support.

NxD Training exposed three usage interfaces, letting developers work at the level of abstraction that suited their needs:

• YAML configuration files — high-level distributed training with minimal code changes.

• PyTorch Lightning APIs — standardized training workflows over NxD Core primitives.

• NxD Core primitives — low-level APIs for custom model integration and advanced use.

Fig. 90 NxD Training

Fig. 91 Using NxD Training through (a) configuration files, (b) PyTorch Lightning APIs, and (c) NxD Core primitives.

Setup and configuration

Installation and configuration guidance for the archived NxD Training library:

• Installation guide

• YAML configuration settings

• Features

Tutorials

End-to-end training tutorials for the XLA torch-neuronx and NxD Training flows. These tutorials are archived and may not run on current Neuron releases.

torch-neuronx training tutorials:

• BERT pretraining

• Multi-layer perceptron (MLP) training

• HuggingFace BERT fine-tuning

• ZeRO-1 GPT-2 pretraining

• Analyze a model for training

NxD Training library tutorials:

• Llama3 8B pretraining

• Llama3 8B supervised fine-tuning (SFT)

• Llama3 8B SFT with LoRA

• Llama3 8B DPO/ORPO alignment

• Llama3 70B pretraining

• Checkpoint conversion

Developer guides

Guides for extending and customizing NxD Training — integrating new models and dataloaders, customizing the optimizer/LR-scheduler flow, and running in CPU mode:

• Integrate a new model

• Integrate a new dataloader

• Register the optimizer and LR scheduler

• CPU-mode developer guide

Migration guides

For teams moving from earlier training stacks to NxD Training:

• Migrate from NeuronX NeMo Megatron (NNM)

• Migrate from NVIDIA NeMo

Application notes

Deep dives on the distributed strategies and memory-optimization techniques NxD Training used:

• Tensor parallelism

• Pipeline parallelism

• Context parallelism

• Activation memory reduction

Reference and troubleshooting

Reference material and troubleshooting for the XLA torch-neuronx training flow:

• PyTorch NeuronX programming guide

• Debugging guide

• Environment variables

• Training troubleshooting

• Known issues

Additional examples

Reference sample repositories on GitHub (external; may also be unmaintained):

• AWS Neuron Reference for NeMo Megatron

• AWS Neuron Samples for EKS

• AWS Neuron Samples for AWS ParallelCluster

• AWS Neuron Samples (torch-neuronx training)

[1]

Distributed strategies are implemented in the NxD Core library.