Profilage de TorchNeuron (Partie 1) : Profilage système sur AWS Trainium
huggingface ·
Cet article présente le profilage au niveau du système pour TorchNeuron sur AWS Trainium, dans le but d'aider les praticiens du machine learning à analyser les goulots d'étranglement des performances d'inférence. Il détaille la configuration de NeuronConfig et NeuronProfiler pour capturer les traces matérielles et d'exécution à l'aide de l'extension PrivateUse1 de PyTorch. Le guide démontre comment utiliser Neuron Explorer pour visualiser le cycle de vie des opérations, depuis leur dispatchement par le framework jusqu'à leur exécution matérielle. Trois implémentations de RMSNorm sont comparées : des opérations PyTorch naïves, la fonction intégrée torch.rms_norm, et une fonction naïve compilée avec torch.compile. Les résultats montrent que la fusion via torch.rms_norm et la compilation réduisent considérablement le temps d'exécution sur le dispositif Neuron en éliminant la matérialisation intermédiaire des tenseurs et les écarts entre les opérations. Les approches fusionnées nécessitent environ un sixième du temps matériel requis par l'implémentation naïve. L'auteur conclut que l'identification des points chauds dans les tables key_averages et leur vérification dans les traces constitue un flux de travail efficace pour l'optimisation.
huggingface