This document is relevant for: Trn2, Trn3

What is TorchTitan Neuron?#

This topic covers TorchTitan Neuron and how it applies to pre-training with the AWS Neuron SDK. TorchTitan is a PyTorch-native library for pre-training large models. On Trainium it runs on TorchNeuron, the native PyTorch backend for Neuron, so you can use TorchTitan’s standard model definitions, parallelism strategies, and torch.compile integration without changing your model code.

TorchTitan expresses every parallelism strategy with standard PyTorch building blocks: DeviceMesh, DTensor, fully_shard (FSDP2), and torch.distributed collectives. TorchNeuron registers a neuron device and a neuron process-group backend that implement these APIs. When you add --compile.enable --compile.backend neuron, TorchTitan’s compiled regions are lowered by the Neuron compiler. Large-model pre-training on Trainium therefore stays inside the PyTorch-native stack, with no separate framework or graph-tracing layer to adopt.

Applies to#

This concept is applicable to:

  • Pre-training workloads on Trn2 and Trn3 instances using TorchTitan.

  • Distributed training using TorchTitan’s built-in parallelism strategies on top of TorchNeuron’s torch.distributed support.

Support matrix#

Component

Supported version or value

Neuron SDK

2.33.0 (TorchNeuron public beta)

PyTorch

2.14.0

torch-neuronx

2.14.3.0 package family

Python

3.12, 3.13, 3.14 (Linux x86_64)

TorchTitan

Use aws-neuron/torchtitan-neuron, the Neuron fork of upstream pytorch/torchtitan commit 7cdf0aa9

Instance types

Trn2 and Trn3. The TorchTitan tutorials use a single trn2.48xlarge with 64 logical NeuronCores (one process per NeuronCore).

Execution modes

Eager, and torch.compile with --compile.backend neuron (static shapes only)
Table 1 Parallelism support#

Strategy

Status

Notes

FSDP2 (fully_shard)

Supported

Default TorchTitan data-parallel sharding (--parallelism.data-parallel-shard-degree). Uses Neuron all_gather and reduce_scatter.

SimpleFSDP

Supported with jit backend

Compiler-friendly FSDP for torch.compile. Available for Llama3 and DeepSeek-V3.

Tensor parallel (TP)

Supported

DTensor ColwiseParallel, RowwiseParallel, and SequenceParallel. The tutorials use TP=4.

Context parallel (CP)

Supported

Ring attention for sequence-sharded SDPA, forward and backward.

Expert parallel (EP)

Supported

Uneven all_to_all_single token dispatch through a native NKI all-to-all-v collective.

Pipeline parallel (PP)

Supported by TorchNeuron

The Neuron backend implements point-to-point send and recv. Not covered by the TorchTitan tutorials.

DDP

Supported

Standard DistributedDataParallel.

Activation checkpointing

Supported

--activation-checkpoint.mode full and selective, including the per-op policy.

How TorchTitan works on Neuron#

TorchTitan is a layer on top of the native PyTorch backend. It does not use PyTorch/XLA, and it doesn’t trace your model into a separate graph framework. Model code, parallelism plans, the optimizer, and checkpointing are ordinary PyTorch running on the neuron device.

┌──────────────────────────────────────────────────────────────┐
│ TorchTitan                                                   │
│   models (Llama3, Qwen3, DeepSeek-V3, GPT-OSS, ...)          │
│   parallelize: FSDP2 / SimpleFSDP, TP, CP, EP, AC            │
│   train loop, optimizer, checkpoint, metrics                 │
├──────────────────────────────────────────────────────────────┤
│ PyTorch 2.14                                                 │
│   DeviceMesh · DTensor · fully_shard · torch.distributed     │
│   torch.compile (Dynamo + AOTAutograd)                       │
├──────────────────────────────────────────────────────────────┤
│ TorchNeuron (torch-neuronx)                                  │
│   "neuron" device (PrivateUse1)   "neuron" process group     │
│   eager ATen dispatch             DTensor sharding rules     │
│   torch.compile backend="neuron"  NKI kernels (SDPA, ...)    │
├──────────────────────────────────────────────────────────────┤
│ Neuron compiler (neuronx-cc) · Neuron runtime · collectives  │
├──────────────────────────────────────────────────────────────┤
│ Trainium (Trn2 / Trn3) NeuronCores                           │
└──────────────────────────────────────────────────────────────┘

A training run goes through four layers:

  • Process and device setup. torchrun starts one process per NeuronCore. Set NEURON_RT_NUM_CORES and --nproc_per_node to the same value. TorchNeuron registers the neuron backend for init_process_group and assigns each rank its NeuronCore from LOCAL_RANK.

  • Parallelism. TorchTitan builds a DeviceMesh and applies its parallelism plan with DTensor and fully_shard. TorchNeuron records every mesh’s replica groups when the mesh is created, which SPMD compilation depends on. It also supplies sharding rules for attention and linear layers so that DTensor operations stay sharded on the device.

  • Execution. Without --compile.enable, every operation runs in eager mode. Each ATen op is compiled to a Neuron executable on first use and then cached. With --compile.enable --compile.backend neuron, TorchTitan wraps part of the model in torch.compile, and those regions are compiled by the Neuron compiler instead:

    • Stock models with FSDP2: each TransformerBlock is compiled separately with fullgraph=True. FSDP2’s parameter all-gather and gradient reduce-scatter run as eager collectives between the compiled blocks.

    • SimpleFSDP: the whole model is compiled as one fullgraph=True graph. The FSDP collectives are inside the graph, so the compiler can schedule communication together with compute. This is why SimpleFSDP is the recommended FSDP choice with torch.compile.

  • Caching. Compiled executables are stored in a persistent cache keyed by the workload, the target, the logical NeuronCore (LNC) configuration, the compiler version, and the compiler options. The first step of a run can take several minutes to compile. Later steps, and later runs with the same configuration, reuse the cache.

