Profiling TorchNeuron (Part 1): System-level profiling on AWS Trainium.
huggingface ·
This article introduces system-level profiling for TorchNeuron on AWS Trainium, aiming to help ML practitioners analyze inference performance bottlenecks. It details the setup of the NeuronConfig and NeuronProfiler to capture device and runtime traces using PyTorch's PrivateUse1 extension. The guide demonstrates how to use Neuron Explorer to visualize the lifecycle of operations from framework dispatch through hardware execution. Three implementations of RMSNorm are compared: naive PyTorch operations, the built-in torch.rms_norm, and a naive function compiled with torch.compile. Results show that both fusion via torch.rms_norm and compilation significantly reduce Neuron device time by eliminating intermediate tensor materialization and gaps between operations. The fused approaches achieve approximately one-sixth of the device time required by the naive implementation. The author concludes that identifying hotspots in key_averages tables and verifying them in traces is an effective workflow for optimization.
huggingface