This document is relevant for: Trn2, Trn3

How to optimize a model with torch.compile#

Task overview#

This topic discusses how to apply torch.compile(backend="neuron") to a working eager-mode model on AWS Trainium and tune it for throughput.

Prerequisites#

Instructions#

1: Wrap your model or callable with the Neuron backend

Apply the backend with the functional form:

import torch

compiled = torch.compile(model, backend="neuron", fullgraph=True)

Or with the decorator form on a function:

@torch.compile(backend="neuron", fullgraph=True)
def compiled_forward(x):
    return model(x)

Pass backend="neuron" and remove CUDA-only settings such as CUDA graphs and mode="max-autotune". The backend supports mode="default" only; other modes emit a warning and fall back to default. Dynamic shapes are not supported, so leave dynamic unset or pass dynamic=False. Use fullgraph=True so graph breaks are raised rather than silently splitting the model.

2: Warm up, then time steady-state iterations

The first call with a given input signature triggers compilation, which can take several minutes for large models. Later calls with the same shapes and dtypes reuse the compiled NEFF. Time a whole training iteration, not only the compiled forward pass, because host work and the optimizer do not appear in a timer that covers the forward and backward passes alone. Run a few warmup iterations, then time the steady-state iterations:

import time
import torch_neuronx

compiled = torch.compile(model, backend="neuron", fullgraph=True)

def train_step(batch):                     # one full training iteration
    optimizer.zero_grad()
    loss = loss_fn(compiled(batch["input"]), batch["target"])
    loss.backward()
    optimizer.step()
    return loss

for batch in warmup_batches:               # warm up: triggers and caches compilation
    train_step(batch)

torch_neuronx.synchronize()
start = time.perf_counter()
for batch in timed_batches:
    train_step(batch)
torch_neuronx.synchronize()
iteration_ms = (time.perf_counter() - start) * 1000 / len(timed_batches)

Synchronize the device before you read the clock. Neuron operations run asynchronously, so a timer without a synchronize measures only the time to enqueue the work.

3: Reuse compiled NEFFs across runs

Compiled NEFFs are cached on disk (TORCH_NEURONX_NEFF_CACHE_DIR, default /tmp/neff_cache). To reuse them across runs and across instances, point the cache directory at a shared mount such as Amazon FSx, and keep the lock and scratch directory on a local disk:

export TORCH_NEURONX_NEFF_CACHE_DIR=/fsx/neuron/neff_cache
export TORCH_NEURONX_NEFF_LOCAL_CACHE_DIR=/tmp/local_cache

TORCH_NEURONX_NEFF_LOCAL_CACHE_DIR (default /tmp/local_cache) holds the backend’s file locks and scratch files only. The cached NEFFs themselves stay in TORCH_NEURONX_NEFF_CACHE_DIR. Keep the lock directory on a local POSIX filesystem, because file locks are not reliable on NFS or Lustre.

The cache key includes the compiler version, the target platform, the logical NeuronCore configuration, and the compiler flags (NEURON_CC_FLAGS). A change to any of these makes a new key, so pin torch-neuronx and the compiler in production. See torch.compile backend reference for all cache-directory variables.

4: Find graph breaks and recompilations

Graph breaks and recompilations are the most common reasons a compiled model runs slower than expected. Use the standard PyTorch logging to list them:

TORCH_LOGS="graph_breaks,recompiles" python train.py

Each graph-break entry names the source line and the reason, for example a print() call, a .item() call, or a Python builtin that Dynamo cannot trace. Each recompilation entry names the guard that failed, for example a changed input shape or a changed Python attribute. Fix these before you tune anything else. See best practices 1, 2, and 3.

5: Measure compilation

Use the backend metrics API to see how many graphs compiled and how long each stage took:

from torch_neuronx.neuron_dynamo_backend import (
    get_dynamo_metrics,
    reset_dynamo_metrics,
)

reset_dynamo_metrics()
compiled = torch.compile(model, backend="neuron", fullgraph=True)
compiled(*example_inputs)          # first call triggers compilation

for m in get_dynamo_metrics():
    print(m.graph_name, m.graph_node_count, m.torch_neuronx_compile_us)

Each entry is one compiled graph, named {model_name}_{aot_id}_{graph_type}_{nth}. For a training step, each compiled region produces a forward graph and a backward graph. More graphs than you expect means that a graph break split a region. To export the metrics as CSV at process exit, set TORCH_NEURONX_ENABLED_METRIC_TABLES=graph_stats (and optionally TORCH_NEURONX_METRICS_DIR).

Confirm your work#

Confirm that the model compiled once and that later calls reuse the compiled graph. Reset the metrics, run the model twice, and check the compiled-graph count each time:

reset_dynamo_metrics()
compiled(*example_inputs)                       # first run
first = len(get_dynamo_metrics())

reset_dynamo_metrics()
compiled(*example_inputs)                       # second run, same shapes
second = len(get_dynamo_metrics())

assert first >= 1 and second == 0, (first, second)

