Profilage de TorchNeuron (Partie 1) : Profilage système sur AWS Trainium
huggingface ·
Cet article présente une approche systématique du profilage des performances d'inférence sur AWS Trainium à l'aide de TorchNeuron, un backend PyTorch open source. Il guide les lecteurs dans la configuration de NeuronConfig et NeuronProfiler afin de capturer les traces matérielles et d'exécution pour analyse. L'auteur illustre ces outils en profilant trois variantes de l'opération RMSNorm : des opérations PyTorch naïves, torch.rms_norm intégré, et une version compilée avec torch.compile. Les résultats montrent que le noyau fusionné et le graphe compilé réduisent significativement le temps d'exécution sur le dispositif Neuron par rapport à l'implémentation naïve, en éliminant la matérialisation des tenseurs intermédiaires. Le guide explique comment interpréter la table key_averages() pour identifier les goulots d'étranglement de performance et utilise Neuron Explorer pour visualiser le cycle d'exécution, depuis la dispatch du framework jusqu'à l'exécution matérielle. Ce flux de travail permet aux praticiens de localiser les points chauds et d'optimiser les modèles en tirant parti des techniques de fusion ou de compilation. L'article se conclut par une annonce de la deuxième partie, qui se concentrera sur le profilage au niveau du dispositif et l'analyse roofline.
huggingface