AWS Shows Trainium Profiling Workflow

In a Hugging Face post published October 21, 2026, AWS author Josh Longenecker demonstrates how to profile RMSNorm inference on AWS Trainium with TorchNeuron, AWS’s open-source PyTorch backend. The Part 1 walkthrough starts from no profiling experience and shows how to capture a trace, find operations consuming Neuron device time, and follow an operation from framework dispatch through runtime scheduling to hardware execution. TorchNeuron supports standard PyTorch workflows, including eager execution, torch.compile, and PyTorch APIs, with minimal code changes; it also connects to the Neuron Kernel Interface for custom operators and optimization.

The setup uses torch.profiler with NeuronConfig to select device and runtime data, NeuronCores, and an output directory. NeuronProfiler exports the framework trace alongside system traces. Setting correlate_device_time=True adds device time to the key_averages() table for comparison with Self CPU time. The example includes ProfilerActivity.PrivateUse1 and synchronizes the device inside the profiling loop so results are ready before export. The schedule uses wait=2, warmup=1, and active=1 across four prof.step() calls, excluding initial setup overhead from the recorded warm iteration. Neuron uses PyTorch’s PrivateUse1 extension point to register as a backend without a dedicated dispatch key such as cuda.

Longenecker compares three RMSNorm implementations with the same inputs and profiler settings: separate PyTorch operations, the built-in torch.rms_norm(), and the original operations compiled with torch.compile. In the naive version, six operator calls take 2,355 µs of Neuron device time and about 125 µs of Self CPU time. Its three aten::mul calls account for 1,164.5 µs, or 49.5% of device time; aten::add takes 434.6 µs (18.5%), aten::rsqrt 415.0 µs (17.6%), and aten::mean 340.8 µs (14.5%). The add also triggers a host-to-device copy because its scalar begins on the CPU.

The built-in torch.rms_norm() lowers the computation to one fused kernel: one operator call, 403 µs of Neuron time, and about 48 µs of Self CPU time. Compiling the naive function with @torch.compile(backend="neuron") also fuses it into one dispatch, reporting 376 µs of device time and about 98 µs of Self CPU time. The post says both fused approaches take about one-sixth of the naive version’s device time in these runs. Fusion removes gaps between hardware executions and avoids intermediate values such as x_sq and mean_sq. Neuron Explorer displays framework, runtime, and hardware lanes, letting users follow an operation through dispatch, scheduling, and execution. The post says Part 2 will examine TensorEngine, VectorEngine, DMA, and whether a kernel is compute-bound or memory-bound.

In a Hugging Face post published October 21, 2026, AWS author Josh Longenecker demonstrates how to profile RMSNorm inference on AWS Trainium with TorchNeuron, AWS’s open-source PyTorch backend. The Part 1 walkthrough starts from no profiling experience and shows how to capture a trace, find operations consuming Neuron device time, and follow an operation from framework dispatch through runtime scheduling to hardware execution. TorchNeuron supports standard PyTorch workflows, including eager execution, torch.compile, and PyTorch APIs, with minimal code changes; it also connects to the Neuron Kernel Interface for custom operators and optimization.

The setup uses torch.profiler with NeuronConfig to select device and runtime data, NeuronCores, and an output directory. NeuronProfiler exports the framework trace alongside system traces. Setting correlate_device_time=True adds device time to the key_averages() table for comparison with Self CPU time. The example includes ProfilerActivity.PrivateUse1 and synchronizes the device inside the profiling loop so results are ready before export. The schedule uses wait=2, warmup=1, and active=1 across four prof.step() calls, excluding initial setup overhead from the recorded warm iteration. Neuron uses PyTorch’s PrivateUse1 extension point to register as a backend without a dedicated dispatch key such as cuda.

Longenecker compares three RMSNorm implementations with the same inputs and profiler settings: separate PyTorch operations, the built-in torch.rms_norm(), and the original operations compiled with torch.compile. In the naive version, six operator calls take 2,355 µs of Neuron device time and about 125 µs of Self CPU time. Its three aten::mul calls account for 1,164.5 µs, or 49.5% of device time; aten::add takes 434.6 µs (18.5%), aten::rsqrt 415.0 µs (17.6%), and aten::mean 340.8 µs (14.5%). The add also triggers a host-to-device copy because its scalar begins on the CPU.

The built-in torch.rms_norm() lowers the computation to one fused kernel: one operator call, 403 µs of Neuron time, and about 48 µs of Self CPU time. Compiling the naive function with @torch.compile(backend="neuron") also fuses it into one dispatch, reporting 376 µs of device time and about 98 µs of Self CPU time. The post says both fused approaches take about one-sixth of the naive version’s device time in these runs. Fusion removes gaps between hardware executions and avoids intermediate values such as x_sq and mean_sq. Neuron Explorer displays framework, runtime, and hardware lanes, letting users follow an operation through dispatch, scheduling, and execution. The post says Part 2 will examine TensorEngine, VectorEngine, DMA, and whether a kernel is compute-bound or memory-bound.