---
title: "Profiling TorchNeuron (Part 1): System-level profiling on AWS Trainium."
description: "The post introduces system-level profiling of inference performance on AWS Trainium using TorchNeuron. It explains how to configure torch.profiler with the Neuron backend by passing a NeuronConfig and using NeuronProfiler to capture traces. The walkthrough profiles three implementations of RMSNorm: naive PyTorch operations, the fused torch.rms_norm, and naive operations compiled with torch.compile. The key_averages() table shows that the naive implementation spends 2,355 µs on the Neuron device across six operator dispatches, while both fused and compiled versions reduce device time to about 400 µs with a single dispatch. Neuron Explorer visually confirms the lifecycle of each operation from framework dispatch through runtime scheduling to hardware execution, revealing the gaps that fusion eliminates."
---

# Profiling TorchNeuron (Part 1): System-level profiling on AWS Trainium.

[Hugging Face](https://yomu.fyi/company/hugging-face) · Josh Longenecker · Oct 21, 2026

**Type:** Tutorial

## Summary

The post introduces system-level profiling of inference performance on AWS Trainium using TorchNeuron. It explains how to configure torch.profiler with the Neuron backend by passing a NeuronConfig and using NeuronProfiler to capture traces. The walkthrough profiles three implementations of RMSNorm: naive PyTorch operations, the fused torch.rms\_norm, and naive operations compiled with torch.compile. The key\_averages() table shows that the naive implementation spends 2,355 µs on the Neuron device across six operator dispatches, while both fused and compiled versions reduce device time to about 400 µs with a single dispatch. Neuron Explorer visually confirms the lifecycle of each operation from framework dispatch through runtime scheduling to hardware execution, revealing the gaps that fusion eliminates.

## Context

Many ML practitioners lack systematic methods to analyze inference performance bottlenecks on custom hardware. This post starts from scratch and assumes no prior profiling experience.

## Approach / What changed

Use torch.profiler with a NeuronConfig to capture CPU and Neuron device traces, then analyze the key\_averages() table and inspect the trace in Neuron Explorer to identify performance bottlenecks.

## Takeaways

- The naive RMSNorm implementation with individual PyTorch operations dispatches six separate ATen operations, each crossing the framework-to-hardware boundary independently, resulting in 2,355 µs of Neuron device time.
- Both torch.rms\_norm and torch.compile fuse the RMSNorm computation into a single kernel dispatch, reducing Neuron device time to approximately 400 µs by eliminating gaps between hardware executions and avoiding materialization of intermediate tensors.
- Neuron Explorer’s dependency chain viewer allows following a single operation from framework dispatch through runtime scheduling to hardware execution, confirming the lifecycle and gaps that fusion removes.

## FAQ

### How do I set up the profiler for TorchNeuron on Trainium?

Create a NeuronConfig with desired modes and output directory, then pass it as experimental\_config to torch.profiler.profile. Use NeuronProfiler to export the trace. Include ProfilerActivity.PrivateUse1 and call torch.neuron.synchronize() inside the profiler loop to ensure device completion.

### What does the key\_averages() table show for a Neuron workload?

The table lists ATen operations with their call count, Self CPU time, and Self NEURON time. Sorting by self\_privateuse1\_time\_total reveals which operations consume the most Neuron device time, allowing comparison of CPU and device costs.

### How does torch.compile affect profiling compared to naive operations?

With torch.compile, TorchDynamo captures the entire graph and the Neuron backend fuses operations into a single NEFF artifact, appearing as one Torch-Compiled Region in the table. This reduces device time from multiple dispatches to one, similar to a purpose-built fused operation.

**Tags:** [Neuron Explorer](https://yomu.fyi/topic/neuron-explorer), [Profiling](https://yomu.fyi/topic/profiling), [PyTorch](https://yomu.fyi/topic/pytorch), [TorchNeuron](https://yomu.fyi/topic/torchneuron), [Trainium](https://yomu.fyi/topic/trainium)

- Source: [Hugging Face](https://huggingface.co/blog/amazon/profiling-torchneuron-part-1)
- Source URL: https://huggingface.co/blog/amazon/profiling-torchneuron-part-1
- Ingested by Yomu: 2026-10-21T18:00:14.148Z

[Read original post](https://huggingface.co/blog/amazon/profiling-torchneuron-part-1)