In a post published on the Hugging Face platform, Amazon launches a series dedicated to analyzing inference performance, that is, the execution of a pre-trained model. This first part explains how to identify where time is spent, using profiling tools applied to AWS’s Trainium chip, designed for deep learning model training and inference.
The guide is aimed at readers with no prior profiling experience. It relies on TorchNeuron, the native open-source PyTorch backend for Trainium, which supports standard PyTorch workflows, including eager execution and torch.compile, with minimal code changes. TorchNeuron also integrates with the Neuron Kernel Interface (NKI), which allows for creating custom operators.
Setting up the profiler
Profiling relies on two objects. NeuronConfig defines the data streams to capture, the NeuronCores to trace, and the output folder, and is then passed to torch.profiler.profile. NeuronProfiler provides an export function that writes the framework trace, in trace.json format, alongside Neuron’s system traces. The parameter correlate_device_time=True allows displaying the time spent on the device in the key_averages() table, where it can be compared to CPU time.
The traces can then be opened in Neuron Explorer, a visualization tool that runs in the browser. It allows tracking an operation from its dispatch in the framework to its execution on the hardware, via scheduling in the runtime. For a NeuronCore, i.e., a compute core of the chip, the NEFF file contains the compiled executable, while the corresponding NTFF file contains the trace of hardware instructions.
Three versions of the same operation
The post examines the RMSNorm function, a simple operation composed of a series of arithmetic calculations and reductions. Three implementations are compared: a version written operation by operation in PyTorch, the call to the built-in function torch.rms_norm(), and the version written operation by operation but compiled with torch.compile.
In the naive version, the table shows that three calls to aten::mul account for the largest share of time spent on the device. The total Neuron time reaches 2,355 µs for six operator calls. According to the author, the built-in function torch.rms_norm(), which is transformed into a single fused compute kernel, occupies only one operation, aten::_fused_rms_norm, with a Neuron time of 403.3 µs. The version compiled with torch.compile appears as a single compiled region, lasting 375.5 µs.
The post specifies that, in these tests, both approaches take about one-sixth of the compute time of the naive version. Fusion removes interruptions between hardware executions and avoids materializing intermediate values, such as x_sq and mean_sq, between separately dispatched operations.
The author indicates that the method can be applied to larger models: identify a hotspot in the table, verify it in the trace, then determine if a fused operation or torch.compile can reduce its cost. The second part of the series will focus on device-level profiling, with analysis of the NEFF file.
