Profiling TorchNeuron, Part 1: System-Level Profiling on AWS Trainium

Profiling TorchNeuron, Part 1: System-Level Profiling on AWS Trainium

Many machine learning practitioners have an intuition for making models run faster, but determining why a model is slow requires systematic analysis. This walkthrough explains how to analyze inference performance and identify where execution time is spent.

The examples use AWS Trainium, an AI chip designed for deep learning training and inference, together with TorchNeuron, its open-source native PyTorch backend. TorchNeuron supports standard PyTorch workflows on Trainium, including eager execution, torch.compile, and standard PyTorch APIs, with limited code changes. It also integrates with the Neuron Kernel Interface (NKI), which supports custom operators and hardware-specific optimization.

The examples demonstrate how to:

• Configure torch.profiler with the Neuron backend and capture a trace.
• Read the key_averages() table to identify operations consuming Neuron device time and compare their Self CPU time.
• Follow an operation in Neuron Explorer from framework dispatch through runtime scheduling to execution on a NeuronCore.
• Compare these signals across three implementations of the same operation.

Neuron as a PyTorch Backend

Before profiling a workload, it helps to understand how Neuron integrates with PyTorch.

The Neuron backend uses PyTorch's PrivateUse1 extension point. This allows an accelerator to integrate with PyTorch without adding a dedicated dispatch key such as cuda. The torch_neuronx package works with upstream PyTorch to register Neuron as a native backend. More information about this extension mechanism is available in PyTorch's documentation on Facilitating New Backend Integration by PrivateUse1.

Setting Up the Profiler

Two objects are used to profile a Neuron workload:

• NeuronConfig controls the Neuron data streams and NeuronCores to capture, as well as the output location. It is passed to torch.profiler.profile as experimental_config.
• NeuronProfiler accepts a NeuronConfig and provides an export_trace callback. The callback writes the framework trace, trace.json, alongside the Neuron system traces.

nc = NeuronConfig(
    modes=[ProfileMode.DEVICE, ProfileMode.RUNTIME],
    profile_output_dir=OUTPUT_DIR,
    capture_enabled_for_nc='0',
    correlate_device_time=True,
)
exporter = NeuronProfiler(nc)

The modes parameter selects the data streams to capture, including device instructions and runtime traces. capture_enabled_for_nc selects the NeuronCores to trace. When correlate_device_time=True, device timing is added to the key_averages() table, allowing CPU and Neuron device time to be compared. The complete parameter reference and available ProfileMode values are documented in the NeuronConfig Profile Modes documentation.

Neuron Explorer

Neuron Explorer is a browser-based viewer for system traces, device timelines, and per-engine breakdowns. It can be accessed through an SSH tunnel or the Visual Studio Code extension. In this workflow, Explorer is used to verify the numerical results from key_averages() and inspect the events between dispatches.

After profiling, the output directory contains the system trace and, when device data is captured, matching .neff and .ntff files. The Neuron Executable File Format, or NEFF, contains the compiled Neuron executable. The corresponding NTFF contains its hardware instruction trace. The two files share a numeric hash that identifies them as a pair.

For a system or combined profile, upload the complete output directory to Neuron Explorer. For a device-only profile, upload the matching .neff and .ntff files individually.

Profiling RMSNorm

The first analysis uses RMSNorm, a small operation with a clear sequence of arithmetic and reduction steps.

Three implementations of the same operation are profiled:

1. A naive RMSNorm implementation using separate PyTorch operations: multiplication, mean, addition, reciprocal square root, and multiplication.
2. torch.rms_norm(), the PyTorch built-in lowered to a single fused Neuron kernel.
3. The naive RMSNorm implementation compiled with torch.compile.

All three scripts use the same inputs, profiler configuration, and schedule, making the results directly comparable. Reproduction scripts are available in the torch-neuronx examples directory.

Profile 1: Separate PyTorch Operations

The first implementation writes RMSNorm using individual PyTorch operations. Each line maps to one or more ATen operations and can create intermediate tensors.

device = torch.device('neuron')

x = torch.randn(4, 512, dtype=torch.bfloat16, device=device)
weight = torch.ones(512, dtype=torch.bfloat16, device=device)

def rms_norm_fn(t):
    x_sq = t * t                                    # square
    mean_sq = x_sq.mean(dim=-1, keepdim=True)       # mean of squares
    rms = torch.rsqrt(mean_sq + 1e-6)               # 1 / sqrt(mean_sq + eps)
    return t * rms * weight                         # normalize + scale

The profiler requires two Neuron-specific settings:

