Les internals de PyTorch (2019)

• Une explication de l’architecture interne de PyTorch, qui sert de guide pour les personnes souhaitant contribuer à la base de code C++ de PyTorch
• L’objectif de cet article est d’aider à comprendre la structure de la bibliothèque de tenseurs de PyTorch ainsi que les techniques de différentiation automatique (autograd), et à s’orienter dans la base de code

Structure de base des tenseurs PyTorch

• Dans PyTorch, le tenseur est la structure de données la plus fondamentale
• Un tenseur est une structure de données en n dimensions, capable de stocker des valeurs scalaires comme des nombres à virgule flottante (float), des entiers (int), etc.
• Un tenseur contient les métadonnées suivantes :
  • taille (size) : informations sur les dimensions du tenseur
  • dtype : type des données stockées (par ex. float32, int64, etc.)
  • device : emplacement où les données sont stockées (CPU, CUDA, etc.)
  • stride : informations de décalage dans la mémoire physique des données

• Rôle du stride

  • le stride sert à convertir un indice logique en emplacement mémoire physique
  • le stride définit un décalage pour chaque dimension, et l’emplacement mémoire physique est déterminé en multipliant l’indice par la valeur de stride
  • grâce au stride, il est possible de visualiser les mêmes données différemment sous forme de view sans créer un nouveau tenseur

Concepts de tenseur et de stockage (Storage)

• Dans PyTorch, le tenseur ne stocke pas directement les données réelles → les données sont gérées dans le stockage (Storage)
• Tensor = taille + dtype + device + stride + offset
• Plusieurs tenseurs peuvent partager un même stockage → prise en charge du concept de vue (View)
• La séparation entre stockage et tenseur permet une utilisation plus efficace de la mémoire

Processus de dispatch des opérations sur les tenseurs

• Dans PyTorch, les opérations passent par deux étapes de dispatch :
  1. Dispatch basé sur le type de device et le layout
     • un code d’implémentation différent est exécuté selon qu’il s’agit d’un tenseur CPU ou d’un tenseur CUDA
  2. Dispatch basé sur le dtype
     • des kernels différents sont appelés selon le type de données, par ex. float vs. int

Modèle d’extension des tenseurs PyTorch

• Les trois principaux éléments d’extension d’un tenseur :

  • Device : définit la manière dont la mémoire est allouée sur CPU, GPU, TPU, etc.
  • Layout : définit la manière dont le tenseur est stocké en mémoire (par ex. stockage contigu, stockage creux (sparse), etc.)
  • dtype : définit le type de données stocké dans chaque élément du tenseur

• Options d’extension :

  • il est possible d’étendre les tenseurs en modifiant directement le code de PyTorch
  • il est possible d’écrire une classe wrapper qui encapsule un tenseur existant
  • si un wrapper est nécessaire pendant la différentiation automatique, une extension directe est requise

Fonctionnement de la différentiation automatique (Autograd)

• PyTorch effectue la différentiation automatique sur la base de la rétropropagation (reverse-mode differentiation)
• Lors des opérations en propagation avant (forward), un graphe est créé → lors de la rétropropagation, ce graphe est parcouru pour effectuer la différentiation
• Autograd gère des informations supplémentaires telles que :
  • AutogradMeta : métadonnées associées au tenseur et utilisées pour la rétropropagation
  • enregistrement des résultats des opérations puis calcul des dérivées lors de la rétropropagation

Structure du code PyTorch et emplacement des fichiers

• Principaux répertoires de la base de code PyTorch :
  • torch/ → module Python (code Python)
  • torch/csrc/ → code de liaison Python/C++, moteur de différentiation automatique, compilateur JIT, etc.
  • aten/ → définition des opérations sur les tenseurs (inclut la plupart des opérations cœur)
  • c10/ → définition des structures cœur comme les tenseurs et les stockages

Processus d’exécution des opérations PyTorch

• Exemple : processus d’exécution lors d’un appel à torch.add() :
  1. conversion des arguments de Python vers le code C++
  2. exécution du dispatch dans VariableType
  3. exécution du dispatch basé sur le device / le layout
  4. exécution du kernel final

Processus d’écriture des kernels et outils

• Dans PyTorch, un kernel s’écrit selon les étapes suivantes :
  1. Rédaction des métadonnées de l’opération : définition de la signature de fonction, des devices pris en charge et des types de données pris en charge
  2. Validation des entrées : vérification des dimensions, des types, etc.
  3. Allocation du tenseur de sortie
  4. Dispatch par dtype : exécution du kernel selon le type de données
  5. Traitement parallèle : sur CPU, utilisation d’OpenMP ; sur CUDA, utilisation de la parallélisation intégrée
  6. Accès aux données et calcul : utilisation de TensorAccessor, TensorIterator, etc.

Principales macros de dispatch

• AT_DISPATCH_ALL_TYPES → effectue le dispatch selon le dtype
• Des macros sont fournies pour différents types de données → optimisation des performances possible

Conseils pour optimiser les performances et améliorer l’efficacité de travail

• Réduire au minimum les modifications des fichiers d’en-tête → une modification entraîne une reconstruction complète du code
• Configurer un environnement de développement local → réduit au minimum le temps perdu à utiliser la CI
• Utiliser ccache → permet de gagner du temps de recompilation
• Utiliser un serveur puissant → permet de réduire le temps de compilation C++ et de build CUDA

Guide de contribution à PyTorch

• Bons points d’entrée pour commencer à contribuer :
  • les issues portant le label triaged → issues déjà vérifiées par les développeurs de PyTorch
  • l’amélioration de la documentation et l’aide à la reproduction des bugs
  • donner son avis sur les RFC (propositions de fonctionnalités) de PyTorch
• PyTorch a grandi grâce aux contributeurs open source, et la participation de la communauté est la bienvenue