TorchTitan Neuron is the AWS Neuron's extension of pytorch/torchtitan. It adds AWS Trainium support through TorchNeuron, the native PyTorch backend for Neuron. The Neuron-specific recipes, NKI kernel integrations, and docs are in torchtitan_neuron/.
| Component | Supported version |
|---|---|
| Neuron SDK | 2.33.0 (TorchNeuron public beta) |
PyTorch / torch-neuronx |
2.14.0 / 2.14.3.0.* |
| Python | 3.12, 3.13, 3.14 (Linux x86_64) |
| Instances | Trn2, Trn3 |
| Execution | Eager, and torch.compile with --compile.backend neuron (static shapes only) |
| Models | Llama3, Qwen3, DeepSeek-V3, GPT-OSS |
| Parallelism | FSDP2, SimpleFSDP, TP, CP, EP, PP, DDP, activation checkpointing |
Before you start, install and verify torch-neuronx on a Trn2 or Trn3 instance by following Quickstart: Install and verify torch-neuronx. Then install this fork:
git clone https://github.com/aws-neuron/torchtitan-neuron.git
cd torchtitan-neuron
git checkout release_2.33.0
pip install -r requirements.txt
pip install -e .Run a short verification on one trn2.48xlarge (64 logical NeuronCores, TP4 × FSDP16) using the bundled test tokenizer and dataset:
NGPU=64 MODULE=torchtitan_neuron.models.qwen3 CONFIG=qwen3_8b_tp4_fsdp16_sac \
./run_train.sh --training.steps 10 --metrics.log_freq 1 \
--dataloader.dataset c4_test --hf_assets_path=./tests/assets/tokenizerStart with the Neuron package overview for the package boundary, layout, included components, and recipe selection.
| Topic | Guide |
|---|---|
| Dense Qwen3 and Llama3 training | Dense model tutorials |
| Qwen3, DeepSeek-V3, and GPT-OSS MoE training | MoE model tutorials |
| Qwen3 NKI mixture-of-experts recipes | Qwen3 NKI MoE |
| Full Qwen3-235B-A22B training | Four-node Trainium3 UltraServer guide |
Qwen3 inference with torch.compile |
Qwen3 inference |
| Neuron extension test suite | Testing torchtitan_neuron |
Report problems with running on Neuron or Trainium in this repository's issues. Report problems that also occur on upstream torchtitan to pytorch/torchtitan.
The Neuron-specific recipes, NKI kernel integrations, and docs are in
torchtitan_neuron/:
torchtitan_neuron/
|-- models/ Neuron modules and Trainer configuration registries
|-- experiments/ GraphTrainer configuration registries
|-- components/ Neuron-specific reusable components
|-- docs/ AWS Neuron documentation
`-- tests/ Neuron kernel, Trainer, and GraphTrainer tests
We provide a detailed look into the parallelisms and optimizations available in torchtitan, along with summary advice on when to use various techniques.
TorchTitan: One-stop PyTorch native solution for production ready LLM pre-training
@inproceedings{
liang2025torchtitan,
title={TorchTitan: One-stop PyTorch native solution for production ready {LLM} pretraining},
author={Wanchao Liang and Tianyu Liu and Less Wright and Will Constable and Andrew Gu and Chien-Chin Huang and Iris Zhang and Wei Feng and Howard Huang and Junjie Wang and Sanket Purandare and Gokul Nadathur and Stratos Idreos},
booktitle={The Thirteenth International Conference on Learning Representations},
year={2025},
url={https://openreview.net/forum?id=SFN6Wm7YBI}
}
Source code is made available under a BSD 3 license, however you may have other legal obligations that govern your use of other content linked in this repository, such as the license or terms of service for third-party data and models.