This document is relevant for: Trn1, Trn2

Developer Guide

This section will go over a variety of developer guides to help users get started with the Neuronx Distributed Training library.

• Integrating a new model
  • Model Building (torch.nn.Module)
  • Model Integration
• Integrating a new dataset/dataloader
  • Building Dataset module
  • Building DataModule
• Registering an optimizer and LR scheduler
  • Setting up the optimizer
  • Setting up the LR scheduler
• Migrating from Neuron-NeMo-Megatron to Neuronx Distributed Training
  • Config migration
  • Model code
  • Checkpointing Save/Load
  • Config Mapping
• NxD Training Compatibility with NeMo
  • Model Integration
  • Dataloader Integration
  • Optimizer/LR Scheduler Integration
  • Optimal Partitioning
  • Fusions/kernels
  • Checkpoint Saving/loading
  • Config Mapping

This document is relevant for: Trn1, Trn2