Profiling TorchNeuron Part 1: System-level Inference Analysis on AWS Trainium
Many ML practitioners understand how to accelerate models, but analyzing why a model is slow requires systematic profiling. This guide walks through how to analyze inference performance on AWS Trainium using TorchNeuron, providing a comprehensive method to identify time consumption hotspots and demonstrate the performance benefits of kernel fusion.
Neuron Integration as a PyTorch Backend
The Neuron backend integrates with PyTorch using the PrivateUse1 extension point. This mechanism allows an accelerator to integrate with the framework without requiring a dedicated dispatch key, such as cuda. The torch_neuronx package registers Neuron as a native backend for PyTorch code, leveraging PyTorch's advanced extension capabilities.
Setting Up the Profiler
To profile a Neuron workload, two objects must be configured: NeuronConfig and NeuronProfiler. The NeuronConfig controls which Neuron data streams and NeuronCores are captured and specifies the output location. The NeuronProfiler takes this configuration and provides an export_trace callback to write the framework trace (trace.json) alongside the Neuron system traces.
Key configuration parameters include:
modes: Selects the data streams to capture, such as device instructions and runtime traces.capture_enabled_for_nc: Selects the specific NeuronCores to trace.correlate_device_time=True: This setting propagates device time from the trace into thekey_averages()table, enabling direct comparison between CPU and Neuron device time.
Tooling: Neuron Explorer
Neuron Explorer is a browser-based viewer designed to visualize the system trace, device timeline, and per-engine breakdown. It allows users to confirm numerical data from key_averages() and inspect the detailed lifecycle of operations.
After profiling, the output directory contains the system trace, along with matched .neff (Neuron Executable File Format) and .ntff (hardware instruction trace) files if device data is captured. To view the profile in Neuron Explorer, the complete output directory must be uploaded for a system or combined profile, or the matching .neff and .ntff files must be uploaded individually for a device-only profile.
Case Study: Profiling RMSNorm
RMSNorm was used for this analysis because it is a small operation with a clear sequence of arithmetic and reduction operations. Three implementations of RMSNorm were profiled to compare performance:
- Naive RMSNorm: Implemented using individual PyTorch operations (multiply, mean, add, reciprocal square root, multiply).
torch.rms_norm(): The PyTorch built-in function, which is lowered to a single fused Neuron kernel.- Naive RMSNorm +
torch.compile: The original naive function compiled usingtorch.compilewith the Neuron backend.
All scripts used identical inputs, profiler configurations, and scheduling to ensure direct comparability of the results.
Profile 1: Naive PyTorch Operations
When RMSNorm is implemented using individual PyTorch operations, each line maps to one or more ATen operations and can create intermediate tensors. The profiling results show that three aten::mul calls account for the largest share of Neuron device time.
| Name | # Calls | Self CPU | Self NEURON | % NEURON |
|---|---|---|---|---|
| aten::mul | 3 | 60.9 μs | 1,164.5 μs | 49.5% |
| aten::add | 1 | 26.6 μs | 434.6 μs | 18.5% |
| aten::rsqrt | 1 | 15.6 μs | 415.0 μs | 17.6% |
| aten::mean | 1 | 21.1 μs | 340.8 μs | 14.5% |
Self NEURON total: 2.355 ms
In this naive implementation, each operation crosses the dispatch boundary independently. The aten::add operation, for example, triggers a host-to-device copy because the scalar epsilon starts on the CPU.
Profile 2: torch.rms_norm()
By replacing individual operations with PyTorch's built-in rms_norm function, Neuron lowers the entire operation into one fused kernel. The key averages table shows this operation as a single device operation.
| Name | # Calls | Self CPU | Self NEURON | % NEURON |
|---|---|---|---|---|
| aten::_fused_rms_norm | 1 | 47.6 μs | 403.3 μs | 100% |
Self NEURON total: 403.3 μs
This fused implementation follows a single framework-to-hardware path, avoiding the overhead of dispatching each arithmetic operation separately.
Profile 3: Naive PyTorch with torch.compile
When the original naive RMSNorm function is compiled using @torch.compile(backend='neuron', fullgraph=True), TorchDynamo captures the Python bytecode as an FX graph. The Neuron backend receives this complete graph, which it fuses and compiles into a single NEFF artifact. This artifact runs as one hardware dispatch.
| Name | # Calls | Self CPU | Self NEURON | % NEURON |
|---|---|---|---|---|
| Torch-Compiled Region: 0/0 | 1 | 98.4 μs | 375.5 μs | 100% |
Self NEURON total: 375.5 μs
This compiled graph also follows a single path through runtime scheduling and hardware execution. The Neuron time achieved here (375 μs) is comparable to the purpose-built torch.rms_norm (403 μs).
Side-by-Side Performance Summary
| Variant | NEURON time | Self CPU time | Operator calls |
|---|---|---|---|
| Naive ops | 2,355 µs | ~125 µs | 6 |
| torch.rms_norm | 403 µs | ~48 µs | 1 |
| Naive + torch.compile | 376 µs | ~98 µs | 1 |
Both torch.rms_norm and torch.compile successfully fuse RMSNorm into one kernel dispatch. These fusion methods reduce the device time to approximately one-sixth of the naive implementation's device time. Fusion eliminates the gaps between hardware executions and prevents the materialization of intermediate values, such as x_sq and mean_sq, between separately dispatched operations.
Conclusion
The key_averages() table identifies which operations consume device time, while the full output provides context on CPU-side setup and transfers. Neuron Explorer provides the corresponding visual trace, allowing users to follow an operation from framework dispatch through runtime scheduling to hardware execution and observe the gaps that fusion removes.
This workflow can be applied to larger models: identify a performance hotspot in the table, verify it visually in the trace, and then determine if a fused operation or torch.compile can reduce its cost. Future analysis will move to device-level profiling, inspecting the TensorEngine, VectorEngine, and DMA breakdown to determine if kernels are compute-bound or memory-bound.