Profiling TorchNeuron (Part 1): System-level profiling on AWS Trainium.
Hugging Face details how to profile PyTorch models on AWS Trainium using TorchNeuron, explaining system-level tracing and performance analysis workflows.
The post introduces systematic profiling for PyTorch models running on AWS Trainium via TorchNeuron, an open source PyTorch backend designed for Trainium hardware. It explains how TorchNeuron integrates with PyTorch using the PrivateUse1 extension point, allowing standard workflows like eager execution and torch.compile with minimal code changes. The Neuron backend leverages the Neuron Kernel Interface to support custom operators and hardware-specific optimizations, enabling deeper performance tuning beyond default PyTorch behavior.
Profiling on Trainium involves configuring NeuronConfig with parameters such as modes, capture_enabled_for_nc, and correlate_device_time to capture device instructions and runtime traces. These settings generate output files including system traces, NEFF compiled executables, and NTFF hardware instruction traces. The Neuron Explorer tool visualizes these traces through an SSH tunnel or VS Code extension, helping users inspect per-engine breakdowns and confirm numerical results from key_averages().
The tutorial profiles RMSNorm, a small PyTorch operation, using a consistent schedule with wait=2, warmup=1, and active=1 across four prof.step() calls. Synchronization with torch.neuron.synchronize() inside the profiler loop ensures device operations complete before exporting data, preventing empty Neuron time columns. The resulting key_averages() table ranks operations by self_privateuse1_time_total, revealing aten::mul as the dominant Neuron device cost due to three multiplication operations in RMSNorm.
Neuron Explorer’s dependency chain viewer traces individual operations like aten::mul from framework dispatch through runtime scheduling to hardware execution on Neuron cores. The system trace organizes events into lanes, showing the lifecycle of each operation from creation in the framework to execution on device hardware. Fused operations such as RMSNorm appear as single entries in the key averages table, reflecting hardware-level fusion that reduces dispatch overhead compared to separate arithmetic operations.