🚀 Profilage AWS Trainium : torch.profiler & Neuron
Découvrez le profilage AWS Trainium avec torch.profiler et Neuron Explorer : identifiez les ops coûteuses, mesurez l’impact et optimisez PyTorch avec TorchNeuron.
Profilage PyTorch sur AWS Trainium : trouver ce qui ralentit vraiment vos modèles
Un modèle qui répond lentement n'a pas forcément besoin d'un modèle plus petit ou d'un accélérateur plus puissant. Avant de changer quoi que ce soit, il faut comprendre où part le temps : le calcul sur l'accélérateur, les appels exécutés par PyTorch, les transferts entre processeur et puce, ou encore les temps morts entre deux exécutions. Sans mesures, on risque surtout d'optimiser la mauvaise pièce.
L'idée de ce guide est simple : profiler une opération PyTorch sur AWS Trainium avec TorchNeuron, puis relier ce qu'on voit dans les traces à ce qui se passe réellement sur le matériel. On s'appuie sur RMSNorm, une opération fréquente dans les modèles de langage, déclinée en trois versions : plusieurs opérations PyTorch distinctes, une fonction intégrée fusionnée, et une version compilée avec torch.compile.
Le cas est volontairement petit. Sur un modèle complet, la mémoire, la concurrence et les tailles de lot compliquent vite l'interprétation. L'objectif ici : capturer une trace exploitable, lire les temps CPU et Neuron, puis juger si la fusion est une piste crédible.
Pourquoi profiler avant d'optimiser
En général, les équipes accélèrent un modèle quand elles savent déjà où se cache le goulot d'étranglement. Le piège, c'est de confondre une impression avec une mesure. Un temps de réponse élevé ne dit pas, à lui seul, que le calcul sur l'accélérateur est lent. Le retard peut venir d'un enchaînement d'opérations lancées séparément, d'un transfert, d'une synchronisation, ou même d'une phase d'initialisation incluse par erreur dans la mesure.
En production, cette distinction change concrètement la nature du travail : réécrire un noyau qui ne domine pas le temps total ne changera pas grand-chose à la latence. À l'inverse, fusionner plusieurs opérations coûteuses peut supprimer des lancements et des intermédiaires inutiles. Bref : le profilage n'est pas un passage « obligatoire », c'est le moyen de s'assurer que l'effort tombe sur la bonne partie du système.
TorchNeuron dans l'écosystème PyTorch
TorchNeuron permet d'utiliser PyTorch sur les accélérateurs AWS Neuron, dont Trainium. L'intégration passe par le backend via PrivateUse1 de PyTorch : le code peut s'insérer dans les flux PyTorch sans avoir une logique comparable à cuda côté distribution. Dans la pratique, les API restent celles de PyTorch, y compris l'exécution en mode eager et la compilation via torch.compile.
Attention toutefois : « PyTorch tourne avec un autre backend » ne veut pas dire « toutes les opérations ont le même comportement ni les mêmes performances ». La façon dont le graphe est distribué, compilé et exécuté dépend de ce qui est supporté et de la configuration. Le profilage sert justement à observer ce chemin d'exécution, au lieu de le supposer.
Une question précise pour chaque profil
Avant de lancer l'outil, posez-vous la question exacte à laquelle vous voulez répondre. Ici, par exemple : « Quelles opérations de RMSNorm consomment du temps sur Neuron, et la fusion réduit-elle ce coût ? ». Cette contrainte évite de collecter une trace énorme sans savoir quoi en extraire.
Le tutoriel original publié par Amazon sur Hugging Face détaille l'exemple et fournit des scripts de reproduction : profilage système de TorchNeuron sur Trainium. Et si une opération n'apparaît pas dans un profil, ne concluez pas trop vite à « zéro coût » : vérifiez d'abord que les bons événements sont bien capturés, et que l'exécution a été correctement synchronisée.
Configurer une trace TorchNeuron exploitable
Le profilage repose sur deux éléments spécifiques à Neuron. NeuronConfig décrit les données système et matérielles à capturer, ainsi que le répertoire de sortie. NeuronProfiler prend cette configuration et exporte la trace PyTorch en parallèle des traces Neuron.
nc = NeuronConfig(
modes=[ProfileMode.DEVICE, ProfileMode.RUNTIME],
profile_output_dir=OUTPUT_DIR,
capture_enabled_for_nc="0",
correlate_device_time=True,
)
exporter = NeuronProfiler(nc)
Les modes sélectionnés déterminent ce qui est collecté. La configuration ci-dessus capture à la fois des informations d'appareil et de runtime. capture_enabled_for_nc choisit le cœur Neuron à tracer ; ici, c'est le cœur logique 0. Enfin, correlate_device_time=True remonte le temps de l'appareil dans le tableau key_averages(). On peut alors comparer le coût côté CPU et celui côté Neuron.
Inclure les activités Neuron et synchroniser
Pour récupérer les événements de la puce, le profil doit demander ProfilerActivity.PrivateUse1, en plus de l'activité CPU. La configuration Neuron passe ensuite à torch.profiler.profile via experimental_config :
with torch.profiler.profile(
activities=[ProfilerActivity.CPU, ProfilerActivity.PrivateUse1],
experimental_config=nc,
record_shapes=False,
on_trace_ready=exporter.export_trace,
schedule=torch.profiler.schedule(wait=2, warmup=1, active=1),
) as prof:
for step in range(4):
with torch.no_grad():
_ = rms_norm_fn(x)
torch.neuron.synchronize()
prof.step()
Le calendrier suit deux étapes d'attente, une phase de chauffe, puis une phase active, le tout réparti sur quatre appels à prof.step(). La chauffe sert à éviter d'inclure certains coûts de démarrage dans l'unique échantillon enregistré. Ce n'est pas une garantie de représentativité pour toutes les charges, mais c'est souvent plus fiable que de mesurer « le premier passage » sans précaution.
Et la synchronisation compte autant que l'outil. Sur l'accélérateur, les opérations peuvent être asynchrones par rapport au code Python. Sans torch.neuron.synchronize() dans la boucle, l'export de trace peut arriver avant la fin du travail sur l'appareil : vous risquez alors de voir un temps Neuron manquant ou incomplet. Une trace propre dépend donc autant de la discipline de mesure que du profilage lui-même.
Conserver les fichiers nécessaires
Après la capture, le répertoire contient la trace système et, si vous demandez les données appareil, des fichiers .neff et .ntff. Le NEFF correspond à l'exécutable compilé pour Neuron ; le NTFF contient la trace des instructions matérielles correspondante. Les deux sont reliés par un identifiant numérique commun, ce qui permet de les associer.
Pour ouvrir un profil système ou un profil combiné dans Neuron Explorer, il faut fournir le répertoire complet. Pour un profil limité à l'appareil, vous pouvez charger individuellement les fichiers NEFF et NTFF correspondants. Charger « seulement une partie » peut empêcher d'établir correctement le lien entre les événements de la trace et les fichiers compilés.
Lire les temps et suivre une opération
Le tableau key_averages() répond à une première question : quelles opérations PyTorch consomment le plus de temps appareil ? Pour pousser les plus coûteuses en haut, triez par self_privateuse1_time_total :
print(
prof.key_averages().table(
sort_by="self_privateuse1_time_total",
row_limit=10,
)
)
Chaque ligne regroupe les appels associés à une opération. Les colonnes indiquent notamment le nombre d'appels et les temps propres côté CPU et côté Neuron. Le « temps propre » ne doit pas être lu tout seul : une opération appelée beaucoup de fois peut dominer le total, même si chaque appel semble raisonnable. Avant de décider où travailler, regardez aussi le nombre d'appels et le total d'exécution.
Dans une version naïve de RMSNorm, l'opération aten::mul apparaît trois fois : mettre l'entrée au carré, multiplier par l'inverse de la racine carrée, puis appliquer le poids. aten::mean calcule la moyenne des carrés, aten::add ajoute epsilon, et aten::rsqrt calcule l'inverse de la racine carrée.
Un point à vérifier : l'ajout du scalaire epsilon peut aussi déclencher une copie entre CPU et appareil, car le scalaire commence côté CPU. Dans ce cas, le tableau peut afficher des événements CPU (allocations, transferts, etc.) qui ne seront pas visibles si vous ne regardez que les opérations de calcul. Avant d'attribuer toute la latence au noyau matériel, inspectez la trace complète.
Passer du tableau Ă Neuron Explorer
Le tableau vous donne des totaux ; Neuron Explorer vous montre la chronologie. Dans une vue système, vous voyez comment les événements se répartissent entre les voies du framework, du runtime et du matériel. Le temps se lit de gauche à droite, et les étapes d'une opération apparaissent dans des voies différentes.
| Voie | Événements observés |
|---|---|
| Distribution du framework | Appels d'opérateurs PyTorch, associés à un fil d'exécution |
| Appareil côté framework | Interactions avec le périphérique, regroupées par flux |
| Runtime Neuron | Appels du runtime, regroupés par cœur logique |
| Matériel Neuron | Exécution effective, regroupée par cœur physique |
En mode eager, une opération peut suivre une chaîne qui va de la distribution côté framework à la planification côté runtime, puis à l'exécution matérielle. Le visualiseur de dépendances aide à suivre un événement précis, par exemple aten::mul.Tensor, le long de cette chaîne. C'est utile pour vérifier si les chiffres du tableau correspondent à plusieurs lancements ou à des périodes d'attente entre exécutions.
L'idée est simple : le tableau localise un coût, la trace explique comment il se déroule. L'un ne remplace pas l'autre. Une trace seule est difficile à résumer sur un modèle complet, et un tableau seul ne raconte pas toujours les dépendances ni les intervalles entre les opérations.
Comparer trois implémentations de RMSNorm
RMSNorm est un bon exemple pour apprendre, parce que le calcul se décompose naturellement : carrés, moyenne, ajout d'un epsilon, racine carrée réciproque, puis multiplication par l'entrée et les poids. Les trois variantes ci-dessous partagent le même type d'entrée, les mêmes données et le même protocole de profilage. Les chiffres proviennent de l'expérience rapportée : ils servent surtout à montrer une méthode, pas à promettre des performances identiques sur votre charge.
Version naïve : une suite d'opérations
La première version exprime chaque étape avec des opérations PyTorch « classiques » :
def rms_norm_fn(t):
x_sq = t * t
mean_sq = x_sq.mean(dim=-1, keepdim=True)
rms = torch.rsqrt(mean_sq + 1e-6)
return t * rms * weight
Pour des entrées de forme (4, 512) et de type bfloat16, le profil indique un temps Neuron total de 2,355 ms. Les trois multiplications cumulent 1,1645 ms, soit 49,5 % du temps Neuron observé. Le reste se répartit sur mean, add et rsqrt.
Le point important n'est pas que aten::mul serait « mauvais » en soi. Dans cette version, il est appelé trois fois, et chaque étape produit des intermédiaires séparés. Résultat : plusieurs événements matériels et des frontières entre opérations à gérer.
Fonction intégrée : une opération fusionnée
La deuxième version remplace la séquence par la fonction PyTorch intégrée :
def rms_norm_fn(t):
return torch.rms_norm(t, [512], weight)
Dans la mesure rapportée, le tableau montre un seul appel à aten::_fused_rms_norm, pour 403,3 µs de temps Neuron. Au lieu de voir les opérations arithmétiques comme des événements séparés, le profil expose un noyau fusionné.
La trace confirme ce changement : une seule distribution et une seule exécution matérielle couvrent le calcul. La fusion évite de matérialiser plusieurs intermédiaires entre noyaux distincts, et limite les transitions. Et surtout : ce gain est mesuré dans ce cas précis. Rien ne garantit qu'une fonction intégrée sera toujours plus rapide sur toutes les formes de tenseurs et toutes les configurations.
torch.compile : laisser le compilateur fusionner
La troisième variante garde la formulation naïve, mais compile avec un backend Neuron :
@torch.compile(backend="neuron", fullgraph=True)
def naive_compiled(t):
x_sq = t * t
mean_sq = x_sq.mean(dim=-1, keepdim=True)
rms = torch.rsqrt(mean_sq + 1e-6)
return t * rms * weight
Dans le flux décrit, TorchDynamo capture le code Python sous forme de graphe FX, puis le backend Neuron reçoit un graphe complet à compiler. Le backend peut fusionner les opérations et produire un artefact NEFF. Le tableau montre une région compilée unique, à 375,5 µs de temps Neuron.
Les chiffres sont proches de la fonction intégrée : 375,5 µs pour la version compilée contre 403,3 µs pour torch.rms_norm. Cet écart ne suffit pas à dire que la version compilée serait « universellement » meilleure. On parle ici d'une expérience donnée, sur une petite charge et une configuration donnée. Le bénéfice intéressant, en revanche, est clair : moins d'exécutions. Là où la version naïve envoie plusieurs appels, les approches fusionnées ramènent ça à une seule région d'exécution par approche.
| Variante | Temps Neuron | Temps CPU propre | Appels rapportés |
|---|---|---|---|
| Opérations naïves | 2 355 µs | environ 125 µs | 6 |
torch.rms_norm |
403 µs | environ 48 µs | 1 |
Opérations naïves avec torch.compile |
376 µs | environ 98 µs | 1 |
Dans ces mesures, les variantes fusionnées consomment environ un sixième du temps appareil de la version naïve. On obtient donc une piste convaincante : réduire les lancements distincts et le trafic autour des intermédiaires. Mais ça ne dit pas encore que le même gain se retrouvera dans un modèle de production. Pour trancher, il faut profiler et benchmarker sur votre scénario réel (dimensions, lots, séquences, charge).
Appliquer la méthode à un modèle réel
Sur un grand modèle, le profilage ne cherche pas forcément l'opération la plus « spectaculaire ». Il vise surtout à construire un cycle court et répétable : localiser, comprendre, modifier, puis re-mesurer. Cette approche limite le risque de passer des jours à optimiser une fraction marginale du système.
Une démarche en cinq étapes
- Définissez le scénario mesuré. Fixez les entrées, la taille de lot, la longueur de séquence et le mode d'exécution. Un profil d'inférence ne répond pas automatiquement à une question d'entraînement.
- Écartez le démarrage. Mettez en place une phase d'attente et de chauffe adaptée, puis synchronisez l'appareil avant de valider le relevé.
- Triez les événements par temps appareil. Repérez les opérations qui dominent en fréquence et en part du total. Ne vous arrêtez pas au nom de la première ligne.
- Suivez les candidats dans Explorer. Vérifiez la chaîne framework-runtime-matériel, les transferts et les intervalles entre exécutions.
- Testez une modification à la fois. Par exemple, essayez une fonction fusionnée ou
torch.compile, puis refaites la mĂŞme mesure dans les mĂŞmes conditions.
Cette méthode aide aussi à éviter une erreur fréquente : confondre la durée d'un noyau avec la latence de bout en bout. Une requête réelle inclut souvent du prétraitement, des transferts, plusieurs couches, et de la synchronisation. Si votre objectif métier est de réduire la latence p95 d'un service, c'est cette métrique qu'il faudra valider, en gardant le profil système comme outil de diagnostic.
Les limites Ă garder en tĂŞte
Le profilage a un coût et peut modifier légèrement le comportement observé. Une trace plus détaillée ajoute de la collecte ; ce n'est pas forcément ce qu'il faut laisser activé en continu en production. Les résultats dépendent aussi du matériel, de la version logicielle, de la forme exacte des tenseurs et du protocole de mesure. Gardez ces paramètres avec vos résultats si vous voulez comparer correctement.
Il faut aussi distinguer temps CPU et temps appareil. Un temps CPU faible ne veut pas dire que l'opération est « gratuite ». Inversement, un temps appareil élevé ne prouve pas que le calcul arithmétique est le seul problème. Les transferts, la planification et le nombre de lancements peuvent peser dans le résultat. La vue système permet d'examiner ces hypothèses, puis le profil appareil (traité dans la suite de la série) sert à creuser les moteurs de calcul et la mémoire.
Pour une équipe produit ou toute organisation qui évalue Trainium, la question pertinente n'est donc pas « quel accélérateur est le plus rapide ? ». C'est plutôt : « quelle configuration répond à nos exigences de latence, de débit et de coût pour notre charge ? ». Le profil donne des éléments techniques utiles pour répondre, sans remplacer ni un test représentatif, ni une analyse économique.
Conclusion : mesurer avant de changer le code
Le profilage système de TorchNeuron donne deux vues complémentaires : key_averages() met en évidence les opérations qui consomment du temps Neuron, et Neuron Explorer montre comment elles transitent par le framework, le runtime et le matériel. Sur RMSNorm, la fusion réduit plusieurs exécutions en une seule et fait passer le temps appareil mesuré de 2,355 ms à environ 0,4 ms.
Ce chiffre ne doit pas être lu comme une promesse à copier-coller. En revanche, la méthode est réutilisable : construisez une trace propre, repérez un point chaud, vérifiez la chronologie, appliquez une modification, puis re-mesurez dans des conditions comparables. C'est plus fiable que d'optimiser à l'intuition.
Pour aller plus loin, commencez par reproduire les trois variantes sur votre environnement, puis appliquez la même démarche à une requête représentative de votre modèle. L'étape suivante logique consiste à passer du profil système à l'analyse des moteurs de calcul et de la mémoire dans le profil appareil.