Profiling PyTorch models on AWS Trainium with TorchNeuron

AWS Trainium is a purpose-built chip for deep learning, but optimizing models for it requires specific profiling techniques. A new tutorial from Hugging Face demonstrates how to use TorchNeuron, the native PyTorch backend for Trainium, to analyze inference performance systematically. The guide walks engineers through capturing traces, reading system-level data, and comparing different implementation strategies for a common operation.

What happened

Many machine learning practitioners rely on intuition to speed up models, but few possess the tools to diagnose exactly why a model is slow. To address this gap, Hugging Face published the first part of a series focused on system-level profiling for AWS Trainium. The tutorial assumes no prior experience with profiling on this hardware and starts from scratch, using RMSNorm as a case study.

The walkthrough utilizes TorchNeuron, an open-source backend that allows standard PyTorch workflows to run on Trainium with minimal code changes. It supports eager execution, torch.compile, and custom operators via the Neuron Kernel Interface. The goal is to teach engineers how to configure the profiler, capture clean traces, and interpret the resulting data to identify time-consuming operations.

By the end of the guide, readers learn to correlate CPU time with device time, follow operations through the software stack, and understand how different coding patterns affect hardware utilization. This foundational knowledge prepares developers to optimize larger models by identifying hotspots and applying appropriate fusion or compilation techniques.

How it works

Profiling on Trainium relies on two main components: NeuronConfig and NeuronProfiler. The NeuronConfig object controls which data streams to capture, such as device instructions and runtime traces, and specifies the output directory. It is passed to torch.profiler.profile as an experimental configuration. The NeuronProfiler uses this config to export both the standard framework trace and Neuron-specific system traces.

To visualize the data, engineers use Neuron Explorer, a browser-based viewer accessible via SSH tunnel or VS Code extension. This tool displays the system trace, device timeline, and per-engine breakdown. It matches compiled executables (NEFF files) with their hardware instruction traces (NTFF files) using a shared numeric hash. This allows users to see exactly how high-level PyTorch operations map to low-level hardware execution.

The profiling process involves synchronizing the device within the profiler loop to ensure accurate timing. Without calling torch.neuron.synchronize(), the profiler might export data before the device finishes its work, leading to empty or incorrect time columns. The resulting traces show the lifecycle of each operation from framework dispatch to runtime scheduling and finally to hardware execution on a NeuronCore.

Key details

• TorchNeuron integrates with PyTorch using the PrivateUse1 extension point, avoiding the need for a dedicated dispatch key like CUDA.
• The NeuronConfig parameter correlate_device_time=True propagates device time into the key_averages() table for direct comparison with CPU time.
• Naive implementations of RMSNorm using individual PyTorch operations result in multiple separate hardware dispatches and intermediate tensor materialization.
• Using torch.rms_norm() lowers the operation to a single fused Neuron kernel, reducing device time significantly by eliminating gaps between operations.
• torch.compile with the Neuron backend can automatically fuse naive operations into a single graph, achieving performance similar to manual fusion without changing the original Python code structure.
• Neuron Explorer organizes events into lanes for framework dispatch, device interaction, runtime API calls, and hardware execution, allowing detailed dependency chain analysis.

Why it matters

For engineers building products on AWS Trainium, understanding the difference between CPU overhead and actual device computation is critical for cost and latency optimization. Naive code often looks efficient in Python but generates excessive overhead on the accelerator due to frequent dispatches and memory transfers. By visualizing these gaps, developers can make informed decisions about when to use built-in fused operations versus relying on compiler optimizations.

The ability to trace an operation from the PyTorch framework down to the NeuronCore hardware provides a level of transparency that is often missing in black-box accelerators. This visibility helps teams debug performance regressions and validate that their optimization efforts are actually reaching the hardware as intended. It shifts optimization from guesswork to evidence-based engineering.

Furthermore, as models grow in complexity, the cumulative effect of small inefficiencies becomes substantial. Identifying and fusing operations early in the development cycle prevents technical debt related to performance. The techniques described here provide a repeatable workflow for auditing model efficiency, ensuring that infrastructure spend aligns with computational needs.

What you can do

• Install the torch_neuronx package and configure NeuronConfig with ProfileMode.DEVICE and ProfileMode.RUNTIME to capture comprehensive traces.
• Use torch.profiler.schedule with warmup steps to exclude initialization overhead from your performance measurements.
• Call torch.neuron.synchronize() inside your profiling loop to ensure device tasks complete before trace export.
• Analyze the key_averages() table sorted by self_privateuse1_time_total to identify the most expensive operations on the Neuron device.
• Upload the generated .neff and .ntff files to Neuron Explorer to visualize the execution timeline and verify operator fusion.
• Compare naive implementations against torch.rms_norm or torch.compile variants to quantify the performance benefits of operation fusion.