Tech · Hugging Face
Profiling TorchNeuron (Part 1): System-level profiling on AWS Trainium.
Compiled by KHAO Editorial — aggregated from 1 source. See llms.txt for citation guidance.
★ Tier-1 Source
Many ML practitioners have some intuition about how to make models run fast, but fewer know how to analyze why a model is slow.
Key facts
- The schedule uses wait=2, warmup=1, and active=1 across four prof.step calls
- In this run, the Neuron time is comparable to torch.rms_norm, at 375 μs compared with 403 μs
- See Step 5: Upload your profile for the UI and command-line instructions
- In Part 2, they will move into the device-level profile
Summary
This post starts from scratch and assumes no prior profiling experience. Additionally, TorchNeuron integrates with Neuron Kernel Interface (NKI) to enable custom operators and performance optimization for Trainium hardware. How to configure torch.profiler with the Neuron backend and capture a clean trace. How to read the key_averages table to identify which operations consume Neuron device time and compare their Self CPU time. How to follow one operation through Neuron Explorer, from framework dispatch through runtime scheduling to hardware execution on a NeuronCore.