AWS Neuron SDK, in depth: the compiler, runtime and tools, the 2026 move from XLA to native PyTorch, compile caches, NKI and version pinning

The layers, from your code down to the silicon

Read the stack from the top down. You write ordinary framework code. A framework integration turns it into something a compiler can see. There are three integrations. TorchNeuron is the native PyTorch backend. torch-neuronx is the older integration built on PyTorch/XLA. JAX NeuronX handles JAX. The graph compiler, neuronx-cc, takes those graphs, together with any custom NKI kernels, and produces NEFFs. A NEFF fixes tensor shapes, assigns operations to the engines inside each NeuronCore and schedules data movement. The runtime and kernel driver load NEFFs and own the cores. They move tensors between host memory and the chip's HBM, and they run collective communication across cores and chips. Host tools sit beside the stack and report what the runtime is doing.

Each layer is versioned, and the layers must agree. A runtime too old for the compiler that built a NEFF will refuse to load it. That is why AWS ships the stack as one numbered SDK release, and why pinning is the most important operational habit on Neuron.

Why a compiler sits in the middle

A NeuronCore contains specialised engines: a tensor engine for matrix multiplies, plus vector, scalar and general-purpose SIMD engines for everything else. Each has its own instruction stream. Keeping them all busy requires deciding in advance which engine runs each operation, where every intermediate lives in on-chip memory, and when DMA transfers start. That whole-program scheduling problem is the compiler's job, so it needs the whole graph, with concrete shapes, before anything runs. Three consequences follow:

• Compilation is slow and happens up front. A large transformer graph can take many minutes to compile.
• Shapes are part of the program. A batch with a different sequence length is a different graph and needs a different NEFF. Unbounded shape variety means unbounded compiles.
• Caching is the main lever. Identical graphs with identical compiler flags produce the same NEFF. Store it once and every later run, node or restart can skip the compile.

If you come from GPUs, the closest analogy is torch.compile with dynamic=False, where compile cost is higher and there is no fallback to fast hand-written kernels for a graph that fails to compile. The site's torch.compile deep dive explains guards and recompiles, and those ideas carry over directly.

The 2026 training transition: what changed at 2.28 and 2.29

Training on Trainium used to mean torch-neuronx on PyTorch/XLA, with NeuronX Distributed (NxD Core and NxD Training) for model parallelism. According to the Neuron end-of-support announcement, release 2.28 (PyTorch 2.9) is the last release that includes the NxD Training library. Starting with 2.29 (PyTorch 2.10), the NxD Core training APIs and the PyTorch/XLA package are no longer supported for training. The PyTorch framework page states that PyTorch 2.9 is the last torch-neuronx version based on PyTorch/XLA, and that torch-neuronx moves to the native implementation from PyTorch 2.10.

One caution about status. On the day this was checked, the TorchNeuron overview page described it as available only through a closed Beta, while other pages describe it as shipping since release 2.27. Check its status for your account and release before committing a production plan to it.

Native PyTorch on Trainium: a training loop

TorchNeuron plugs into PyTorch through the standard out-of-tree device mechanism. Code moves to Trainium the same way it moves to a GPU: put tensors on the "neuron" device and optionally wrap the model in torch.compile(backend="neuron"). Eager mode is for debugging. Compiled mode hands whole graphs to the Neuron compiler and is where the throughput is. The Neuron-specific parts below are the device name, the compile backend and the shape discipline in the collate function.

# Native PyTorch on Trainium (TorchNeuron, PyTorch 2.10+, Trn2/Trn3).
# Device and compile backend names are from the TorchNeuron overview page; everything else is plain PyTorch.
import torch
from torch.utils.data import DataLoader

SEQ = 2048                                   # one fixed sequence length: every new shape is a new compile

def collate(batch):
    ids = torch.full((len(batch), SEQ), PAD_ID, dtype=torch.long)
    for i, row in enumerate(batch):
        row = row[:SEQ]
        ids[i, :len(row)] = torch.tensor(row)
    return ids

