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

In a Hugging Face enterprise article published October 21, 2026, Josh Longenecker of AWS opens a series on analyzing why models run slowly on AWS Trainium, Amazon's purpose-built deep learning chip. The tutorial, aimed at readers with no profiling experience, uses TorchNeuron, the open source native PyTorch backend for Trainium, which plugs into PyTorch through its PrivateUse1 extension point.

Longenecker configures torch.profiler with Neuron-specific settings, reads per-operation device time from the key_averages() table, and uses the browser-based Neuron Explorer to trace an operation from framework dispatch through runtime scheduling to hardware execution. He warns that omitting torch.neuron.synchronize() inside the profiler loop can leave the Neuron time column empty.

The test case is RMSNorm, implemented three ways with identical inputs and profiler settings. Written as individual PyTorch operations, it took 2,355 microseconds of device time across six operator calls, with three multiplications accounting for about half. The built-in torch.rms_norm, lowered to one fused kernel, took 403 microseconds, while the original naive code wrapped in torch.compile took 376 microseconds, also as a single dispatch.

The author says either route cuts device time to roughly one-sixth of the naive version by removing gaps between hardware executions and avoiding intermediate values. The compiled approach used about 98 microseconds of CPU time versus 48 for the built-in. He suggests the same workflow for larger models: find a hotspot, confirm it in the trace, then test fusion or compilation. Part 2 will cover device-level profiling, including whether a kernel is compute-bound or memory-bound.