Profiling TorchNeuron (Partie 1) : Profilage au niveau système sur AWS Trainium.

Beaucoup de praticiens du ML ont une certaine intuition de la façon de faire fonctionner les modèles rapidement, mais moins savent comment analyser pourquoi un modèle est lent. Dans la Partie 1 de cette série, vous apprendrez à analyser systématiquement les performances d'inférence et à identifier où le temps est passé.

Cet article part de zéro et ne suppose aucune expérience préalable du profilage. Le tutoriel utilise AWS Trainium, une puce d'IA conçue sur mesure pour les charges de travail d'entraînement et d'inférence en deep learning, ainsi que TorchNeuron, son backend PyTorch natif open source. TorchNeuron prend en charge les workflows PyTorch standards sur Trainium, y compris l'exécution en mode eager, torch.compile, et les API PyTorch standards, avec des modifications de code minimales. De plus, TorchNeuron s'intègre avec Neuron Kernel Interface (NKI) pour permettre des opérateurs personnalisés et l'optimisation des performances pour le matériel Trainium. À la fin, vous devriez savoir :

• Comment configurer torch.profiler avec le backend Neuron et capturer une trace propre.
• Comment lire le tableau key_averages() pour identifier quelles opérations consomment du temps sur le dispositif Neuron et comparer leur Self CPU time.
• Comment suivre une opération à travers Neuron Explorer, depuis l'envoi par le framework en passant par l'ordonnancement du runtime jusqu'à l'exécution matérielle sur un NeuronCore.
• Comment ces signaux évoluent entre trois implémentations de la même opération.

Neuron comme backend PyTorch

Avant de profiler une charge de travail, il est utile de comprendre comment Neuron s'intègre avec PyTorch.

Le backend Neuron utilise le point d'extension PrivateUse1 de PyTorch. Ce mécanisme permet à un accélérateur de s'intégrer au framework sans ajouter de clé de dispatch dédiée telle que cuda. Le paquet torch_neuronx fonctionne avec le PyTorch amont pour enregistrer Neuron comme backend natif pour le code PyTorch. Pour plus de détails sur ce mécanisme d'extension, voir Facilitating New Backend Integration by PrivateUse1.

Configurer le profiler

Vous utilisez deux objets pour profiler une charge de travail Neuron :

• NeuronConfig contrôle quels flux de données Neuron et quels NeuronCores capturer et où écrire la sortie. Vous le passez à torch.profiler.profile comme experimental_config.
• NeuronProfiler prend un NeuronConfig et fournit un callback export_trace . Le callback écrit la trace du framework, trace.json, à côté des traces système 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)

Le modes paramètre sélectionne les flux de données à capturer, comme les instructions du dispositif et les traces d'exécution. Le capture_enabled_for_nc paramètre sélectionne les NeuronCores à tracer. Définir correlate_device_time=True propage le temps du dispositif depuis la trace dans la table key_averages() , où vous pouvez comparer le temps CPU et le temps du dispositif Neuron. Pour la référence complète des paramètres et les ProfileMode valeurs disponibles, consultez la documentation NeuronConfig Profile Modes.

Outillage : Neuron Explorer

Neuron Explorer est une visionneuse basée sur un navigateur pour la trace système, la chronologie du dispositif et la répartition par moteur. Vous pouvez y accéder via un tunnel SSH ou l'extension Visual Studio Code. Dans cet article, vous utilisez Explorer pour confirmer ce que key_averages() signale numériquement et inspecter ce qui se passe entre les dispatches.

Après le profilage, le répertoire de sortie contient la trace système et, lorsque vous capturez des données du dispositif, les fichiers .neff et .ntff correspondants. Le format Neuron Executable File Format (NEFF) contient l'exécutable Neuron compilé, tandis que le fichier NTFF correspondant contient sa trace d'instructions matérielle. Les fichiers partagent un hachage numérique qui les identifie comme une paire. Pour ouvrir le profil dans Neuron Explorer, téléchargez le répertoire de sortie complet d'un profil système ou combiné. Pour un profil du dispositif uniquement, téléchargez individuellement les fichiers .neff et .ntff correspondants. Consultez Step 5: Upload your profile pour les instructions via l'interface et la ligne de commande.

