This document is relevant for: Trn2, Trn3

Tutorial: Run a basic TorchTitan training on Neuron#

This tutorial walks you through a first, complete TorchTitan training run on AWS Trainium using TorchNeuron. When you have completed it, you will have a working end-to-end TorchTitan training example that trains Qwen3-8B on Trainium (on a single node and on multiple nodes).

Overview#

TorchTitan is PyTorch’s native platform for large-model pre-training. This tutorial runs it on Trainium using Qwen3-8B as the reference model: large enough to exercise tensor parallelism and FSDP, while staying small enough that the focus stays on the training example itself.

The training path below runs in eager mode — no torch.compile. Eager mode is the simplest way to confirm the example works end-to-end; once it is running, the torch.compile and Simple FSDP paths in Tutorial: Train Llama and DeepSeek with SimpleFSDP on Neuron are the next step.

Before you start#

This tutorial assumes that you have experience in the following areas:

  • Running a single-device PyTorch training loop on Neuron.

  • Basic familiarity with torchrun and distributed data-parallel training.

You do not need a prior TorchTitan installation — the setup doc linked in “Prepare your environment” below walks you through it.

Model details#

This tutorial trains Qwen3-8B with the following configuration:

  • Module: torchtitan_neuron.models.qwen3

  • Config: qwen3_8b_tp4_fsdp16_sac

  • Parallelism: tensor-parallel degree 4, data-parallel shard degree -1 (FSDP fills the remaining cores) — baked into the config

  • Execution: eager mode (no torch.compile), with selective activation checkpointing (SAC)

The qwen3_8b_tp4_fsdp16_sac config already sets the TP4 × FSDP16 parallelism layout, the batch size, and the sequence length, so you do not pass those as command-line flags. You select it by name through the CONFIG environment variable (see Step 2).

Limitations:

  • Validated on a single trn2.48xlarge node (64 NeuronCores, TP4 × FSDP16), and on two trn2.48xlarge nodes (128 NeuronCores, TP4 × FSDP32).

  • The batch size shipped in the config is the tested configuration; raising it may require more aggressive activation checkpointing to stay within device memory.

Prerequisites#

Prepare your environment#

Install and configure TorchTitan on your Trainium instance by following How to set up TorchTitan Neuron. Return here once import torchtitan succeeds and the setup doc’s verification run passes.

Step 1: Download the tokenizer#

From the torchtitan-neuron directory you installed in the setup step, download the Qwen3 tokenizer from Hugging Face:

python scripts/download_hf_assets.py --repo_id Qwen/Qwen3-8B --assets tokenizer

Tokenizers are saved to ./assets/hf/<model-name>/ by default. TorchTitan streams the C4 dataset from Hugging Face during training, so you do not download it separately. If the download fails with HTTP 429, see Common issues.

Step 2: Launch the training run#

TorchTitan ships a launcher script, run_train.sh, that wraps torchrun. You select the model and configuration through the MODULE and CONFIG environment variables and set the NeuronCore count with NGPU. MODULE is the model’s module path, torchtitan_neuron.models.qwen3:

# Train Qwen3-8B on a trn2.48xlarge instance (64 NeuronCores):

NGPU=64 \
MODULE=torchtitan_neuron.models.qwen3 \
CONFIG=qwen3_8b_tp4_fsdp16_sac \
./run_train.sh

run_train.sh passes any extra arguments straight through to torchtitan.train, so you can override a config value on the command line without editing the config — for example, to shorten the run:

NGPU=64 MODULE=torchtitan_neuron.models.qwen3 CONFIG=qwen3_8b_tp4_fsdp16_sac \
./run_train.sh --training.steps 10 --metrics.log_freq 1

Neuron environment variables of interest#

  • NGPU — the number of NeuronCores the run uses. run_train.sh passes this to torchrun as --nproc_per_node (see the GPU/Neuron contrast in Step 4). On a trn2.48xlarge this is 64.

  • MODULE / CONFIG — select the model (torchtitan_neuron.models.qwen3) and the registered training config (qwen3_8b_tp4_fsdp16_sac).

For the full set of TorchNeuron environment variables — compile-cache location, CPU-fallback controls, and debug directories — see TorchNeuron environment variables. The torch.compile backend and its Neuron-specific flags are covered in Tutorial: Train Llama and DeepSeek with SimpleFSDP on Neuron.

Step 3: Confirm the run succeeds#

Watch the training log for per-step loss lines. A healthy step looks like:

step:  3  loss: 11.04064  grad_norm:  5.3200  memory:  9.38GiB(26.07%)  tps: 587  tflops: 30.92  mfu: 19.57%

Confirm the run is healthy by checking that:

  • Training begins stepping and per-step loss lines appear.

  • Loss values are finite (not NaN) and trend downward over the first several steps.

Congratulations! You have now completed a first TorchTitan training run on AWS Trainium. If you encountered any issues, see Common issues.

Step 4 (optional): Scale to multiple nodes#

