Profilage des modèles PyTorch sur AWS Trainium avec TorchNeuron

AWS Trainium est une puce dédiée à l'apprentissage profond, mais l'optimisation des modèles pour cette plateforme nécessite des techniques de profilage spécifiques. Un nouveau tutoriel de Hugging Face montre comment utiliser TorchNeuron, le backend PyTorch natif pour Trainium, afin d'analyser systématiquement les performances d'inférence. Le guide accompagne les ingénieurs dans la capture de traces, la lecture des données au niveau système et la comparaison de différentes stratégies d'implémentation pour une opération courante.

Ce qui s'est passé

De nombreux praticiens du machine learning se fient à leur intuition pour accélérer leurs modèles, mais peu disposent des outils permettant de diagnostiquer précisément pourquoi un modèle est lent. Pour combler ce manque, Hugging Face a publié la première partie d'une série consacrée au profilage au niveau système pour AWS Trainium. Le tutoriel ne suppose aucune expérience préalable du profilage sur ce matériel et commence par les bases, en utilisant RMSNorm comme étude de cas.

La démonstration utilise TorchNeuron, un backend open source qui permet aux workflows PyTorch standards de fonctionner sur Trainium avec des modifications minimales du code. Il prend en charge l'exécution eager, torch.compile et les opérateurs personnalisés via l'interface Neuron Kernel Interface. L'objectif est d'enseigner aux ingénieurs comment configurer le profiler, capturer des traces propres et interpréter les données résultantes pour identifier les opérations chronophages.

À la fin du guide, les lecteurs apprennent à corréler le temps CPU avec le temps device, à suivre les opérations à travers la pile logicielle et à comprendre comment différents patterns de codage affectent l'utilisation du matériel. Ces connaissances fondamentales préparent les développeurs à optimiser des modèles plus vastes en identifiant les points chauds et en appliquant des techniques appropriées de fusion ou de compilation.

Comment cela fonctionne

Le profilage sur Trainium repose sur deux composants principaux : NeuronConfig et NeuronProfiler. L'objet NeuronConfig contrôle quels flux de données capturer, tels que les instructions du device et les traces runtime, et spécifie le répertoire de sortie. Il est transmis à torch.profiler.profile en tant que configuration expérimentale. Le NeuronProfiler utilise cette configuration pour exporter à la fois la trace standard du framework et les traces système spécifiques à Neuron.

Pour visualiser les données, les ingénieurs utilisent Neuron Explorer, une visionneuse basée sur le navigateur accessible via un tunnel SSH ou une extension VS Code. Cet outil affiche la trace système, la timeline du device et la décomposition par moteur. Il associe les exécutables compilés (fichiers NEFF) à leurs traces d'instructions matérielles (fichiers NTFF) grâce à un hash numérique partagé. Cela permet aux utilisateurs de voir exactement comment les opérations PyTorch de haut niveau correspondent à l'exécution matérielle de bas niveau.

Le processus de profilage implique la synchronisation du device dans la boucle du profiler pour garantir une mesure précise du temps. Sans appel à torch.neuron.synchronize(), le profiler pourrait exporter les données avant que le device n'ait terminé son travail, entraînant des colonnes de temps vides ou incorrectes. Les traces résultantes montrent le cycle de vie de chaque opération, depuis la dispatch du framework jusqu'à la planification runtime et enfin à l'exécution matérielle sur un NeuronCore.

Détails clés

• TorchNeuron s'intègre à PyTorch en utilisant le point d'extension PrivateUse1, évitant ainsi la nécessité d'une clé de dispatch dédiée comme CUDA.
• Le paramètre correlate_device_time=True de NeuronConfig propage le temps device dans la table key_averages() pour une comparaison directe avec le temps CPU.
• Les implémentations naïves de RMSNorm utilisant des opérations PyTorch individuelles entraînent plusieurs dispatches matériels séparés et la matérialisation de tenseurs intermédiaires.
• L'utilisation de torch.rms_norm() abaisse l'opération à un seul kernel Neuron fusionné, réduisant considérablement le temps device en éliminant les lacunes entre les opérations.
• torch.compile avec le backend Neuron peut automatiquement fusionner les opérations naïves en un seul graphe, atteignant des performances similaires à la fusion manuelle sans modifier la structure du code Python original.
• Neuron Explorer organise les événements en pistes pour la dispatch du framework, l'interaction avec le device, les appels API runtime et l'exécution matérielle, permettant une analyse détaillée des chaînes de dépendances.

Pourquoi c'est important

Pour les ingénieurs développant des produits sur AWS Trainium, comprendre la différence entre la surcharge CPU et le calcul réel sur le device est crucial pour l'optimisation des coûts et de la latence. Le code naïf semble souvent efficace en Python mais génère une surcharge excessive sur l'accélérateur en raison de dispatches fréquents et de transferts mémoire. En visualisant ces lacunes, les développeurs peuvent prendre des décisions éclairées quant à l'utilisation d'opérations fusionnées intégrées plutôt qu'en comptant sur les optimisations du compilateur.

La capacité de tracer une opération depuis le framework PyTorch jusqu'au hardware NeuronCore offre un niveau de transparence souvent absent des accélérateurs boîte noire. Cette visibilité aide les équipes à déboguer les régressions de performance et à valider que leurs efforts d'optimisation atteignent réellement le hardware comme prévu. Elle fait passer l'optimisation du tâtonnement à l'ingénierie basée sur des preuves.

De plus, à mesure que les modèles gagnent en complexité, l'effet cumulatif des petites inefficacités devient substantiel. Identifier et fusionner les opérations tôt dans le cycle de développement prévient la dette technique liée à la performance. Les techniques décrites ici fournissent un workflow reproductible pour auditer l'efficacité des modèles, garantissant que les dépenses d'infrastructure sont alignées avec les besoins computationnels.

Ce que vous pouvez faire

• Installez le package torch_neuronx et configurez NeuronConfig avec ProfileMode.DEVICE et ProfileMode.RUNTIME pour capturer des traces complètes.
• Utilisez torch.profiler.schedule avec des étapes de warmup pour exclure la surcharge d'initialisation de vos mesures de performance.
• Appelez torch.neuron.synchronize() dans votre boucle de profilage pour vous assurer que les tâches du device sont terminées avant l'exportation de la trace.
• Analysez la table key_averages() triée par self_privateuse1_time_total pour identifier les opérations les plus coûteuses sur le device Neuron.
• Téléversez les fichiers .neff et .ntff générés dans Neuron Explorer pour visualiser la timeline d'exécution et vérifier la fusion des opérateurs.
• Comparez les implémentations naïves avec les variantes torch.rms_norm ou torch.compile pour quantifier les bénéfices de performance de la fusion des opérations.