This document is relevant for: Trn2, Trn3

Best practices for TorchNeuron#

This topic covers best practices for running PyTorch workloads efficiently on AWS Trainium with TorchNeuron. The practices apply to eager mode, torch.compile, distributed training, and custom kernels on Neuron hardware.

Overview#

These practices help you maximize throughput, reduce compilation time, and avoid common performance problems when running PyTorch on Trainium. They apply to both training and inference workloads on TorchNeuron’s native PyTorch backend.

This guidance targets the TorchNeuron public beta in Neuron SDK 2.33.0: PyTorch 2.14.0, torch-neuronx 2.14.3.0, Python 3.10–3.14, on Trainium2 and Trainium3.

TorchNeuron registers a neuron device with PyTorch through the PrivateUse1 mechanism. Supported ATen ops on a neuron tensor are lowered to StableHLO, compiled to a NEFF by neuronx-cc, and cached. Unsupported ops fall back to CPU by default, which is convenient for bring-up but costs performance (see practice 1). torch.compile(backend="neuron") follows the same path, but compiles a whole FX graph at a time instead of one op at a time:

Eager:    ATen op -> NeuronDispatcher -> StableHLO -> neuronx-cc -> NEFF (cached) -> NRT
Compiled: Dynamo -> AOTAutograd -> FX passes -> StableHLO -> neuronx-cc -> NEFF (cached) -> NRT

General principles#

Start with eager mode to validate correctness, then add torch.compile for performance.
With TorchTitan, use SimpleFSDP for torch.compile workloads; it is more compiler-friendly than FSDP1 or FSDP2.
Profile before optimizing — use Neuron Explorer to identify bottlenecks.
Minimize graph breaks in torch.compile to maximize optimization opportunities.

Execution mode best practices#

1: Start development in eager mode, move to torch.compile for production

Follow this practice when: Beginning a new workload on Trainium

How to apply:

Move the model and inputs to the device and run in eager mode first:

import torch

model = MyModel().to("neuron")
x = torch.randn(8, 1024, device="neuron")
ref = model(x)

Confirm that ops run on Neuron, not on the CPU. Unsupported ops fall back to CPU silently, which costs performance:
```
torch_neuronx.clear_op_tracking()
model(x)
print("CPU fallback ops:", torch_neuronx.get_fallback_ops())
```
To make fallbacks visible instead of silent, set TORCH_NEURONX_FALLBACK_ONLY_FOR_UNIMPLEMENTED_OPS=1. Execution errors are then raised, and only ops with no Neuron implementation fall back.

Switch to torch.compile using the supported settings and compare against the eager reference:

compiled = torch.compile(model, backend="neuron", fullgraph=True, dynamic=False)
out = compiled(x)
torch.testing.assert_close(out.cpu(), ref.cpu(), rtol=1e-2, atol=1e-2)

Supported torch.compile parameters:

Parameter	Support	Notes
`fullgraph`	Supported	`True` recommended for best performance
`dynamic`	`False` only	Dynamic shapes and symbolic ints are not supported
`mode`	`"default"` only	Other modes log a warning and run as `"default"`; CUDA-graph modes give no speedup on Neuron
`options`	Not supported	Inductor tuning options do not configure Neuron; pass compiler flags through `NEURON_CC_FLAGS` instead
`disable`	Supported	Useful for A/B testing

Tips:

Keep shapes static. Pad variable-length inputs to a fixed set of bucket sizes; each distinct shape triggers a separate compile. A SymInt in the FX graph, or a neuronx-cc compilation error that contains Unbounded dynamism is disabled, means a dynamic shape got through.
Prefer float32 or bfloat16 for floating point: float64 is cast to float32 internally. int64 is supported. In eager mode, ops with native int64 support run in 64 bits (TORCH_NEURONX_INT64_MODE=native, the default). Under torch.compile, int64 graph inputs are converted to int32, so values outside the 32-bit range lose precision there.
Keep the model and inputs on the same device. CPU inputs to a compiled graph are copied to the device automatically, which hides transfer costs. Set TORCH_NEURONX_DYNAMO_DISABLE_CPU_AUTOCOPY=1 to make this an error instead.
Use contiguous tensors. Only contiguous tensors are supported under torch.compile, and non-contiguous inputs are copied on every call (see Additional recommendations).
Compiled graphs re-raise errors by default rather than silently falling back to eager (TORCH_NEURONX_DISABLE_FALLBACK_EXECUTION defaults to 1). Leave this default in place, because a silent fallback can make a broken graph look like it works.
Compile only the hot regions. torch.compile can wrap a single function or submodule, through the decorator or nn.Module.compile(), while the rest of the application stays eager.
(Experimental) TORCH_NEURONX_MODULE_AUTO_COMPILE=1 finds and compiles eligible leaf modules automatically, starting from Module.to("neuron"). A leaf that hits a graph break, guard miss, or unprofitable graph stays in eager mode. Set TORCH_NEURONX_MODULE_AUTO_COMPILE_DEBUG=1 to see which modules were promoted to compiled execution.