Profiler RMSNorm

Pour cette première analyse, vous profilez RMSNorm, une petite opération avec une séquence claire d'opérations arithmétiques et de réduction.

Vous profilez trois implémentations de la même opération :

1. RMSNorm naïf avec des opérations PyTorch individuelles : multiplication, moyenne, addition, racine carrée réciproque et multiplication.
2. torch.rms_norm(), la fonction intégrée de PyTorch abaissée en un seul noyau Neuron fusionné.
3. RMSNorm naïf avec torch.compile, en utilisant les mêmes opérations que la première implémentation.

Les scripts utilisent les mêmes entrées, la même configuration du profileur et le même ordonnancement afin que les résultats soient directement comparables.

Tous les scripts pour reproduire ce travail sont disponibles dans le torch-neuronx répertoire examples.

Profil 1 : opérations PyTorch naïves

Commencez par écrire RMSNorm avec des opérations PyTorch individuelles. Chaque ligne correspond à une ou plusieurs opérations ATen et peut créer des tenseurs intermédiaires.

device = torch.device('neuron')

x = torch.randn(4, 512, dtype=torch.bfloat16, device=device)
weight = torch.ones(512, dtype=torch.bfloat16, device=device)

def rms_norm_fn(t):
    x_sq = t * t                                    # square
    mean_sq = x_sq.mean(dim=-1, keepdim=True)       # mean of squares
    rms = torch.rsqrt(mean_sq + 1e-6)               # 1 / sqrt(mean_sq + eps)
    return t * rms * weight                         # normalize + scale

Appelez le profileur avec deux arguments requis pour le backend Neuron :

1. Incluez ProfilerActivity.PrivateUse1 pour collecter l'activité du dispositif Neuron.
2. Passez NeuronConfig comme experimental_config pour activer les fonctionnalités Kineto spécifiques à Neuron.

L'ordonnancement utilise wait=2, warmup=1, et active=1 sur quatre appels prof.step() . Cela enregistre une itération d'échauffement sans inclure les frais initiaux de configuration.

Appelez torch.neuron.synchronize() dans la boucle du profileur. Sans cette synchronisation, le profileur peut exporter avant que le dispositif ne termine et laisser la colonne de temps Neuron vide.

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()

Une fois le profilage terminé, examinez la table key_averages() . Triez par self_privateuse1_time_total pour placer en premier les opérations ayant le coût le plus élevé sur le dispositif. La table présentée ici est filtrée pour ne montrer que les opérations RMSNorm qui indiquent un temps sur le dispositif Neuron. La sortie complète de key_averages() inclut également les événements côté CPU tels que les allocations et les transferts hôte-dispositif.

print(prof.key_averages().table(sort_by="self_privateuse1_time_total", row_limit=4))

La table des moyennes clés montre que trois appels à aten::mul représentent la plus grande part du temps sur le dispositif Neuron.

La table montre quelles opérations ATen ont été exécutées, à quelle fréquence, et combien de temps elles ont passé sur le CPU et sur le dispositif Neuron. aten::mul a le temps total sur le dispositif le plus élevé car la fonction effectue trois multiplications.

Total Self NEURON : 2.355 ms

Chaque ligne correspond directement à la fonction. aten::mul s'exécute pendant x * x, x * rms, et la multiplication finale par weight. aten::mean réduit la dernière dimension. aten::add ajoute la valeur epsilon et déclenche une copie hôte-vers-périphérique car le scalaire démarre sur le CPU. aten::rsqrt calcule la racine carrée inverse.

Neuron Explorer suit une opération aten::mul depuis la répartition au niveau du framework jusqu'à la planification à l'exécution et l'exécution sur les unités de calcul du matériel Neuron.

Lire la trace dans Explorer

Dans Explorer, le temps s'écoule de gauche à droite et le cycle de vie des événements de haut en bas. Chaque opération ATen suit le même schéma : création dans la voie du framework, planification dans la voie du runtime sous forme de nrta_execute_schedule, et exécution sur le périphérique dans la voie matérielle sous forme de nc_exec_running.

La vue de trace système organise les événements en voies. Dans le mode de regroupement Default, les libellés suivent le schéma instance/PID/sous-système/identifiant :