Supported models#

TorchTitan’s model families run on TorchNeuron through the same model definitions used on other hardware. The following configurations have TorchNeuron tutorials:

Model

Workload

Validated configuration

Llama3-8B

Pre-training

TP4 × FSDP16 on one trn2.48xlarge. Eager, torch.compile, and SimpleFSDP + torch.compile.

Qwen3-32B (dense)

Pre-training

TP4 × FSDP16, full activation checkpointing. Eager and torch.compile.

DeepSeek-V3

Pre-training

Mixture of experts (MoE) with multi-head latent attention. SimpleFSDP + torch.compile.

GPT-OSS

Pre-training

MoE with grouped experts.

Qwen3-8B

Inference

Hugging Face transformers model with DTensor tensor parallelism and torch.compile(backend="neuron").

Other TorchTitan models (for example Llama4 and the Qwen3 MoE variants) use the same building blocks. They are expected to run, but they don’t have a validated TorchNeuron configuration yet.

Distributed strategies#

TorchTitan composes its parallelism strategies on a multi-dimensional DeviceMesh. Each rank of the mesh is one process, and each process owns one logical NeuronCore. On a trn2.48xlarge with the default LNC=2, that is 64 ranks per node: 16 Trainium2 devices with 4 logical cores each.

FSDP2 and SimpleFSDP. Both shard parameters, gradients, and optimizer state across the data-parallel shard dimension. FSDP2 (fully_shard) uses runtime hooks around each module and works in eager mode and with per-block compilation. SimpleFSDP expresses sharding as DTensor parametrizations, so the all-gather and reduce-scatter are traced into the compiled graph. Use SimpleFSDP when you combine FSDP with torch.compile.

Tensor parallelism (DTensor). TP shards attention heads and MLP projections. TorchNeuron provides DTensor sharding strategies for scaled dot-product attention (replicated, head-sharded, batch-sharded, and sequence-sharded) and for linear forward and backward, so a TP plan stays sharded through attention.

Context parallelism. Sequence-sharded attention runs as ring attention. By default, TorchNeuron uses a fused NKI ring-attention kernel that drives the whole ring inside one kernel, and falls back to PyTorch’s templated ring otherwise.

Expert parallelism. MoE token dispatch uses all_to_all_single with uneven split sizes. TorchNeuron implements this as a native NKI all-to-all-v collective that sends exactly the requested per-rank counts without dynamic shapes. Set NEURON_USE_NKI_ALLTOALLV=0 to use the padding-based fallback.

Mapping onto NeuronCores. Collective cost depends on where the ranks of a group sit. Ranks on the same Trainium device communicate fastest, then ranks on the same node, then ranks across nodes. Put the most communication-heavy dimension on the innermost mesh dimension:

  • TP communicates in every layer, so keep each TP group on one device. For example, TP=4 matches the 4 logical NeuronCores of a Trainium2 device at LNC=2.

  • FSDP communicates once per layer or per step and can span devices and nodes. --parallelism.data-parallel-shard-degree -1 uses all remaining ranks.

trn2.48xlarge, 64 ranks, TP=4 × FSDP=16

device 0             device 1                   device 15
┌─────────────────┐  ┌─────────────────┐        ┌─────────────────┐
│ r0  r1  r2  r3  │  │ r4  r5  r6  r7  │  ...   │ r60 r61 r62 r63 │
└─────────────────┘  └─────────────────┘        └─────────────────┘
    TP group 0           TP group 1                 TP group 15
└────────────── FSDP shards across the 16 TP groups ──────────────┘

The LNC setting changes the number of ranks per node and is part of the compilation cache key, so changing it triggers a recompile.

Known limitations#

The following limitations are specific to running TorchTitan on Neuron. For general TorchNeuron constraints (operator coverage, data types, and torch.compile parameters), see the TorchNeuron API compatibility guide.

  • Static shapes only. torch.compile on Neuron requires static shapes. Keep --training.seq-len and --training.local-batch-size fixed for a run. Each new input shape triggers a new compilation.

  • Long first step. The first training step compiles every graph and can take several minutes. Throughput is steady after that. The persistent NEFF cache avoids paying this cost again on later runs.

  • Compile mode. Only mode="default" is supported. TorchTitan’s default is compatible.

  • Forced graph breaks under fullgraph=True. TorchTitan compiles with fullgraph=True. TorchNeuron forces a graph break on operations Neuron can’t lower in a compiled graph, and under fullgraph=True each of those becomes an error. This affects:

    • torch._grouped_mm / grouped_mm, used by MoE grouped experts. It always forces a graph break; it runs on Neuron in eager mode.

    • sort / argsort on a dimension longer than 16,384, and topk with k ≥ 16,384.

    For MoE models, keep the expert computation out of the compiled region, or run without --compile.enable.

  • 64-bit types. float64 and int64 are computed as 32-bit types on the device. Keep training in bfloat16 / float32.

  • Pinned TorchTitan revision. The tutorials are validated against aws-neuron/torchtitan-neuron, the Neuron fork of upstream pytorch/torchtitan commit 7cdf0aa9.

  • Profiling. Neuron device profiling through TorchTitan’s built-in profiler needs support for the PrivateUse1 profiler activity, which isn’t in the pinned revision. The Llama3 tutorial passes --profiling.no-enable-profiling. To profile, use torch.profiler with torch_neuronx.profiling.NeuronConfig and Neuron Explorer.

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.

Further reading#

This document is relevant for: Trn2, Trn3