The first run compiles at least one graph. The second run with the same input shapes compiles nothing, which confirms that the compiled graph was reused. If the second count is not zero, the inputs changed shape or dtype between the calls.

To confirm on-disk cache reuse across processes, run the script a second time. The second process starts much faster because it loads the NEFFs from TORCH_NEURONX_NEFF_CACHE_DIR instead of compiling them again.

Common issues#

Uh oh! Did you encounter an error or other issue while working through this task? Here are some commonly encountered issues and how to address them.

The model silently runs slower than expected

The model recompiles every iteration

The model recompiles on the second or third iteration only

  • Possible solution: Python state that the compiled code reads changed after the first call. See best practice 3.

A small code change makes one region much slower

  • Possible solution: The change added a graph break inside the region. Find it with step 4. See best practice 1.

The compiled model gives different results from eager mode

  • Possible solution: Disable dropout and compare again. See best practice 8.

A larger batch size does not increase throughput

  • Possible solution: Step time can grow linearly with batch size. See best practice 6.

Best practices#

These best practices target fast steady-state throughput (after compilation), effective cache reuse, and predictable compilation behavior from the Neuron TorchDynamo backend.

In these practices, a region is a module or function that you pass to torch.compile. A region without graph breaks compiles to one graph, and a graph break splits it into more than one graph.

General principles#

  • Validate correctness in eager mode first, then apply torch.compile for performance.

  • Keep each region free of graph breaks (fullgraph=True) and give it fixed input shapes.

  • Choose how many regions to compile by measurement, not by assumption. See best practice 5.

  • Measure every change on the full training iteration. The sum of isolated measurements is not equal to the full-iteration result.

  • Keep the default fail-loud behavior in production.

Compilation best practices#

1: Minimize graph breaks

  • Follow this practice when: Using torch.compile(backend="neuron") on any model.

Each graph break splits a region into separately compiled graphs, which reduces the scope the compiler can optimize across. Common causes include data-dependent control flow, unsupported operators, and calls into code Dynamo cannot trace. Calls that run only during tracing also break the graph, for example a print() or logging call in a helper that builds a configuration object on its first call. Refactor those paths, or move them out of the compiled region.

To read a scalar from the device, such as a loss value or a counter, read it outside the compiled region. .item() inside the region breaks the graph and synchronizes the host with the device. NKI kernels registered as PyTorch operators integrate into the graph without forcing a break.

TorchNeuron also forces a graph break on operations that Neuron cannot lower in a compiled graph. With fullgraph=True, each of these raises an error instead:

  • torch.sort and torch.argsort when the sorted dimension is longer than 16,384.

  • torch.topk when k is 16,384 or more. The limit applies to k, not to the size of the dimension you select from.

  • torch.nn.functional.grouped_mm and torch._grouped_mm, always.

  • Any other operation that has no Neuron lowering, always. TorchNeuron builds this list when it is imported, from the installed PyTorch, torch-mlir, and Neuron decompositions.

Keep sorts and top-k selections under these limits, and move the other operations out of the hot compiled region. For other operations that commonly cause graph breaks, see Graph breaks.

2: Keep input shapes static to avoid recompilation

  • Follow this practice when: Running compiled models with variable batch or sequence sizes.

Each new input shape or dtype triggers a new compilation, which stalls the step that sees it. Limit the shapes the model sees:

  • Use a fixed batch size. For a DataLoader, set drop_last=True so the last batch is not smaller than the others.

  • Pad variable-length inputs to a small set of bucket sizes, and pass a mask for the padding. Each bucket compiles once.

import torch.nn.functional as F

BUCKETS = (128, 256, 512, 1024)

def pad_to_bucket(input_ids):              # input_ids: [batch, seq_len]
    seq_len = input_ids.shape[-1]
    target = next(b for b in BUCKETS if b >= seq_len)
    return F.pad(input_ids, (0, target - seq_len))

Run each bucket once during warmup, so that no compilation occurs in the timed steps. The graph count from get_dynamo_metrics() should match the number of buckets.

3: Initialize lazy state before the first compiled call

  • Follow this practice when: The model or its custom operators build caches, kernel configurations, or index tables on first use.

Dynamo guards on the Python state that compiled code reads. If that state is built lazily during the first call, the first trace sees one value and later calls see another, so the region recompiles. If the code that builds the state prints, logs, or calls an untraceable function, the region also breaks. Build the state before the first compiled call:

# Build caches, kernel configs, and static index tables in eager mode first.
# prepare_static_state is a placeholder for your model's own setup method.
model.prepare_static_state(example_inputs)

compiled = torch.compile(model, backend="neuron", fullgraph=True)

The same applies to data-independent tensors such as gather indices, masks, and geometry tables that depend only on input shapes. Compute them once, move them to the neuron device, and pass them to the compiled region, instead of recomputing them in every iteration.

Graph partitioning best practices#

4: Compile large, stable regions separately

  • Follow this practice when: Training a model with several large components, such as a backbone, task heads, and a loss, or when only some components benefit from compilation.