device = torch.device("neuron")              # replaces "cuda"
model = build_model().to(device)
opt = torch.optim.AdamW(model.parameters(), lr=2e-5)
step_fn_model = torch.compile(model, backend="neuron")   # eager works too; compile for throughput

loader = DataLoader(dataset, batch_size=8, collate_fn=collate, drop_last=True)  # drop_last: no odd last shape
for step, ids in enumerate(loader):
    ids = ids.to(device)
    logits = step_fn_model(ids[:, :-1])
    loss = torch.nn.functional.cross_entropy(
        logits.reshape(-1, logits.size(-1)), ids[:, 1:].reshape(-1), ignore_index=PAD_ID)
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    opt.step()
    opt.zero_grad(set_to_none=True)
    if step % 50 == 0:
        print(step, loss.item())             # .item() forces a sync: keep it off the hot path

The fixed SEQ and drop_last=True matter. Without them every distinct batch length is a distinct graph and a compiler run. If lengths really vary, pad each batch up to the nearest of a few buckets, such as 512, 1024 and 2048.

To scale out, the documentation lists DDP, FSDP, DTensor and tensor parallelism, with the NeuronCore as the unit of distribution. The FSDP sharding model you know from GPUs applies. Take the process-group backend name and launcher settings from the TorchNeuron documentation for your release rather than guessing them. Precision follows the usual rules, as described in the site's mixed precision guide.

Pinned at 2.28: the XLA flow and its compile cache

Many fleets will run the XLA path for a while yet. In PyTorch/XLA, tensor operations are recorded into a graph, and xm.mark_step() cuts the graph, compiles it (or fetches it from the cache) and runs it:

import torch_xla.core.xla_model as xm

device = xm.xla_device()
model = build_model().to(device)
for ids in loader:                          # tensors are recorded lazily into a graph...
    loss = loss_fn(model(ids.to(device)), ids.to(device))
    loss.backward()
    opt.step(); opt.zero_grad()
    xm.mark_step()                          # ...which is cut, compiled (or fetched from cache) and run here

A cold first run compiles every graph it meets, one after another. neuron_parallel_compile runs your script in a graph-extraction mode for a few steps and compiles every graph in parallel into the persistent cache, so the real run starts warm:

# Pinned at Neuron 2.28 / PyTorch 2.9: the XLA training flow (no longer supported from 2.29).
export NEURON_COMPILE_CACHE_URL="s3://ml-artifacts/neuron-cache/llama-ft/"   # shared across nodes and runs
export NEURON_CC_FLAGS="--model-type transformer"                          # check the flag list for your compiler version
export XLA_DOWNCAST_BF16=1        # documented: float32 -> bf16, float64 -> float32 (XLA_USE_BF16 is deprecated)

# 1) Extract every graph the job will produce and compile them in parallel, filling the cache.
neuron_parallel_compile torchrun --nproc_per_node=32 train.py --max_steps 10

# 2) The real run finds every NEFF in the cache and starts training immediately.
torchrun --nproc_per_node=32 train.py

The documented environment variables are worth knowing by name. NEURON_COMPILE_CACHE_URL points the cache at a local directory (the default is /var/tmp/neuron-compile-cache) or at S3 when the value starts with s3://. NEURON_CC_FLAGS passes compiler options; a --cache_dir given there takes precedence over the URL. NEURON_EXTRACT_GRAPHS_ONLY dumps graphs without executing them. The cache key includes the compiler flags, so changing NEURON_CC_FLAGS on one node and not the others silently splits the cache and puts compile time back on the critical path.

Inference: where each serving path fits

Serving is a separate branch with its own release notes. torch_neuronx.trace compiles a model for fixed example inputs into a module with an embedded NEFF. NxD Inference serves large models with tensor parallelism and KV caching. vLLM Neuron, Beta and limited to Trn2 and Trn3 in release 2.32, adds vLLM's OpenAI-compatible server. Bucketing and sizing are covered in the Inferentia deep dive, including the release where NxD Inference stopped supporting Inf2.