1. Include ProfilerActivity.PrivateUse1 to collect activity from the Neuron device.
2. Pass NeuronConfig as experimental_config to enable Neuron-specific Kineto features.

The schedule uses wait=2, warmup=1, and active=1 across four prof.step() calls. This records one warmed-up iteration without including initial setup overhead. torch.neuron.synchronize() is called inside the loop so the profiler waits for device execution to finish before exporting the trace.

with torch.profiler.profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.PrivateUse1],
    experimental_config=nc,
    record_shapes=False,
    on_trace_ready=exporter.export_trace,
    schedule=torch.profiler.schedule(wait=2, warmup=1, active=1),
) as prof:
    for step in range(4):
        with torch.no_grad():
            _ = rms_norm_fn(x)
        torch.neuron.synchronize()
        prof.step()

After profiling, inspect the key_averages() table and sort it by self_privateuse1_time_total to place the operations with the highest device cost first.

print(prof.key_averages().table(sort_by="self_privateuse1_time_total", row_limit=4))

The relevant results are:

Self NEURON total: 2.355 ms

The table shows which ATen operations ran, how often they ran, and how much time they spent on the CPU and Neuron device. aten::mul has the highest total device time because the function performs three multiplications.

Each row corresponds to part of the function:

• aten::mul runs for x * x, x * rms, and the final multiplication by weight.
• aten::mean reduces the last dimension.
• aten::add adds the epsilon value and triggers a host-to-device copy because the scalar begins on the CPU.
• aten::rsqrt computes the reciprocal square root.

Reading the Trace in Neuron Explorer

In Explorer, time moves from left to right, while the event lifecycle moves from top to bottom. Each ATen operation follows the same general sequence: creation in the framework lane, scheduling in the runtime lane as nrta_execute_schedule, and device execution in the hardware lane as nc_exec_running.

The system trace is organized into lanes. In the default grouping mode, labels use the pattern instance/PID/subsystem/identifier.

In eager mode, each operation runs independently. The dependency chain viewer can be used to follow aten::mul.Tensor from the framework layer through runtime scheduling to device execution.

Profile 2: torch.rms_norm

The second implementation replaces the individual operations with PyTorch's built-in rms_norm function. Neuron lowers it to one fused kernel.

def rms_norm_fn(t):
    return torch.rms_norm(t, [512], weight)

The resulting table contains one device operation:

Self NEURON total: 403.3 μs

Comparing the Lowering Path

In the naive implementation, each operation crosses the dispatch boundary independently and follows its own framework-to-runtime-to-hardware path. In the fused implementation, one Python call covers the complete computation. The trace contains one dispatch and one hardware event, with no gaps between the component operations.

Profile 3: Naive PyTorch with torch.compile

The final example compiles the original naive RMSNorm function with torch.compile.

With @torch.compile(backend="neuron"), TorchDynamo captures the Python bytecode as an FX graph, and AOTAutograd lowers it to ATen operations. The Neuron backend receives the complete graph instead of dispatching each operation separately. It fuses the operations, optimizes memory access, and compiles the graph into a single NEFF artifact. The artifact runs as one hardware dispatch.

@torch.compile(backend='neuron', fullgraph=True)
def naive_compiled(t):
    x_sq = t * t
    mean_sq = x_sq.mean(dim=-1, keepdim=True)
    rms = torch.rsqrt(mean_sq + 1e-6)
    return t * rms * weight

The compiled function appears as one device execution:

Self NEURON total: 375.5 μs

The table shows one device execution, as in the torch.rms_norm profile, but the code follows a different path. Rather than calling a purpose-built fused operation, the compiler discovers the fusion from the original Python function. In this run, Neuron time is comparable to torch.rms_norm: 375 μs versus 403 μs.

Side-by-Side Summary

Both torch.rms_norm and torch.compile fuse RMSNorm into one kernel dispatch. In these runs, each approach takes about one-sixth of the naive implementation's device time. Fusion removes gaps between hardware executions and avoids materializing intermediate values such as x_sq and mean_sq between separately dispatched operations.

Conclusion

key_averages() identifies the operations consuming device time, while its full output also exposes CPU-side setup and transfers. Neuron Explorer provides the corresponding visual trace, allowing an operation to be followed from framework dispatch through runtime scheduling to hardware execution.

The same workflow can be applied to larger models: locate a hotspot in the table, verify it in the trace, and determine whether a fused operation or torch.compile can reduce its cost.

The next stage is device-level profiling. That analysis examines the NEFF, breaks down TensorEngine, VectorEngine, and direct memory access (DMA) activity, and places the kernel on a roofline to determine whether it is compute-bound or memory-bound.