2: Use Adaptive Eager Execution for a balance of debuggability and performance

Follow this practice when: Running in eager mode and not actively debugging

In plain eager mode every op is its own compile and launch, so per-op overhead dominates for small ops. Adaptive eager execution buffers consecutive eligible ops and compiles and executes them as one fused unit. Buffered ops are flushed by a timeout, by torch.neuron.synchronize(), or by an op that cannot be fused. You keep the eager programming model and reduce the number of separate compiles and launches.

Adaptive eager execution is experimental and off by default in 2.33.0.

How to apply: set the variable before Python starts. It is read once, when the execution engine is created.

export TORCH_NEURONX_ENABLE_CONCATENATION=1
python train.py

Inspect what was fused:

import torch_neuronx

print(torch_neuronx.get_concatenation_status())
print(torch_neuronx.get_concatenated_ops())              # fused group -> original ops
print(torch_neuronx.get_concatenated_op_execution_counts())

Tuning knobs (advanced): TORCH_NEURONX_CONCATENATION_BUFFER_SIZE_LIMIT and TORCH_NEURONX_CONCATENATION_QUEUE_LOW_THRESHOLD. To profile the fusion itself, set TORCH_NEURONX_CONCATENATION_PROFILING=1 or call torch_neuronx.metrics.set_concatenation_profiling_enabled(True).

When to turn it off: while debugging op-level errors or numerics, also set NEURON_LAUNCH_BLOCKING=1 so that each op runs synchronously and errors point at the op that caused them. TORCH_NEURONX_ENABLE_STACK_TRACE=1 adds Python stack traces per op, at extra cost.

3: Minimize graph breaks when using torch.compile

Follow this practice when: Using torch.compile(backend=”neuron”)

Each graph break splits the model into separate NEFFs. The compiler cannot fuse or schedule across the split, and every boundary adds host round-trips and launch overhead. fullgraph=True turns any graph break into an error, so you find them early.

Breaks that TorchNeuron inserts on its own. When you import torch_neuronx, it patches Dynamo to force a graph break on ops that Neuron cannot lower. With the default fullgraph=False, those ops run in eager mode between compiled graphs. With fullgraph=True, they raise Unsupported. The list includes:

Op	Breaks when
`torch.sort`, `torch.argsort`	The sorted dimension is longer than 16384
`torch.topk`	`k` ≥ 16384
`torch.nn.functional.grouped_mm`, `torch._grouped_mm`	Always
Ops with no Neuron lowering	Always (the list is collected at import from the installed torch, torch-mlir, and Neuron decompositions)

Keep sorts and top-k selections under these limits. Otherwise every step pays for a graph split and an eager round-trip.

How to apply:

Compile with fullgraph=True during development to find breaks. Use TORCH_LOGS="graph_breaks" to list where they happen and why.
Remove common break sources:
- Host syncs inside the model: .item(), .tolist(), .cpu(), print(tensor), or Python control flow that depends on tensor values.
- Data-dependent output shapes, such as nonzero, boolean mask indexing, or unique. Rewrite them with fixed-shape alternatives like torch.where or masked arithmetic.
- Object collectives such as dist.all_gather_object. Inside compiled regions, use functional collectives (torch.distributed._functional_collectives), which trace into the graph.
- Calls into untraceable extensions. All torch_neuronx._C bindings are wrapped with torch._dynamo.disable, so call them outside the compiled region.
Write custom kernels with @nki.jit. A kernel call on Neuron tensors is integrated into the compiled graph instead of causing a break. Wrap the call in @nki_op only when you need autograd or a named PyTorch operator (see Additional recommendations).
If a break can’t be avoided, place it at a natural boundary, such as between the forward pass and the optimizer step, rather than inside a hot layer.

Inspecting what was compiled:

export TORCH_NEURONX_PRESERVE_COMPILATION_ARTIFACTS=1
export TORCH_NEURONX_DEBUG_DIR=./neuron_artifacts

Each graph writes neuron_dynamo_backend/rank<N>/<job_id>/{fx_graphs,stablehlo}/ and neuronx-cc/<hash>/log-neuron-cc.txt under the debug directory. If there are more NEFFs than you expect, look for graph breaks or recompiles.

