TpuGraphs: Ein Datensatz zur Leistungsvorhersage von großen Tensor-Rechengraphen

Präzise Hardware-Leistungsmodelle spielen eine entscheidende Rolle bei Code-Optimierungen. Sie können Compilern bei heuristischen Entscheidungen helfen oder Autotuner bei der Identifizierung der optimalen Konfiguration für ein gegebenes Programm unterstützen. Zum Beispiel entdeckte der Autotuner für XLA, einen Compiler für maschinelles Lernen, eine Geschwindigkeitssteigerung von 10-20 % bei den neuesten Modellen, die erheblichen Produktionsverkehr bei Google bedienen. Obwohl es einige Datensätze zur Leistungsvorhersage von Programmen gibt, richten sie sich auf kleine Unterprogramme wie Basisblöcke oder Kernels. In diesem Artikel wird TpuGraphs vorgestellt, ein Datensatz zur Leistungsvorhersage vollständiger Tensor-Programme, die als Rechengraphen auf Tensor Processing Units (TPUs) ausgeführt werden. Jeder Graph im Datensatz repräsentiert die Hauptberechnung einer maschinellen Lerntaufe, z.B. einen Trainingszyklus oder einen Inferenzschritt. Jedes Datenbeispiel enthält einen Rechengraphen, eine Kompilierungskonfiguration und die Ausführungszeit des Graphen nach Kompilierung mit dieser Konfiguration. Die Graphen im Datensatz stammen aus Open-Source-Maschinenlernprogrammen und umfassen bekannte Modellarchitekturen wie ResNet, EfficientNet, Mask R-CNN und Transformer. TpuGraphs bietet 25-mal mehr Graphen als der größte Datensatz zur Vorhersage von Grapheneigenschaften (mit vergleichbaren Graphengrößen) und ist durchschnittlich 770-mal größer als bestehende Datensätze zur Leistungsvorhersage von Maschinenlernprogrammen. Diese graphbasierte Vorhersageaufgabe für große Graphen wirft neue Herausforderungen in Bezug auf Skalierbarkeit, Trainings-effizienz und Modellqualität auf.