Ce qui anime réellement un réseau neuronal : les compilateurs d’IA sous Python

Toute personne ayant une expérience en compilation et quelques connaissances en apprentissage automatique n’a peut-être jamais fait le lien entre ces deux domaines. Les analyseurs lexicaux, les analyseurs syntaxiques, les arbres syntaxiques et LLVM IR appartiennent à un monde, tandis que les réseaux neuronaux en font partie d’un autre. Ils se rencontrent dans les compilateurs d’IA/apprentissage automatique, et la première surprise est la suivante : un modèle d’apprentissage profond ne s’exécute pas vraiment en Python.

Cela semble étrange, car une ligne comme celle-ci donne l’impression que Python fait tout le travail :

output = model(input)

Cependant, Python ne effectue pas lui-même des milliards de multiplications de matrices. Un processus bien plus long a lieu en arrière-plan.

De la formule aux instructions matérielles

Imaginez un petit réseau qui passe des données d’entrée par une multiplication de matrices avant de continuer :

Input
  ↓
Matrix Multiplication
  ↓
ReLU
  ↓
Matrix Multiplication
  ↓
Output

En Python, son noyau pourrait ressembler à ceci :

y = relu(x @ W + b)

Cela semble simple, mais réfléchissez à ce dont a besoin la machine. Une CPU ou une GPU n’a aucune notion de cette fonction :

relu()

ou de cet opérateur :

x @ W

telles qu’elles apparaissent dans le code. Le matériel exécute les instructions. À un moment donné, la description de haut niveau du modèle doit être transformée en quelque chose que la puce peut exécuter efficacement, et c’est là que interviennent les compilateurs d’IA/ML.

Représentation simplifiée de la chaîne, du modèle jusqu’au processeur :

AI Model
   ↓
PyTorch / TensorFlow
   ↓
Computational Graph
   ↓
Intermediate Representation (IR)
   ↓
Optimisations
   ↓
Lowering
   ↓
Hardware-specific Code
   ↓
Machine Instructions
   ↓
CPU / GPU / NPU

C’est beaucoup d’étapes, il vaut donc mieux les aborder une par une.

Ce que fait le framework

Avec PyTorch, la même idée est écrite de cette manière :

y = torch.relu(x @ W + b)

Un développeur ne souhaite pas s’occuper des registres GPU, de la disposition de la mémoire, de l’ordonnancement des instructions ou de milliers de threads matériels. Il indique simplement quelles calculs il souhaite effectuer. Python et PyTorch offrent un moyen pratique pour le formuler, tandis qu’un autre outil doit traduire cette description en instructions compréhensibles par le matériel.

Le graphe de calcul

Un réseau neuronal peut être considéré comme un graphe de calcul. L’expression :

y = relu(x @ W + b)

on peut le représenter sous forme de graphe des opérations connectées :

Désormais, le compilateur dispose de quelque chose de concret sur quoi travailler. Une fois les opérations visibles, il peut déterminer si le calcul peut être accéléré. C’est là que l’optimisation devient essentielle.

Fusion des opérateurs et trafic mémoire

Prenons un modèle contenant une chaîne comme celle-ci :

MatMul
   ↓
Add
   ↓
ReLU

Un compilateur peut remarquer que ces étapes peuvent être fusionnées ou adaptées au matériel cible. Une technique clé est la fusion des opérateurs : au lieu de traiter chaque opération comme une étape distincte, certaines sont combinées, ce qui réduit le travail inutile, en particulier le transfert de données entre la mémoire et les unités de calcul.

Cela est important car le transfert de données est coûteux dans les charges de travail d’intelligence artificielle modernes. Le goulot d’étranglement ne réside pas seulement dans la vitesse à laquelle un résultat peut être calculé, mais aussi dans l’efficacité avec laquelle les données d’entrée parviennent au calcul. C’est l’une des raisons pour lesquelles ce domaine est intéressant.

Représentations intermédiaires

Vient ensuite un concept familier des cours classiques de compilation : la représentation intermédiaire (RI). Un compilateur traditionnel passe rarement directement par celle-ci :

Source Code

vers ceci :

Machine Code

Il passe par les IR. Les compilateurs d’IA suivent le même principe, mais leurs représentations doivent retenir les idées importantes pour l’apprentissage automatique. Là où un compilateur général prend en compte :

variables
functions
loops

un compilateur d’IA doit également pouvoir représenter des éléments tels que :

tensors
shapes
matrix operations
convolutions
attention
memory layouts