To debug numerics, TORCH_NEURONX_DYNAMO_TENSOR_DUMPS="inputs=all;outputs=0" saves each compiled graph’s inputs and outputs as pickled NumPy arrays. You can then compare them against an eager or CPU run.

For how graph breaks work on Neuron, see the Graph breaks section of How torch.compile works on Neuron. For a step-by-step workflow to find and remove them, see How to optimize a model with torch.compile.

Distributed training best practices#

4: Choose TorchTitan’s SimpleFSDP for torch.compile workloads

Follow this practice when: Combining distributed training with torch.compile

DDP and FSDP2 (fully_shard) are the supported distributed APIs in 2.33.0, and both work with eager execution. FSDP1 (FullyShardedDataParallel) also runs. However, FSDP’s runtime hooks and parameter swapping are hard for the compiler to see through. SimpleFSDP, available in TorchTitan as the simple_fsdp module rather than as a torch.distributed API, expresses sharding as DTensor parametrizations with collectives in the traced graph. This lets the Neuron backend compile communication and compute together.

How to apply (TorchTitan, Llama3-8B, trn2.48xlarge, TP4 × FSDP16):

export NEURON_RT_NUM_CORES=64
torchrun --nnodes 1 --nproc_per_node $NEURON_RT_NUM_CORES \
  --local-ranks-filter 0 --role rank --tee 3 \
  --rdzv_backend c10d --rdzv_endpoint "localhost:0" \
  -m torchtitan.train --module simple_fsdp.llama3 --config simple_fsdp_llama3_8b \
  --parallelism.tensor-parallel-degree 4 \
  --parallelism.data-parallel-shard-degree -1 \
  --training.global-batch-size 16 \
  --training.local-batch-size 1 \
  --compile.enable --compile.backend neuron

Tips:

Use one Neuron device per process. Initialize with dist.init_process_group(backend="neuron") and call torch.neuron.set_device(local_rank).
If you stay on FSDP2 or DDP, compile the per-layer blocks rather than the outer wrapper, so the FSDP hooks remain outside the compiled region.
Inside compiled code, use functional collectives (_functional_collectives.all_reduce, all_gather_tensor, reduce_scatter_tensor), not the in-place dist.* APIs.
Create meshes with DeviceMesh / init_device_mesh. TorchNeuron hooks mesh creation to record every replica group, which SPMD compilation depends on. Process groups built by hand, outside a mesh, are not visible to it.

5: Configure sharding based on NeuronCore topology

Follow this practice when: Setting up any distributed training on Trainium

Collective cost depends on where the ranks in a group sit physically. Ranks on the same chip communicate much faster than ranks on different chips, and those are faster again than ranks on different nodes. Place the most communication-heavy dimension on the closest cores.

How to apply:

Put tensor parallelism on the innermost mesh dimension. TP communicates on every layer, so keep each TP group within one Trainium device, for example TP4 on adjacent ranks. Data-parallel/FSDP communication happens once per layer or step and can span devices and nodes.
```
from torch.distributed.device_mesh import init_device_mesh
mesh = init_device_mesh("neuron", (world_size // 4, 4), mesh_dim_names=("dp_shard", "tp"))
```
Match the rank count to the NeuronCore count. Set NEURON_RT_NUM_CORES and --nproc_per_node to the same value. TorchNeuron sets NEURON_RT_VISIBLE_CORES for each rank from LOCAL_RANK.
Choose the Logical NeuronCore (LNC) configuration on purpose. The LNC setting changes the number of visible cores and is part of the NEFF cache key, so switching it forces a recompile. See Logical NeuronCore Configuration.
Pin host threads to the device’s NUMA node. TorchNeuron does this automatically on first device use (NEURON_NUMA_PIN=1). On many-rank jobs, also bind memory from process start:
```
torchrun --no-python --nproc_per_node 64 ... neuron-numa-launch ./train.py --args
```
Check it with torch_neuronx.numa.get_numa_status().
Tune collective bucketing only after profiling. The defaults (COLLECTIVE_BUCKETSIZE_IN_MB=512, split all-gather and split reduce-scatter enabled) are a reasonable starting point.

Mixed precision best practices#

6: Use torch.autocast for automatic mixed precision

Follow this practice when: Training models where mixed precision improves throughput

Trainium’s matmul engines run fastest in BF16. torch.autocast with device_type="neuron" runs matmul- and convolution-heavy ops in lower precision. Numerically sensitive ops such as reductions, softmax, and losses stay in FP32. You don’t need to rewrite the model to get this.

