Introduction

Dans cette partie de la série sur le profilage dans PyTorch, nous allons explorer comment remplacer une paire matmul-add écrite à la main par une couche linéaire (nn.Linear) avec biais. Cela constitue le bloc de base de tout modèle d'apprentissage profond. Nous allons également empiler trois de ces couches pour former un bloc Multilayer Perceptron (MLP) avec une activation entre elles.

Contexte Technique

Le script pour cet article est disponible et peut être exécuté sur un GPU NVIDIA A100-SXM4-80GB. La couche linéaire nn.Linear est un module qui enveloppe la même multiplication matricielle et l'addition que celles que nous avons déjà profilées dans la partie 1. La différence réside dans le fait qu'elle possède son propre poids et biais en tant que paramètres et expose une méthode forward familière aux utilisateurs de PyTorch.

En utilisant nn.Linear avec bias=True, nous émulons les opérations de multiplication et d'addition que nous avons vues dans la partie 1 de la série. L'opération peut être écrite comme y = x @ w.T + b, où x est l'entrée, w est le poids, et b est le biais.

Analyse et Implications

Lorsque nous exécutons le script et vérifions le profil, nous remarquons qu'il y a une opération aten::t (transposition) avant l'opération aten::addmm (multiplication et addition). Cela indique que nn.Linear transpose le paramètre de poids puis le multiplie par l'entrée. L'ajout du biais est intégré dans le noyau de multiplication matricielle via ce qu'on appelle un épilogue, ce qui évite de charger ou d'écrire à nouveau les données en mémoire.

En comparant les traces de profil pour un appel forward d'une couche linéaire unique avec et sans compilation, nous constatons que le noyau cuBLAS GEMM sur le GPU et l'opération aten::addmm sur le CPU sont les mêmes. Cela signifie que pour une seule opération GEMM avec biais, la compilation n'a pas grand-chose à faire.

Perspective

Il est important de noter que la compilation est plus utile lorsqu'il y a plusieurs opérations à fusionner. En regardant un MLP, nous pouvons voir comment la compilation peut réellement faire une différence en réduisant les frais généraux de dispatching sur le CPU et en améliorant les performances globales. Il est essentiel de comprendre les mécanismes sous-jacents pour utiliser efficacement les outils de compilation et d'optimisation dans PyTorch.