Compile each large component as its own region with fullgraph=True. Each region produces its own NEFFs, so you can change one component without a full recompilation, and you can measure each region separately. Include the loss computation in a compiled region where possible. A loss computed in eager mode runs as many small operations and can be a large part of the step time.

model.backbone = torch.compile(model.backbone, backend="neuron", fullgraph=True)
for i, head in enumerate(model.heads):
    model.heads[i] = torch.compile(head, backend="neuron", fullgraph=True)
loss_fn = torch.compile(loss_fn, backend="neuron", fullgraph=True)

Components that you do not compile stay in eager mode. Compiled and eager code share the same neuron tensors, so no copy occurs at the boundary. Keep each compiled region large enough to justify its compilation time, and time it against its eager version before you keep it.

5: Measure graph boundaries instead of assuming one graph is fastest

  • Follow this practice when: A compiled region is large, or the region mixes NKI kernels with compiler-generated code.

One large graph gives the compiler the most scope, but it is not always the fastest option. A graph boundary in the right place can reduce step time, and a boundary in the wrong place can increase it. If a large region is slower than you expect, compile its parts as separate regions and compare the step time. Keep the partitioning that is faster, and record why it is there so a later change does not remove it by mistake.

A common example is the backward pass of a data-parallel model. Both partitionings let the all-reduce of a layer’s weight gradients start as soon as those gradients are available, so both can overlap communication with the backward pass of the next layer. What differs is where that overlap can happen. When the whole model compiles as one backward graph, the overlap is confined to that graph. With smaller graphs, or with FSDP v2 and host-side collectives, a collective can also overlap with work outside the backward pass, such as the optimizer, which gives the scheduler more overlap budget. Measure both partitionings on the full iteration before you choose.

6: Measure throughput per sample when you change the batch size

  • Follow this practice when: Increasing the batch size to raise throughput.

A larger batch size increases throughput only if step time grows more slowly than the batch size. For workloads that are limited by data movement, such as gather-heavy or scatter-heavy models, step time can grow linearly with batch size. Compute samples per second for each batch size, and check device memory use with torch_neuronx.memory_summary(). If throughput per sample does not improve, keep the smaller batch size and use data parallelism across NeuronCores instead. When you run one process per NeuronCore, compare configurations by the sum of the per-process rates, not by the rate of one process.

Custom kernel best practices#

7: Replace profiled hotspots with NKI kernels, one site at a time

  • Follow this practice when: A profile shows a few operations that dominate the step, such as large gathers, scatter-adds, or segment reductions.

Profile first, and use real inputs. A profile captured with random inputs does not show the real cost of data-dependent operations such as gathers and scatters, because random indices give a different memory access pattern. Profile the full step with torch.profiler and view it in Neuron Explorer. See How to profile memory usage in torch-neuronx.

For an operation that dominates the profile, write an NKI kernel and register it with torch_neuronx.nki_op (and an autograd backward for training). A kernel registered as a PyTorch operator stays inside the compiled graph. See How to work with NKI kernels from torch-neuronx.

An NKI kernel is not faster at every call site. For an operation that appears at many sites, such as a convolution, benchmark each site’s shape separately against the compiler-generated version, and enable the kernel only where it is faster. Keep an environment variable or flag that turns each kernel off, so you can compare performance and numerics at any time.

Numerical validation best practices#

8: Validate each compiled region against eager mode

  • Follow this practice when: Compiling a model for the first time, or after any performance change.

Compare the compiled outputs, losses, and gradients against eager mode on the same inputs after every optimization. To make the comparison deterministic, disable dropout and other random operations, because random numbers generated in a compiled graph do not match the eager random stream. If the results differ, compile smaller regions until you find the region that diverges, and compare it against an eager reference, or against a CPU reference in float64.

Check gradients as well as forward outputs, and check for values that are finite but wrong, not only for NaN and Inf. Expect small differences from bfloat16 and from reductions that run in a different order. In models with discrete steps, such as Hungarian matching or topk, a small difference can change a decision and cause a large change in gradients. For these models, compare losses and the gradients of the affected components, or hold the matching or topk result constant during the comparison.

Caching best practices#

9: Compile once, then start the rest of the fleet

  • Follow this practice when: Running the same model on many instances.

Set up the shared NEFF cache as in step 3. Compile on one instance first, then start the other instances, so that they load the NEFFs instead of compiling them again. Keep the compiler version, NEURON_CC_FLAGS, and the logical NeuronCore configuration identical across the fleet, because a difference in any of them makes a new cache key.

Fallback and reliability best practices#

10: Keep the default fail-loud behavior in production

  • Follow this practice when: Promoting a compiled workload to production or CI.

By default (TORCH_NEURONX_DISABLE_FALLBACK_EXECUTION=1) the backend re-raises compilation and execution failures instead of silently falling back to eager execution. Keep this default in production and CI so a compilation regression surfaces as an error rather than a silent slowdown. Re-enable fallback with TORCH_NEURONX_DISABLE_FALLBACK_EXECUTION=0 only where degrading to eager execution is preferable to failing, and expect reduced performance when a region falls back.