C’est ce qui rend des projets comme MLIR remarquables. MLIR fournit une infrastructure permettant de représenter et de transformer des calculs à plusieurs niveaux d’abstraction.

Descendre pas à pas

Le compilateur passe progressivement d’une forme de haut niveau à des formes plus proches du matériel cible. Imaginez que vous descendez un escalier :

High-level ML operation
          ↓
Tensor operation
          ↓
Lower-level operation
          ↓
Hardware-oriented operation
          ↓
Machine instructions

À chaque niveau, les idées abstraites deviennent un peu plus concrètes, jusqu’à ce que le compilateur génère quelque chose que le matériel cible peut utiliser.

Le matériel de base

Nous arrivons maintenant à la partie à laquelle personne ne pense lors de l’écriture :

model(input)

Le calcul final peut être effectué sur un CPU, une GPU, une NPU, une TPU ou un autre accélérateur d’IA. Aucun d’eux ne se soucie du fait que le code original indiquait :

y = relu(x @ W + b)

Ils ont besoin d’instructions exécutables et de données organisées de manière appropriée pour ce dispositif. Une GPU peut nécessiter du code et des modèles d’exécution spécifiques à cette carte. Une CPU utilise des instructions CPU et la vectorisation. Un accélérateur peut avoir besoin d’un backend complètement différent. Le même modèle de haut niveau peut donc être compilé différemment en fonction de sa cible.

Alors, quelle place pour Python ?

Python n’a pas disparu, et il reste très utile. C’est l’interface pour définir les modèles, contrôler l’exécution, charger des données et créer des applications. Il ne s’occupe simplement pas nécessairement lui-même de toutes les opérations numériques. Un modèle mental simplifié :

Cela change la façon dont on perçoit l’IA. Vue uniquement du côté du modèle, l’image est concise :

Data → Model → Prediction

Ajoutez le compilateur, et il y a une autre couche entre eux :

Model
  ↓
Compiler
  ↓
Hardware

Le modèle indique quelles calculs doivent être effectués. Le compilateur détermine comment les représenter et les optimiser. Le matériel s’occupe de leur exécution.

Pourquoi un autre niveau est nécessaire

Puisque des frameworks existent déjà, pourquoi ajouter un niveau de compilateur ? Parce que le matériel d’IA devient de plus en plus spécialisé. Il existe des CPU, des GPU, des TPU, des NPU, des accélérateurs d’inférence et des puces personnalisées, chacune ayant des caractéristiques différentes. Un modèle qui fonctionne bien sur une architecture peut nécessiter des ajustements très différents sur une autre. Le compilateur fait le lien entre ces deux aspects. D’un côté, il y a la question de ce que le modèle souhaite calculer :

WHAT
The ML model wants to compute

L’autre aspect concerne la manière dont un appareil donné doit le faire de manière efficace :

HOW
A specific piece of hardware should execute it efficiently

Ce pont gagne constamment en importance.

Une philosophie familière dans un nouveau domaine

Tout ceux qui connaissent les compilateurs traditionnels reconnaissent le pipeline de base :

Source Code
    ↓
Lexer
    ↓
Parser
    ↓
AST
    ↓
IR
    ↓
Optimisation
    ↓
Machine Code

Apprentissage automatique : une philosophie similaire appliquée aux modèles :

ML Model
    ↓
ML Representation
    ↓
ML IR
    ↓
Optimisation
    ↓
Lowering
    ↓
Hardware Code

Le domaine est différent et les problèmes sont bien plus difficiles, mais la mentalité reste la même : représenter, transformer, optimiser, réduire, exécuter.

L’idée suivante logique est de créer un petit compilateur d’IA pour comprendre son fonctionnement. Il ne rivaliserait pas avec LLVM, XLA ou MLIR, et son objectif serait simplement d’être suffisamment petit pour être compris, peut-être en partant d’un langage ML simple :

Simple ML Language
        ↓
       Parser
        ↓
      ML AST
        ↓
       ML IR
        ↓
   Simple Optimiser
        ↓
    Lowering
        ↓
   LLVM / Backend

La pile en une seule ligne

L’IA se présente généralement à travers une interface très simple :

model(input)

Derrière cet appel unique se trouve toute une pile logicielle et matérielle : des frameworks, des graphes, des représentations intermédiaires, des compilateurs, des optimisations, des backends, du matériel et enfin des instructions machine. Exécuter un modèle signifie faire fonctionner un écosystème complet de compilateurs et de matériel, avec Python uniquement en couche supérieure.