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

De nombreux 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 dépensé.

Cet article part de zéro et ne suppose aucune expérience préalable du profiling. La démonstration utilise AWS Trainium, une puce d'IA conçue spécifiquement pour les charges de travail d'entraînement et d'inférence en deep learning, ainsi que TorchNeuron, son backend natif PyTorch open source. TorchNeuron prend en charge les workflows PyTorch standard sur Trainium, y compris l'exécution en eager, torch.compile, et les API PyTorch standard, 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 temps Self CPU.
• Comment suivre une opération à travers Neuron Explorer, depuis la répartition au niveau du framework jusqu'à la planification par le runtime et l'exécution matérielle sur un NeuronCore.
• Comment ces signaux changent 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 à 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 une clé de répartition dédiée telle que cuda. Le paquet torch_neuronx fonctionne avec 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 en tant que experimental_config.
• NeuronProfiler prend un NeuronConfig et fournit un callback export_trace . Le callback écrit la trace du framework, trace.json, aux côtés 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, tels que les instructions du device 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 device de la trace dans la key_averages() table, où vous pouvez comparer le temps CPU et le temps device Neuron. Pour la référence complète des paramètres et les ProfileMode valeurs disponibles, consultez la documentation NeuronConfig Profile Modes.

Tooling : Neuron Explorer

Neuron Explorer est une visionneuse basée sur un navigateur pour la trace système, la timeline du device 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() rapporte numériquement et inspecter ce qui se passe entre les dispatchs.

Après le profiling, le répertoire de sortie contient la trace système et, lorsque vous capturez des données du device, les fichiers .neff et .ntff correspondants. Le Neuron Executable File Format (NEFF) contient l'exécutable Neuron compilé, tandis que le 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 device uniquement, téléchargez individuellement les fichiers .neff et .ntff correspondants. Consultez Step 5: Upload your profile pour les instructions via l'interface utilisateur 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 inverse et multiplication.
2. torch.rms_norm(), la fonction native PyTorch abaissée vers un seul kernel 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 de profileur et le même calendrier 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.

Le calendrier utilise wait=2, warmup=1, et active=1 sur quatre prof.step() appels. Cela enregistre une itération d'échauffement sans inclure la surcharge de la configuration initiale.

Appelez torch.neuron.synchronize() à l'intérieur de la boucle du profileur. Sans cette synchronisation, le profileur peut exporter avant que le dispositif ait terminé 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 de dispositif le plus élevé. La table présentée ici est filtrée pour n'afficher que les opérations RMSNorm qui rapportent un temps de 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-vers-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 de dispositif Neuron.

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

Self NEURON total : 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-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 le dispatch du framework, en passant par la planification au runtime, jusqu'à l'exécution sur les lanes 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 lane du framework, planification dans la lane du runtime sous la forme nrta_execute_schedule, et exécution sur le périphérique dans la lane matérielle sous la forme nc_exec_running.

La vue de trace système organise les événements en lanes. Dans le mode de regroupement Default, les étiquettes suivent le modèle 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 framework, en passant par la planification au 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 lanes framework, runtime et 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 sur le périphérique dans la table des moyennes clés.

Self NEURON total : 403.3 μs

Le RMSNorm fusionné suit un chemin unique du framework au 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-runtime-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 le réduit 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 device, 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, avec 375 μs contre 403 μs.

Résumé comparatif

Les deux torch.rms_norm et torch.compile fusionnent fuse RMSNorm en un seul dispatch de kernel. Dans ces exécutions, chaque approche prend environ un sixième du temps device de l'implémentation naïve. La fusion supprime les écarts 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 aussi la configuration et les transferts côté CPU. 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 écarts que la fusion supprime.

Vous pouvez utiliser le même workflow sur des modèles plus grands : trouvez un point chaud dans la table, vérifiez-le dans la trace, puis déterminez si une opération fusionnée ou torch.compile peut réduire son coût.

Dans la Partie 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 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).

Un merci tout particulier à 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 !