NKI: when to write your own kernel

The Neuron Kernel Interface (NKI) is a Python-embedded language for writing kernels that run directly on a NeuronCore. You control tiling, on-chip memory placement and which engine runs each step. The compiler accepts NKI kernels alongside framework graphs, so a custom kernel becomes one more node in the compiled program. AWS also publishes an NKI Library of pre-optimised kernels.

Treat NKI like writing CUDA or Triton kernels. Profile first, and only write a kernel when one operation dominates step time and the compiler's lowering is clearly poor. Check the NKI Library first. Keep a reference PyTorch implementation and a numerical test against it, run on every SDK upgrade. NKI syntax is not shown here because the API has moved between releases; take it from the docs for your pinned release.

Observing a Neuron host

neuron-ls                         # devices, NeuronCores and which PID holds them
neuron-top                        # live per-core utilisation and memory
neuron-monitor | your-exporter    # JSON metrics stream to ship to Prometheus or CloudWatch

# Give two processes on one host disjoint cores instead of letting both grab everything:
NEURON_RT_VISIBLE_CORES=0-15  python serve_a.py &
NEURON_RT_VISIBLE_CORES=16-31 python serve_b.py &

neuron-ls answers "which process owns which cores"; run it first when a job fails because a device is busy, usually a leftover process. neuron-top answers "are the cores busy"; low utilisation usually means data loading or graph tracing on the host is the bottleneck. neuron-monitor emits a machine-readable stream for dashboards and alerts. NEURON_RT_VISIBLE_CORES restricts a process to a range of cores, which is how you run several independent models on one instance. On Kubernetes, the Neuron device plugin plays the same role through resource requests, as described in the site's EKS deep dive.

Failure modes and how to recognise them

Version pinning: the operational playbook

Pin the SDK as a unit. Run from a Neuron Deep Learning AMI or container that matches one release, and record the release number in the image tag. If you build your own image, pin the host driver and the runtime, compiler and framework packages to the same release. Put the release in the cache path too, for example s3://ml-artifacts/neuron-cache/2.28/, so an upgrade starts with a clean cache.

Upgrade as a project, not a patch. Read the release notes for deprecations first; 2.29 shows that a release can remove whole workflows. Then rebuild the image, warm the new cache, compare a short run against a known-good loss curve, and only then move the fleet.

Worked example: planning the migration of a 2.28 fine-tuning job

Suppose a team fine-tunes a 7B model on trn1.32xlarge instances (16 Trainium chips with two NeuronCores each, 32 cores in total) using torch-neuronx and NxD Training on release 2.28. They want PyTorch 2.10 features.

1. Freeze the current job. Record the image tag, the compile cache path and a reference loss curve for the first 2,000 steps.
2. Check hardware. TorchNeuron documents Trn2 and Trn3, so this is also an instance change. Price trn2 capacity first.
3. Port the code. Replace NxD parallel layers with standard modules plus FSDP or DTensor. Remove every xm.mark_step() and XLA environment variable.
4. Keep the fixed sequence length or buckets, and set drop_last on every loader, evaluation included.
5. Run eager on one node until the early loss matches the reference, then enable the compile backend.
6. Scale out and compare the loss curve and tokens per second with the frozen job. Only then retire the 2.28 image.

What to do next

1. Run neuron-ls on one of your hosts and write down the SDK release you are actually running.
2. If you train with torch-neuronx or NxD Training, pin release 2.28 explicitly in your images today.
3. Set NEURON_COMPILE_CACHE_URL to a release-specific S3 path shared by every node.
4. Make every data loader produce fixed or bucketed shapes, and set drop_last.
5. Add neuron-monitor output to your dashboards next to step time.
6. Prototype one training step with TorchNeuron on Trn2 to size the migration, after confirming its current availability.
7. Read the Inferentia deep dive before you design a serving path.