TorchNeuron Profiling: How One Fix Cut Trainium Time by 6x

On AWS Trainium, a plain five-step PyTorch RMSNorm spent 2,355 microseconds on the chip. The fused version did the same maths in about 376. Amazon’s new guide to TorchNeuron profiling shows how to find gaps like that yourself, even if you have never opened a profiler.

The guide is Part 1 of a series by Josh Longenecker, published on Hugging Face on 21 October 2026. It covers system-level profiling. Below is what it teaches, why the numbers matter and what Part 2 will add.

Key takeaways

• TorchNeuron profiling runs through the standard torch.profiler API plus a Neuron config object. You don’t need a separate tool to capture a trace.
• On Trainium, a naive RMSNorm made 6 operator calls and used 2,355 µs of device time.
• Both torch.rms_norm (403 µs) and torch.compile (376 µs) brought that down to about one-sixth.
• Traces open in Neuron Explorer, AWS’s browser-based viewer, which replaced the older Neuron Profiler tools.
• Part 2 will go down to the hardware engines and a roofline analysis.

What is TorchNeuron profiling, and why should you care?

TorchNeuron is Amazon’s open-source native PyTorch backend for its Trainium AI chips. It supports normal eager execution and torch.compile. We covered the launch, and the bet against Nvidia that comes with it, in Amazon TorchNeuron: The Free Tool Aimed at Nvidia’s CUDA Moat.

Getting a framework to run on new hardware is one job. Making it fast is another. Profiling means recording what your code actually did and how long each step took. Without that, you are guessing.

A system-level profile follows one operation from end to end: PyTorch dispatches it, the Neuron Runtime schedules it, and a NeuronCore runs it. A device-level profile zooms into the chip itself. Part 1 deals only with the system view.

How do you capture a trace on Trainium?

If you have profiled on a GPU, the setup will look familiar. The guide wraps a short loop in torch.profiler.profile and adds a few Neuron-specific pieces:

1. A NeuronConfig with both ProfileMode.DEVICE and ProfileMode.RUNTIME turned on, passed in as experimental_config.
2. Two activities, CPU and PrivateUse1. PrivateUse1 is the slot PyTorch reserves for out-of-tree hardware backends, so it is where Neuron’s device events appear.
3. A schedule of wait=2, warmup=1, active=1, so only the fourth step is recorded and the first-run noise is skipped.
4. torch.neuron.synchronize() inside the loop. The author warns that without it, “the profiler can export before the device finishes.”
5. A NeuronProfiler that writes the trace, which you then upload to Neuron Explorer.

One cost is worth knowing about. According to AWS’s Neuron Explorer documentation, device mode reserves about 5 GB of high-bandwidth memory on Trn2. If you only need the system view, you can leave it off.

Where did the 2,355 microseconds go?

The test workload is small on purpose: RMSNorm, a normalisation layer used in most modern LLMs, run on random bf16 tensors of shape (4, 512). Written naively, it is a multiply, a mean, an add, a reciprocal square root and two more multiplies. That is six separate operator calls.

The profiler table makes the cost plain. aten::mul alone took 1,164.5 µs of device time, or 49.5% of the total. The aten::add call had a hidden extra: the epsilon value started on the CPU, so it caused a host-to-device copy.

Think of a chef who walks to the fridge once per ingredient. The cooking itself is quick. The trips add up. Each small op sends work to the chip, waits, and writes an intermediate tensor to memory that the next op has to read back.

Naive vs fused: the three versions compared

Version	Device time	Self CPU time	Operator calls
Naive PyTorch ops	2,355 µs	~125 µs	6
torch.rms_norm	403 µs	~48 µs	1
Naive code + torch.compile	376 µs	~98 µs	1

Both fused versions land at roughly one-sixth of the naive device time. The guide credits two things: fewer gaps between hardware executions, and no intermediate tensors such as x_sq and mean_sq making round trips through memory.

The torch.compile result is the useful one. Running torch.compile(backend='neuron', fullgraph=True) on the same naive code turned it into one compiled graph and one dispatch. You don’t have to know that a fused built-in exists. The compiler can do the fusing for you.

Our take: a 6x gap on a toy layer will not become a 6x speed-up for a whole model. It does show where to look first, which is anywhere your model makes lots of small, separate calls.

What does Neuron Explorer show you?

Neuron Explorer is now the only supported profiler for Trainium and Inferentia. AWS’s profiling tools page says the older neuron-profile and Neuron Profiler 2.0 reached end of support in Neuron 2.29. You can open Explorer in a browser over an SSH tunnel or through the VS Code extension.

In the system trace, each layer of the stack gets its own lane: framework/TID and framework/Stream for PyTorch, neuron_rt/NC for the runtime, and neuron_hw/NC for the hardware. Empty space between blocks on the hardware lane is time your chip spent idle.

What happens next

Part 2 will open the NEFF, the compiled Neuron executable, and read it alongside its NTFF hardware trace. The plan is a per-engine breakdown across the TensorEngine, VectorEngine and DMA (the units that move data), then a roofline. A roofline is a chart that tells you whether a kernel is limited by raw compute or by memory bandwidth.

That step matters more than it sounds. The system view tells you that time is being lost. The device view tells you why, and whether a custom kernel is worth writing. If you want to see how kernel-level tricks pay off on the GPU side, read our breakdown of PolyGraph CUDA Graphs.

FAQ

Do I need a special tool to profile TorchNeuron?

No. You capture with the standard torch.profiler plus Neuron’s NeuronConfig and NeuronProfiler, then view the results in Neuron Explorer.

Why was the naive RMSNorm so slow on Trainium?

It made six separate operator calls. Each one left gaps on the hardware and wrote intermediate tensors to memory, and one call also triggered a host-to-device copy.

Should I use torch.rms_norm or torch.compile?

In the guide’s test, both fused the work into one call: 403 µs and 376 µs respectively. torch.compile has the edge because it can fuse code you wrote yourself.

If you run models on Trainium, take one hot layer, wrap four steps in torch.profiler, and count the operator calls. That count tells you where to start optimising.