How to apply:

model = MyModel().to("neuron")              # keep master weights in fp32
opt = torch.optim.AdamW(model.parameters())
compiled = torch.compile(model, backend="neuron", dynamic=False)

with torch.autocast(device_type="neuron", dtype=torch.bfloat16):
    loss = loss_fn(compiled(x), y)
loss.backward()
opt.step()

Tips:

Prefer bfloat16. It has the same exponent range as FP32, so it doesn’t need a GradScaler. Supported AMP dtypes are float32, float16, and bfloat16, but not float64.
For FSDP, set the precision through the FSDP policy (MixedPrecisionPolicy for FSDP2, MixedPrecision for FSDP1), not by casting the whole model.
Validate accuracy against an FP32 eager run. Bitwise equality with CPU results is not expected; compare with tolerances.
Don’t combine framework autocast with a compiler --auto-cast flag in NEURON_CC_FLAGS unless you mean to. The compiler default is auto_cast="none".

Performance profiling best practices#

7: Profile with Neuron Explorer before optimizing

Follow this practice when: Investigating performance bottlenecks

TorchNeuron plugs into torch.profiler and produces three aligned traces: the PyTorch framework trace, the Neuron Runtime system trace, and the Neuron device trace. You can open them in Neuron Explorer to see where time goes: host dispatch, compilation, data movement, collectives, or device compute.

How to apply:

import os, torch
from torch.profiler import ProfilerActivity, profile, schedule
import torch_neuronx
from torch_neuronx.profiling import NeuronConfig, NeuronProfiler, ProfileMode

out_dir = "./profile_out"
os.makedirs(out_dir, exist_ok=True)
cfg = NeuronConfig(
    modes=[ProfileMode.DEVICE, ProfileMode.RUNTIME],
    profile_output_dir=out_dir,
    correlate_device_time=True,   # adds NEURON device-time columns to key_averages()
)
exporter = NeuronProfiler(cfg)

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.PrivateUse1],
    experimental_config=cfg,
    schedule=schedule(wait=1, warmup=2, active=1),
    on_trace_ready=exporter.export_trace,   # writes next to the NRT traces
) as prof:
    for _ in range(4):
        train_step()
        torch.neuron.synchronize()
        prof.step()

Tips:

Warm up first. The first iterations include compilation. Use a schedule with wait/warmup so only steady-state steps are profiled. precache_node_info=True (the default) collects per-NEFF metadata during the wait steps, which keeps profiling overhead out of the active step.
Limit capture on large jobs with capture_enabled_for_nc="0" and max_events_per_nc, so the trace stays a manageable size.
Call torch.neuron.synchronize() before stopping timers, because execution is asynchronous.
Decide what to fix based on the trace:
- Many small NEFFs, or long gaps between them, point to graph breaks or eager overhead (see practices 2 and 3).
- Collective-bound steps point to sharding layout or bucket sizes (see practice 5).
- Compute-bound hot ops are candidates for NKI kernels.
Measure counters as well as time. Set TORCH_NEURONX_ENABLED_METRIC_TABLES, or use torch_neuronx.metrics.report(), to track compile counts, cache hits, and fallbacks.
Profile memory as well. torch_neuronx.memory_summary() and torch_neuronx.memory_stats() report current and peak allocation, allocation counts, and OOM counts. After recording memory history, NeuronProfiler.export_snapshot() writes a memory snapshot next to the runtime trace, so Neuron Explorer can load both together. torch_neuronx.empty_cache() releases cached blocks without freeing live tensors.

Additional recommendations#

