This document is relevant for: Trn1, Trn2, Trn3

Train models on Neuron#

Train and fine-tune models on AWS Trainium. This page organizes training by path — native PyTorch (TorchNeuron), TorchTitan Neuron, and Hugging Face — so you can start from your task instead of choosing a library first.

On Neuron, you train and fine-tune models in PyTorch. You can serve supported LLM checkpoints with vLLM Neuron. To serve a trained model, see Inference on AWS Neuron.

Choose your training path#

Neuron supports three training paths on Trainium:

  • Native PyTorch (TorchNeuron) — the recommended path for new work. TorchNeuron is the Neuron native PyTorch backend; it supports eager mode, torch.compile with the neuron backend, and standard distributed APIs (FSDP, DTensor, DDP, functional collectives). Most PyTorch code runs with minimal changes, primarily moving device placement from cuda to neuron.

  • TorchTitan Neuron — for large-scale pre-training built on TorchNeuron. Use TorchTitan when you want a distributed training stack with recipes for models such as Llama, DeepSeek (FSDP), and GPT-OSS.

  • Hugging Face libraries — for fine-tuning and training with Hugging Face transformers, diffusers, PEFT, and accelerate, which run natively on TorchNeuron with no Neuron-specific path.

Learn the concepts#

overview

About TorchNeuron

Understand TorchNeuron’s execution modes and how it maps to the PyTorch ecosystem.

overview

TorchNeuron landing

The native PyTorch section home: quickstarts, tutorials, guides, and reference.

deep dive

torch.compile on Neuron

How JIT compilation with backend="neuron" works, including the compilation pipeline and caching.

overview

TorchTitan Neuron

Large-scale pre-training on Trainium with TorchNeuron and TorchTitan.

overview

Training paths (native and legacy XLA)

Compare the native PyTorch path with the legacy PyTorch/XLA (torch-neuronx XLA) path.

Start training#

quickstart

Get started with TorchNeuron

Launch a Trainium instance and run your first TorchNeuron training check.

tutorial

TorchNeuron tutorials

End-to-end walk-throughs: a simple training loop and distributed training with torch.distributed.

how-to

Distributed training how-to

Use FSDP, DTensor, DDP, and functional collectives on Neuron.

tutorial

TorchTitan tutorials

Pre-train Llama, DeepSeek (FSDP), and GPT-OSS on Trainium with TorchTitan.

how-to

Training best practices

Guidance for training performance on Trainium.

Where NKI fits in training

Custom kernels written with the Neuron Kernel Interface (NKI) plug into the training path as PyTorch custom operators — called like any other operator in eager mode and under torch.compile. Reach for NKI when a performance-critical path needs finer control than the compiler provides; start from the NKI Library of pre-optimized kernels.