TpuGraphs : Un jeu de données pour la prédiction des performances sur de grands graphes de calcul tensoriel

Les modèles de performance matérielle précis jouent un rôle crucial dans les optimisations de code. Ils peuvent aider les compilateurs à prendre des décisions heuristiques ou assister les autotuneurs à identifier la configuration optimale pour un programme donné. Par exemple, l'autotuneur d'XLA, un compilateur d'apprentissage automatique, a découvert une accélération de 10 à 20 % sur des modèles de pointe servant un trafic de production considérable chez Google. Bien que quelques ensembles de données existent pour la prédiction de la performance des programmes, ils ciblent principalement de petits sous-programmes tels que des blocs basiques ou des noyaux. Cet article introduit TpuGraphs, un ensemble de données pour la prédiction de la performance sur des programmes tensoriels complets, représentés sous forme de graphes computationnels, s'exécutant sur des Unités de Traitement Tensoriel (TPUs). Chaque graphe dans l'ensemble de données représente le calcul principal d'une charge de travail d'apprentissage automatique, par exemple une époque d'entraînement ou une étape d'inférence. Chaque échantillon de données contient un graphe computationnel, une configuration de compilation et le temps d'exécution du graphe lorsqu'il est compilé avec cette configuration. Les graphes dans l'ensemble de données sont collectés à partir de programmes d'apprentissage automatique open source, mettant en œuvre des architectures de modèles populaires telles que ResNet, EfficientNet, Mask R-CNN et Transformer. TpuGraphs fournit 25 fois plus de graphes que le plus grand ensemble de données pour la prédiction des propriétés graphiques (de taille comparable), et les graphes y sont en moyenne 770 fois plus grands comparativement aux ensembles de données existants pour la prédiction de la performance sur des programmes d'apprentissage automatique. Cette tâche de prédiction au niveau des graphes sur des grands graphes soulève de nouveaux défis en apprentissage, allant du problème d'évolutivité à l'efficacité du entraînement et jusqu'à la qualité du modèle.