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

Hugging Face has started a profiling series that does something a lot of ML writeups skip: it shows how to explain slowness, not just how to make numbers look better. Part 1 begins from zero and walks through inference profiling on AWS Trainium using TorchNeuron, the open source PyTorch backend for that hardware. The goal is simple: figure out where the time goes, then prove it in a trace.

TorchNeuron fits into normal PyTorch code through the PrivateUse1 extension point, so it does not need a special dispatch key like cuda. The post also leans on NeuronConfig and NeuronProfiler to capture system traces, device data, and runtime events. With correlate_device_time turned on, the device time from the trace shows up in key_averages(), which makes it easier to compare CPU work with Neuron work in the same table.

The example workload is RMSNorm, and the post profiles three versions of it. The first uses separate PyTorch ops: multiply, mean, add, rsqrt, and another multiply. That version produces six operator calls and 2,355 microseconds of Neuron time, with aten::mul alone showing up three times. The same function also shows a little bit of CPU setup time, but the heavier cost is clearly on the device side.

The second version swaps in torch.rms_norm(), which TorchNeuron lowers to a single fused kernel. That drops the device time to 403 microseconds. The third keeps the original Python math but wraps it in torch.compile(backend='neuron', fullgraph=True). That path lands at 375 microseconds of Neuron time, also as one hardware dispatch. In both fused cases, the intermediate values vanish as separate device trips.

The post is as much about reading the trace as reading the table. In Neuron Explorer, each operation moves from framework dispatch to runtime scheduling and then to hardware execution on a NeuronCore. That makes it possible to see the gaps that fusion removes. Part 2 will go deeper, into the device-level profile, with NEFF, TensorEngine, VectorEngine, DMA, and a roofline view to decide whether a kernel is compute-bound or memory-bound.