Neuron distributed launches differ from GPU in how the process count maps to hardware:

  • Single-node: NGPU=64 on a single trn2.48xlarge is 64 NeuronCores, arranged here as TP4 × FSDP16. On GPU the equivalent --nproc_per_node maps to the number of GPUs; on Neuron it maps to NeuronCores.

  • Multi-node: run_train.sh reads NNODES and RDZV_ENDPOINT and passes them to torchrun. A launcher, for example MPI, starts run_train.sh once on each node.

Multi-node with MPI#

Use mpirun to start run_train.sh once on each node. run_train.sh passes NNODES and RDZV_ENDPOINT to torchrun, and the ranks of all nodes meet at the rendezvous endpoint. The example trains Qwen3-8B on two trn2.48xlarge nodes (128 NeuronCores, TP4 × FSDP32).

Before you start, make sure that:

  • Each node has the same environment: torch-neuronx, the torchtitan-neuron directory at the same path, and the tokenizer from Step 1.

  • A hostfile lists each node with one slot, for example:

    node-0 slots=1
node-1 slots=1

  • The launcher can start a command on each node. This command prints the host name of each node:

    mpirun --allow-run-as-root -np 2 --hostfile <hostfile> hostname

  • Each node has EFA. This command prints the number of EFA endpoints on each node, which must be greater than 0:

    mpirun --allow-run-as-root -np 2 --hostfile <hostfile> \
  bash -c 'echo $(hostname) $(/opt/amazon/efa/bin/fi_info -p efa -t FI_EP_RDM | grep -c provider:)'

Run this command on the launcher:

mpirun --allow-run-as-root -np 2 --hostfile <hostfile> \
  -x PATH=<venv-path>/bin:$PATH -x CCOM_SOCKET_IFNAME=eth0 \
  -x NEURON_LOGICAL_NC_CONFIG=2 -x NEURON_RT_NUM_CORES=64 \
  bash -c 'cd /path/to/torchtitan-neuron && \
  NGPU=64 NNODES=2 RDZV_ENDPOINT=<node-0 address>:29500 \
  MODULE=torchtitan_neuron.models.qwen3 CONFIG=qwen3_8b_tp4_fsdp16_sac \
  ./run_train.sh'

  • -np and NNODES are the number of nodes. NGPU is the number of NeuronCores on each node.

  • RDZV_ENDPOINT is an address of the first node that all nodes can reach, plus a free port.

  • -x PATH=<venv-path>/bin:$PATH sets PATH on each node, so that each node finds torchrun in the Python environment of the setup step.

  • Use the full path to torchtitan-neuron. mpirun starts each node’s process in the launcher’s current directory.

  • CCOM_SOCKET_IFNAME sets the network interface for the Neuron collectives bootstrap. Set it to the interface of your nodes. It is necessary when the interface cannot be found automatically, for example in a container.

  • NEURON_LOGICAL_NC_CONFIG=2 sets the logical NeuronCore size (LNC 2), and NEURON_RT_NUM_CORES=64 sets the number of NeuronCores per node. NEURON_LOGICAL_NC_CONFIG is valid on both Trn2 and Trn3.

The config sets the data-parallel shard degree to -1, so FSDP fills the remaining NeuronCores. For four nodes, change -np and NNODES to 4. The command does not change in other ways.

Note

  • Each rank streams the C4 dataset from Hugging Face. With many ranks, an anonymous download can fail with HTTP 429 (rate limit). If you see this error, set HF_TOKEN on each node, for example with -x HF_TOKEN.

Common issues#

  • TorchTitan cannot find the Neuron device: Confirm torch-neuronx is installed and device='neuron' resolves in a plain PyTorch script first. See Quickstart: Install and verify torch-neuronx.

  • OOM at this batch size: Enable activation checkpointing before lowering the batch size. See Best practices for TorchTitan Neuron.

  • Dataset download fails with HTTP 429: Hugging Face limits anonymous downloads. Set HF_TOKEN to your Hugging Face access token.

  • Multi-node run fails at process group initialization with NET/OFI errors: Check EFA on each node with the command in Multi-node with MPI. If EFA is present on each node, contact your cluster administrator about the EFA connection between the nodes.

Clean up#

Stop the training processes with Ctrl-C on the launcher: torchrun for a single-node run, or the mpirun process for a multi-node run — mpirun forwards the interrupt to the per-node processes and tears the job down on every node. Then, if you no longer need them, remove the checkpoint output directory TorchTitan wrote under its output path.

Note

Running on Trn3: The instance-dependent values on this page, such as the number of ranks per node, the device memory (HBM) per rank, NEURON_RT_NUM_CORES, and the device mesh shapes, are for Trn2 (trn2.48xlarge). Trn3 servers have a different number of Trainium chips, so these values do not apply to Trn3 without change. Before you run on Trn3, calculate them again from the Amazon EC2 Trn3 architecture and Trainium3 architecture specifications. For the effect of the LNC setting on the number of visible cores, see Logical NeuronCore configuration.

Next steps#

Further reading#

This document is relevant for: Trn2, Trn3