En mode eager, chaque opération s'exécute indépendamment. Utilisez le visualiseur de chaîne de dépendances pour suivre aten::mul.Tensor depuis la couche du framework, en passant par la planification du runtime, jusqu'à l'exécution sur le périphérique.

Le visualiseur de chaîne de dépendances met en évidence le cycle de vie d'une opération à travers les voies du framework, du runtime et du matériel.

Profil 2 : torch.rms_norm

Remplacez les opérations individuelles par la fonction rms_norm intégrée de PyTorch. Neuron convertit l'opération en un seul kernel fusionné :

def rms_norm_fn(t):
    return torch.rms_norm(t, [512], weight)

Le RMSNorm fusionné apparaît comme une seule opération de périphérique dans le tableau des moyennes des clés.

Total Self NEURON : 403.3 μs

Le RMSNorm fusionné suit un unique chemin framework-vers-matériel au lieu de dispatcher chaque opération arithmétique séparément.

Comparer le chemin de lowering

Dans l'implémentation naïve, chaque opération franchit indépendamment la frontière de dispatch et suit son propre chemin framework-vers-runtime-vers-matériel. Dans l'implémentation fusionnée, un seul appel Python couvre le calcul complet. La trace contient un seul dispatch et un seul événement matériel, sans intervalles entre les opérations composantes.

Profil 3 : PyTorch naïf avec torch.compile

Pour le dernier exemple, compilez la fonction RMSNorm naïve d'origine avec torch.compile.

Avec @torch.compile(backend="neuron"), TorchDynamo capture le bytecode Python sous forme de graphe FX et AOTAutograd l'abaisse en opérations ATen. Le backend Neuron reçoit le graphe complet au lieu de dispatcher chaque opération séparément. Il fusionne les opérations, optimise l'accès mémoire et compile le graphe en un unique artefact NEFF. L'artefact s'exécute comme un seul dispatch matériel.

@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

La fonction compilée apparaît comme une seule Torch-Compiled Region dans le tableau des moyennes clés.

Self NEURON total: 375.5 μs

Le graphe compilé suit un seul chemin à travers l'ordonnancement runtime et l'exécution matérielle.

Le tableau montre une seule exécution sur le périphérique, comme dans le Profil 2, mais le code emprunte un chemin différent. Au lieu d'appeler une opération fusionnée spécialement conçue, le compilateur découvre la fusion à partir de la fonction Python d'origine. Dans cette exécution, le temps Neuron est comparable à torch.rms_norm, à 375 μs contre 403 μs.

Résumé comparatif

Both torch.rms_norm and torch.compile fusionnent RMSNorm en une seule dispatch de kernel. Dans ces exécutions, chacune des deux approches prend environ un sixième du temps device de l'implémentation naïve. La fusion supprime les espaces entre les exécutions matérielles et évite de matérialiser des valeurs intermédiaires telles que x_sq et mean_sq entre des opérations dispatchées séparément.

Conclusion

key_averages() montre quelles opérations consomment du temps device, tandis que sa sortie complète révèle également la configuration côté CPU et les transferts. Explorer fournit la trace visuelle correspondante. Vous pouvez suivre une opération depuis le dispatch du framework, en passant par l'ordonnancement du runtime, jusqu'à l'exécution matérielle, et voir les espaces que la fusion supprime.

Vous pouvez utiliser le même workflow sur des modèles plus grands : repérer un hotspot dans le tableau, le vérifier dans la trace, puis déterminer si une opération fusionnée ou torch.compile peut réduire son coût.

Dans la Part 2, nous passerons au profil au niveau du device. Vous ouvrez le NEFF, examinez la répartition du TensorEngine, du VectorEngine et de l'accès direct à la mémoire (DMA), et placez le kernel sur le roofline pour déterminer s'il est limité par le calcul (compute-bound) ou par la mémoire (memory-bound).

Remerciements particuliers à Aritra Roy Gosthipaty pour l'inspiration de commencer cette série, et à mes pairs Pinak Panigrahi, Samir Araujo, Akhil Azhikodan, Esha Lakhotia et tous les autres pour leurs relectures !