Pin package versions in production environments to avoid unexpected behavior from updates. Pin torch, torch-neuronx, neuronx-cc, nki, and the Neuron runtime/collectives packages together. The compiler version is part of the NEFF cache key, so an upgrade also invalidates cached NEFFs.
Pass contiguous tensors to torch.compile regions. The Neuron backend supports only contiguous tensor layouts in compiled graphs. If you pass a non-contiguous input, such as a transposed, permuted, sliced, or expanded view, the backend does not fail. Instead, it makes a contiguous copy before every call. That costs an extra device copy and allocation each step, and the cost doesn’t show up in your model code.
```
# Avoid: x.T is a non-contiguous view -> hidden copy on every call
out = compiled(x.T)

# Prefer: materialize the layout once, outside the hot loop
x_t = x.T.contiguous()
for _ in range(steps):
    out = compiled(x_t)
```
- Store weights, KV caches, and other persistent buffers in the layout the graph uses, rather than permuting them on each call.
- Check layouts at the graph boundary with t.is_contiguous(). For channels_last and other non-default memory formats, convert to torch.contiguous_format before the compiled region.
- DTensor: local shards (dtensor.to_local()) can be non-contiguous after redistribution, particularly in tensor-parallel attention. The TorchTitan tutorials patch the SDPA path with a force_local_contiguous workaround for this (pytorch/pytorch#173041). Call .contiguous() on local shards before they enter a compiled region or a custom kernel.
- Layout changes inside the graph, such as transpose, permute, and view, are fine, because the compiler handles them. The requirement applies only to tensors that cross the graph boundary.
Use the TorchDynamo cache to avoid recompilation across runs. TorchNeuron caches at two levels:
- The persistent NEFF cache (on by default) reuses compiled NEFFs whenever the HLO and compiler config haven’t changed. Point it at shared storage so all instances benefit:
```
export TORCH_NEURONX_NEFF_CACHE_DIR=/fsx/neff_cache        # shared (e.g. FSx)
export TORCH_NEURONX_NEFF_LOCAL_CACHE_DIR=/tmp/local_cache # must be local POSIX FS (locks)
```
- The Dynamo precompile cache skips re-tracing torch.compile regions on warm starts:
```
export TORCH_CACHING_PRECOMPILE=1        # stored under TORCHINDUCTOR_CACHE_DIR
```
- NEFF cache keys are built from the lowered HLO plus the compiler config (TORCH_NEURONX_NEFF_KEY_FROM_HLO=1 is the default). Python edits that don’t change the lowered graph still hit the cache, and entries carry across sessions. Changing the target, the LNC setting, the compiler version, or compiler flags produces a new key.
- Fill the cache ahead of time with torch_neuronx.compiler.compile_only() (or TORCH_NEURONX_COMPILE_ONLY=1). This compiles NEFFs without running them. On a host without a Neuron device, also set NEURON_RT_FAKE_INSTANCE_TYPE. The outputs are placeholders, so don’t use this mode to check results.
- If a compiler process crashes, it can leave a stale lock at {local_cache_dir}/{hash}/locks/compile.lock, which blocks other processes for up to TORCH_NEURONX_NEFF_CACHE_LOCK_TIMEOUT seconds (default 1200). Delete it if that happens.
Consider NKI kernels for performance-critical operations that standard PyTorch operators cannot express efficiently. Write the kernel with @nki.jit and call it on Neuron tensors. The call works in eager mode and inside torch.compile without a graph break. To train through the kernel, or to call it as a named PyTorch operator, wrap the call with @torch_neuronx.nki_op and register a backward pass:
```
import nki, nki.language as nl, nki.isa as nisa
from torch_neuronx import nki_op

@nki.jit
def add_kernel(x1, x2, y):
    a = nl.ndarray(x1.shape, x1.dtype, nl.sbuf)
    b = nl.ndarray(x2.shape, x2.dtype, nl.sbuf)
    c = nl.ndarray(y.shape, y.dtype, nl.sbuf)
    nisa.dma_copy(a, x1); nisa.dma_copy(b, x2)
    nisa.tensor_tensor(c, a, b, nl.add)
    nisa.dma_copy(y, c)
    return y

@nki_op("mylib::nki_add", mutates_args={"y"})
def nki_add(x1: torch.Tensor, x2: torch.Tensor, y: torch.Tensor) -> None:
    add_kernel(x1, x2, y)
```
- Declare every input the kernel writes in mutates_args. An operator that writes into an input returns None, because a torch.library custom operator cannot return a tensor that aliases an input.
- For a backward pass, call register_autograd on an operator that does not write into its inputs, as in step 4 of How to work with NKI kernels from torch-neuronx. PyTorch does not attach a backward pass to an operator that writes into an input, such as the in-place nki_add above.
- You do not need torch_neuronx.wrap_nki or the fake_fn argument of nki_op for kernels that you call on Neuron tensors. Both are for code that runs on meta tensors.
For the full workflow, see How to work with NKI kernels from torch-neuronx.

The built-in NKI flash-attention path for scaled_dot_product_attention is enabled by default when seqlen % 2048 == 0. Set TORCH_NEURONX_ENABLE_NKI_SDPA=0 to disable it. Pass any value that changes the kernel’s control flow as a kernel argument; otherwise the NKI trace cache can return a stale kernel.

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.

Best practices for TorchNeuron

Contents

Best practices for TorchNeuron#

Overview#

General principles#

Execution mode best practices#

Distributed training best practices#

Mixed precision best practices#

Performance profiling best practices#

Additional recommendations#

Further reading#