Progrès du machine learning pour le machine learning
(blog.research.google)- À mesure que les modèles de ML grossissent, leur efficacité d’exécution dépend fortement non seulement du matériel, mais aussi de l’optimisation par le compilateur ; Google Research et Google DeepMind présentent une approche qui améliore à nouveau ce domaine grâce au ML
- Sous des frameworks comme TensorFlow, JAX et PyTorch, les décisions du compilateur peuvent modifier fortement le temps d’exécution et l’utilisation des ressources, même pour un même modèle
- TpuGraphs est un jeu de données contenant les graphes de calcul de programmes ML pour TPU, les paramètres de compilation et les temps d’exécution, utilisable pour la recherche sur les modèles de coût fondés sur l’apprentissage
- Pour réduire la difficulté d’entraîner un modèle sur de grands graphes de calcul d’un seul coup, Graph Segment Training découpe les graphes en segments, diminue l’usage mémoire et réduit par 3 le temps d’entraînement
- Lors d’un concours Kaggle, 792 participants répartis en 616 équipes et issus de 66 pays ont validé des techniques concrètes d’amélioration des modèles de prédiction de coût : compression de graphes, ajustement des valeurs de padding, ajout de caractéristiques de nœuds, attention entre configurations, etc.
Pourquoi les compilateurs ML déterminent les performances d’exécution
- Les modèles ML récents réalisent des tâches comme la compréhension du langage naturel, la conversation, la génération d’images et la génération de vidéos, et sont écrits et entraînés avec des frameworks de programmation ML comme TensorFlow, JAX et PyTorch
- Ces frameworks fournissent des opérations d’algèbre linéaire comme la multiplication de matrices et la convolution, ainsi que des couches de réseaux de neurones comme les couches de convolution 2D et les couches Transformer
- Les utilisateurs n’ont pas besoin de gérer directement les détails de l’exécution efficace des modèles sur le matériel ; le compilateur situé sous le framework optimise automatiquement le modèle
- Toutefois, les compilateurs résolvent souvent des problèmes d’optimisation complexes au moyen d’heuristiques, ce qui rend difficile l’obtention systématique des performances optimales
Graphes de calcul et optimisation en deux étapes
- Un compilateur ML transforme les instructions mathématiques écrites par l’utilisateur en instructions exécutables sur le matériel réel
- Un programme ML peut être représenté sous forme de graphe de calcul
- Les nœuds représentent des opérations sur des tenseurs, comme une multiplication de matrices
- Les arêtes représentent les tenseurs qui circulent d’un nœud à un autre
- Les optimisations du compilateur se divisent globalement en deux catégories
- Optimisation au niveau du graphe : prend des décisions en tenant compte du contexte de l’ensemble du graphe et transforme le graphe dans son ensemble
- Optimisation au niveau du noyau : transforme un noyau, c’est-à-dire un fused subgraph, indépendamment des autres noyaux
Compromis de performance liés à la disposition mémoire
- Un tenseur 2D comme une matrice peut être stocké en mémoire sous la forme
[A B C a b c]ou[A a B b C c], ce qui correspond respectivement aux dispositions row-major et column-major - L’une des optimisations importantes d’un compilateur ML consiste à attribuer une disposition mémoire à tous les tenseurs intermédiaires du programme
- Une disposition donnée peut être la plus efficace pour une opération individuelle, mais si les dispositions ne correspondent pas entre une addition et une convolution, le compilateur doit insérer une opération de copie supplémentaire
- À l’inverse, une configuration qui offre des performances légèrement inférieures pour les opérations individuelles, mais ne nécessite pas de conversion de disposition, peut être meilleure pour l’exécution globale
- Dans la suite de benchmarks XLA, un gain de vitesse allant jusqu’à 32 % a été observé lorsque la configuration de disposition optimale était choisie à la place des paramètres par défaut du compilateur
Jeu de données TpuGraphs
- TpuGraphs est un jeu de données de modèles de coût fondés sur l’apprentissage pour des programmes exécutés sur les TPU custom de Google
- L’objectif est d’entraîner un modèle de coût qui reçoit en entrée un programme et des paramètres de compilateur, puis prédit le temps d’exécution du programme
- Le jeu de données couvre deux paramètres du compilateur XLA
- layout : paramètre qui généralise les notions row-major et column-major des matrices aux tenseurs de dimension supérieure
- tiling : paramètre de taille des tuiles
- Chaque exemple contient le graphe de calcul d’une charge de travail ML, des paramètres de compilation et le temps d’exécution lorsque le programme est compilé avec ces paramètres
- Les graphes ont été collectés à partir de programmes ML open source et incluent des architectures de modèles comme ResNet, EfficientNet, Mask R-CNN et Transformer
- Les instructions de téléchargement et le code de démarrage sont disponibles sur TpuGraphs GitHub
- TpuGraphs contient 25 fois plus de graphes que le plus grand jeu de données existant de prédiction de propriétés de graphes de taille comparable, et sa taille moyenne de graphe est 770 fois supérieure à celle des jeux de données existants de prédiction des performances de programmes ML
Modèle de coût de référence et architecture GNN
- TpuGraphs fournit aussi un modèle de coût de référence fondé sur l’apprentissage et utilise un GNN, car le programme en entrée est représenté sous forme de graphe
- Les caractéristiques de nœuds se composent de deux parties
- opcode id : l’information de nœud la plus importante, qui représente le type d’opération sur tenseur
- Autres caractéristiques de nœuds
- Le modèle de référence transforme l’opcode id en opcode embedding via une table d’embeddings
- L’opcode embedding et les autres caractéristiques de nœuds sont combinés puis utilisés comme entrée du GNN
- Les embeddings de nœuds produits par le GNN sont combinés en un embedding de graphe de taille fixe au moyen d’opérations simples de graph pooling reduction, comme sum et mean
- L’embedding final du graphe passe par une couche feedforward et est transformé en une sortie scalaire unique
Entraîner de grands graphes avec Graph Segment Training
- Graph Segment Training est une technique de passage à l’échelle pour l’entraînement de GNN destinée à traiter de grands graphes sur des appareils à capacité mémoire limitée
- Cette méthode vise les situations de prédiction au niveau du graphe, où la cible de prédiction n’est ni un nœud ni une arête, mais le graphe entier
- Les graphes de calcul pouvant contenir des centaines de milliers de nœuds, le Full Graph Training, qui utilise le graphe entier d’un seul coup, peut être impossible sur le plan computationnel
- GST découpe les grands graphes en petits segments et ne sélectionne qu’un sous-ensemble aléatoire de segments pour mettre à jour le modèle
- Les autres segments génèrent des embeddings sans stocker les activations intermédiaires, ce qui réduit l’usage mémoire
- Tous les embeddings de segments sont combinés pour former l’embedding du grand graphe d’origine, puis utilisés pour la prédiction
- Une historical embedding table et du segment dropout sont également introduits afin d’atténuer le staleness des historical embeddings
- La méthode complète réduit le temps d’entraînement end-to-end par 3
Techniques d’amélioration validées lors du concours Kaggle
- Le concours Kaggle Fast or Slow? Predict AI Model Runtime, basé sur le jeu de données TpuGraphs, a réuni 792 participants répartis en 616 équipes et issus de 66 pays
- Il y a eu 10 507 soumissions ; pour 153 participants, il s’agissait de leur premier concours Kaggle, et 47 d’entre eux ont terminé dans le top 100
- Les équipes participantes ont expérimenté plusieurs techniques
- Pruning et compression de graphes : au lieu de GST, elles ont testé des méthodes de compression des grands graphes ; une approche consistant à ne conserver que le subgraph incluant les nœuds configurables et leurs voisins immédiats a été utilisée
- Changement de la valeur de padding : comme la valeur de padding par défaut, 0, entre en conflit avec des valeurs de features valides, utiliser -1 pouvait améliorer fortement la précision du modèle
- Ajout de caractéristiques de nœuds et modification de l’encodage : des caractéristiques de nœuds supplémentaires, comme les contracting dimensions de dot general, sont importantes, et la méthode d’encodage des caractéristiques de nœuds peut aussi influencer les résultats
- cross-configuration attention : l’équipe gagnante a conçu une couche simple permettant au modèle de comparer explicitement les configurations, ce qui a donné des résultats nettement meilleurs qu’une approche où chaque configuration est inférée séparément
- Les résultats du concours et la solution gagnante devaient être présentés le 16 décembre 2023 lors de la session concours du workshop ML for Systems à NeurIPS
Session liée à NeurIPS Expo
- Pour les lecteurs intéressés par les données structurées et la recherche en intelligence artificielle, le panel NeurIPS Expo Graph Learning Meets Artificial Intelligence a lieu le 9 décembre 2023
- Ce panel traite notamment des progrès des modèles de coût fondés sur l’apprentissage
1 commentaires
Commentaires sur Hacker News
Les compilateurs ML sont survendus. C’est le même compromis que pour les compilateurs traditionnels : on obtient bien plus de débit que si l’on embauche un programmeur spécialiste des performances, mais ce dernier est généralement bien plus rapide, parfois avec plusieurs ordres de grandeur d’écart
Il y a des insuffisances à plusieurs niveaux. Au niveau algorithmique, cela ne renvoie pas vraiment à l’humain les astuces qui rendraient le réseau plus rapide, seulement des signaux très basiques. L’intention se perd aussi. Le concepteur d’un réseau ML décrit l’architecture en Python, mais après plusieurs étapes de transformation descendante, on peut aboutir à un résultat complètement aberrant. Récemment, j’ai vu un compilateur faire une mise à jour par tranche en construisant toute la plage d’indices possible du tableau, puis en la découpant pour obtenir l’indice à mettre à jour, avant de faire un scatter ; je l’ai remplacé par un simple appel à
memcpy. Les kernels sont eux aussi inefficaces. Chaque fois que ce genre de sortie de compilateur affronte un programmeur assembleur expérimenté, le compilateur perd, généralement avec plus de 30 % d’écart. Cela semble facile à corriger, mais si personne n’a vraiment réussi à résoudre ça en 50 ans, c’est sans doute moins simple qu’on le ditLes compilateurs ont eux aussi beaucoup d’heuristiques, comme l’inlining, le déroulage de boucle ou la vectorisation ; un réseau neuronal pourrait donc aider, et être plus facile à maintenir qu’un grand nombre d’heuristiques écrites à la main
Même ainsi, c’est toujours bien mieux que de ne pas avoir du tout ce type d’optimisation
Quelqu’un peut expliquer ça de manière plus concrète ? J’aimerais savoir où en sont réellement les compilateurs ML aujourd’hui, et ce qu’on peut raisonnablement attendre à court terme
Il suffit d’écrire
model = torch.compile(model). « Sur 163 modèles open source, torch.compile a fonctionné dans 93 % des cas et a accéléré l’entraînement de 43 % sur des GPU NVIDIA A100. Nous avons observé une accélération moyenne de 21 % en précision Float32 et de 51 % en précision AMP. »[1] Google semble vouloir faire en sorte que davantage de personnes participent à la R&D sur ce type d’approche[1] https://pytorch.org/get-started/pytorch-2.0/
En simplifiant beaucoup une réalité complexe, l’idée consiste à compiler le graphe de calcul vers une représentation intermédiaire, puis à implémenter un backend adapté. Parmi les projets liés, on peut regarder stableHLO, IREE et openXLA. Le compilateur
jitde Jax peut aussi être vu comme une forme de ce type de compilateur. Il abaisse les opérations tracées vers XLA, puis XLA fait encore toute une série de tours de magie pour les exécuter sur le backend. Au final, plus on descend, plus ce sont des transformations et abstractionstorch.compileEn résumé, il s’agit d’améliorer la prédiction des performances en temps d’exécution d’un graphe de calcul avec un réseau neuronal de graphes (GNN). On utilise à la fois un dictionnaire d’embeddings pour les opcodes de chaque nœud et d’autres caractéristiques de nœud comme la shape, les bits ou la taille de fenêtre ([1])
Ils ont publié dans [2] un grand jeu de données de graphes contenant les performances obtenues sur TPU avec différentes configurations de compilation XLA, et dans [3] ils améliorent la prédiction sur des graphes plus grands grâce à une méthode de découpage de graphe (je voyais pour la première fois une partition de graphe METIS) et à diverses techniques d’apprentissage. Il s’agit ici de prédire les performances d’un graphe donné, pas d’améliorer, proposer ou modifier un nouveau graphe équivalent. Comme FunSearch, un modèle ayant un pouvoir prédictif correct peut être utilisé avec de la recherche évolutionnaire
[1] https://github.com/google-research-datasets/tpu_graphs#featu...
[2] TpuGraphs: A Performance Prediction Dataset on Large Tensor Computational Graphs https://arxiv.org/abs/2308.13490
[3] Learning Large Graph Property Prediction via Graph Segment Training https://arxiv.org/abs/2305.12322
Quelqu’un peut expliquer comment la convolution fonctionne dans ce graphe ? On applique une convolution avec un kernel de shape
[4,16,8]à un tenseur de shape[2,4,16], et on obtient un tenseur[2,8]; comment est-ce possible ?[2,4,16],2peut être vu comme la taille de batch,4comme la dimension des caractéristiques d’entrée, et16comme la dimension des canaux d’entréeDans le kernel
[4,16,8],4est la taille de la fenêtre du filtre,16correspond à la dimension des canaux d’entrée, et8est la dimension des canaux de sortie. Dans la sortie[2,8],2reste la taille de batch, et8correspond à la dimension des canaux de sortie du kernel. À première vue, les dimensions semblent incompatibles, mais dans un graphe, la convolution exploite la structure de voisinage. Le kernel se déplace sur le graphe et applique des poids au nœud courant ainsi qu’aux caractéristiques des voisins dans un certain rayon, puis agrège ces sommes pondérées pour produire de nouvelles caractéristiques pour chaque canal de sortie. La structure du graphe, les poids des arêtes, ainsi que des détails d’implémentation comme le padding ou le stride peuvent aussi influer sur la shape de sortieOù en est Gemini ?
Rien que parmi ceux qui me viennent à l’esprit, il y a au moins cinq modèles de base — Llama, Claude, Gemini, Falcon, Mistral — qui se dépassent tour à tour, mais GPT reste encore un cran au-dessus, et cela dure depuis un an. Les grands modèles de langage basés sur les Transformers se sont révélés suffisamment simples pour que n’importe qui, avec un million de dollars de temps GPU, puisse en construire un, mais cela ne suffit pas à rattraper complètement OpenAI. Quel est leur ingrédient secret ?
Et le Transformer lui-même ? Y a-t-il des indices suggérant qu’il est optimal dans un certain sens ?
J’ai l’impression que le point essentiel a été noyé dans le premier paragraphe, mais le reste est excellent
La vitesse actuelle des progrès du ML est stupéfiante. Je ne crois pas à la singularité, mais cela transforme déjà le logiciel et la société d’une manière que personne n’avait prévue
https://deepmind.google/discover/blog/funsearch-making-new-d...