Fortschritte bei Machine-Learning-Techniken für Machine Learning

• Je größer ML-Modelle werden, desto stärker hängt ihre Ausführungseffizienz nicht nur von der Hardware, sondern auch von Compiler-Optimierungen ab; Google Research und Google DeepMind haben einen Ansatz vorgestellt, der diesen Bereich wiederum mit ML verbessert
• Compiler-Entscheidungen unter Frameworks wie TensorFlow, JAX und PyTorch können selbst beim selben Modell Laufzeit und Ressourcennutzung erheblich verändern
• TpuGraphs ist ein Datensatz für Forschungsarbeiten zu lernbasierten Kostenmodellen; er enthält Berechnungsgraphen von ML-Programmen für TPUs, Compiler-Einstellungen und Laufzeiten
• Um die Einschränkung zu verringern, dass große Berechnungsgraphen schwer auf einmal zu trainieren sind, teilt Graph Segment Training Graphen in Segmente auf, senkt so den Speicherbedarf und verkürzt die Trainingszeit um den Faktor 3
• In einem Kaggle-Wettbewerb prüften 792 Teilnehmende in 616 Teams aus 66 Ländern praktische Techniken zur Verbesserung von Kostenprognosemodellen, darunter Graph-Kompression, Anpassung von Padding-Werten, zusätzliche Node-Features und Attention zwischen Konfigurationen

---

Warum ML-Compiler die Ausführungsleistung bestimmen

• Moderne ML-Modelle erledigen Aufgaben wie Sprachverständnis, Dialog, Bildgenerierung und Videogenerierung und werden mit ML-Programmierframeworks wie TensorFlow, JAX und PyTorch geschrieben und trainiert
• Frameworks stellen lineare Algebraoperationen wie Matrixmultiplikation und Faltungen sowie neuronale Netzwerkschichten wie 2D-Convolution-Layer und Transformer-Layer bereit
• Nutzer müssen sich nicht selbst um die Details kümmern, wie ein Modell effizient auf Hardware ausgeführt wird; der Compiler unterhalb des Frameworks optimiert das Modell automatisch
• Allerdings lösen Compiler komplexe Optimierungsprobleme häufig mit Heuristiken, weshalb sie nicht immer die optimale Leistung erzielen

Berechnungsgraphen und zweistufige Optimierung

• ML-Compiler wandeln mathematische Anweisungen, die Nutzer geschrieben haben, in Befehle um, die auf realer Hardware ausgeführt werden können
• ML-Programme lassen sich als Berechnungsgraphen darstellen
  • Nodes stehen für Tensoroperationen wie matrix multiplication
  • Edges stehen für Tensoren, die von einem Node zu einem anderen fließen
• Compiler-Optimierungen lassen sich grob in zwei Arten einteilen
  • Optimierung auf Graph-Ebene: trifft Entscheidungen im Kontext des gesamten Graphen und transformiert den gesamten Graphen
  • Optimierung auf Kernel-Ebene: transformiert einen einzelnen Kernel, also einen fused subgraph, unabhängig von anderen Kerneln

Performance-Trade-offs beim Speicherlayout

• 2D-Tensoren wie Matrizen können im Speicher in der Form [A B C a b c] oder [A a B b C c] abgelegt werden; dies entspricht jeweils Row-Major- und Column-Major-Layout
• Eine wichtige Optimierung von ML-Compilern besteht darin, allen Zwischentensoren eines Programms ein Speicherlayout zuzuweisen
• Ein bestimmtes Layout kann für eine einzelne Operation am effizientesten sein, doch wenn die Layouts zwischen add und convolution nicht zusammenpassen, muss der Compiler zusätzliche Copy-Operationen einfügen
• Umgekehrt kann eine Konfiguration, die bei einzelnen Operationen etwas weniger performant ist, in der Gesamtausführung besser sein, wenn keine Layout-Transformation nötig ist
• In der XLA Benchmark Suite wurde bei der Wahl optimaler Layout-Konfigurationen statt der Standard-Compiler-Einstellungen eine Beschleunigung von bis zu 32 % beobachtet

TpuGraphs-Datensatz

• TpuGraphs ist ein Datensatz für lernbasierte Kostenmodelle für Programme, die auf Googles kundenspezifischen TPUs ausgeführt werden
• Ziel ist es, ein Kostenmodell zu trainieren, das ein Eingabeprogramm und Compiler-Einstellungen entgegennimmt und die Laufzeit des Programms vorhersagt
• Der Datensatz deckt zwei XLA-Compiler-Einstellungen ab
  • layout: eine Einstellung, die das Konzept von Row-Major und Column-Major bei Matrizen auf höherdimensionale Tensoren verallgemeinert
  • tiling: Einstellungen für die Kachelgröße
• Jedes Beispiel enthält den Berechnungsgraphen eines ML-Workloads, Compiler-Einstellungen und die Laufzeit bei Kompilierung mit diesen Einstellungen
• Die Graphen wurden aus Open-Source-ML-Programmen gesammelt und umfassen Modellarchitekturen wie ResNet, EfficientNet, Mask R-CNN und Transformer
• Download-Anleitung und Einstiegscode sind auf TpuGraphs GitHub verfügbar
• TpuGraphs enthält 25-mal mehr Graphen als der bisher größte Datensatz für Graph Property Prediction mit ähnlicher Graphgröße, und die durchschnittliche Graphgröße ist 770-mal größer als bei bisherigen Datensätzen zur Performance-Vorhersage von ML-Programmen

Baseline-Kostenmodell und GNN-Architektur

• TpuGraphs liefert ein lernbasiertes Baseline-Kostenmodell mit; da das Eingabeprogramm als Graph dargestellt wird, nutzt es ein GNN
• Node-Features bestehen aus zwei Teilen
  • opcode id: die wichtigste Node-Information, die die Art der Tensoroperation angibt
  • weitere Node-Features
• Das Baseline-Modell wandelt die opcode id über eine Embedding Lookup Table in ein opcode embedding um
• Das opcode embedding wird mit den übrigen Node-Features kombiniert und als Eingabe für das GNN verwendet
• Die vom GNN erzeugten Node-Embeddings werden mit einfachen Graph-Pooling-Reduktionen wie sum und mean zu einem Graph-Embedding fester Größe kombiniert
• Das finale Graph-Embedding wird über eine Feedforward Layer in eine einzelne skalare Ausgabe umgewandelt

Große Graphen mit Graph Segment Training trainieren

• Graph Segment Training ist eine Technik zur Skalierung des GNN-Trainings, um große Graphen auf Geräten mit begrenztem Speicher zu verarbeiten
• Die Methode zielt auf Graph-Level-Prediction-Szenarien ab, bei denen nicht Nodes oder Edges, sondern der gesamte Graph vorhergesagt wird
• Berechnungsgraphen können Hunderttausende Nodes enthalten, sodass Full Graph Training, bei dem der gesamte Graph auf einmal verwendet wird, rechnerisch unmöglich sein kann
• GST teilt große Graphen in kleinere Segmente auf und aktualisiert das Modell nur anhand einer zufälligen Teilmenge der Segmente
• Für die übrigen Segmente werden Embeddings erzeugt, ohne Zwischen-Aktivierungen zu speichern, wodurch der Speicherbedarf sinkt
• Alle Segment-Embeddings werden kombiniert, um das Embedding des ursprünglichen großen Graphen zu erstellen; dieses wird für die Vorhersage genutzt
• Zusammen mit einer historical embedding table und segment dropout wird die Staleness historischer Embeddings abgeschwächt
• Insgesamt verkürzt die Methode die End-to-End-Trainingszeit um den Faktor 3

Im Kaggle-Wettbewerb geprüfte Verbesserungstechniken

• Der Kaggle-Wettbewerb Fast or Slow? Predict AI Model Runtime basierte auf dem TpuGraphs-Datensatz; 792 Teilnehmende in 616 Teams aus 66 Ländern nahmen teil
• Es gab 10.507 Einreichungen; für 153 Personen war es der erste Kaggle-Wettbewerb, und 47 von ihnen kamen unter die Top 100
• Die teilnehmenden Teams experimentierten mit mehreren Techniken
  • Graph Pruning und Compression: Statt GST wurden Methoden zur Komprimierung großer Graphen erprobt; verwendet wurde ein Ansatz, der nur den Subgraphen mit konfigurierbaren Nodes und deren unmittelbaren Nachbarn beibehält
  • Änderung des Padding-Werts: Da der Standard-Padding-Wert 0 mit gültigen Feature-Werten kollidiert, konnte die Verwendung von -1 die Modellgenauigkeit deutlich verbessern
  • Zusätzliche Node-Features und geänderte Kodierung: Zusätzliche Node-Features wie die contracting dimensions von dot general sind wichtig, und auch die Kodierungsweise von Node-Features kann das Ergebnis beeinflussen
  • Cross-Configuration Attention: Das Siegerteam entwarf eine einfache Layer, mit der das Modell Konfigurationen explizit vergleichen kann; dies lieferte deutlich bessere Ergebnisse als ein Ansatz, der jede Konfiguration einzeln inferiert
• Die Wettbewerbsergebnisse und die Gewinnerlösung sollten am 16. Dezember 2023 in der Wettbewerbs-Session des ML for Systems workshop auf der NeurIPS behandelt werden

Session zur NeurIPS Expo

• Für Leser, die sich für strukturierte Daten und KI-Forschung interessieren, fand am 9. Dezember 2023 das NeurIPS-Expo-Panel Graph Learning Meets Artificial Intelligence statt
• Dieses Panel behandelt unter anderem Fortschritte bei lernbasierten Kostenmodellen

[GN⁺]

Meinungen auf Hacker News

• ML-Compiler werden überbewertet. Es ist derselbe Kompromiss wie bei klassischen Compilern: Man bekommt viel mehr Durchsatz, als wenn man spezialisierte Performance-Programmierer einstellt, aber Letztere sind normalerweise deutlich schneller und können in manchen Fällen um Größenordnungen vorn liegen.
  Auf mehreren Ebenen unzureichend. Auf algorithmischer Ebene geben sie dem Menschen keine Tricks zurück, wie man das Netzwerk schneller machen kann, sondern nur sehr grundlegende Signale. Auch die Intention geht verloren. ML-Netzwerk-Designer spezifizieren die Struktur in Python, aber nach mehreren Stufen der Herabtransformation kann etwas völlig anderes herauskommen. Kürzlich habe ich gesehen, wie ein Compiler bei einem Slice-Update alle möglichen Indexbereiche eines Arrays erzeugte, diese dann zuschnitt, um die zu aktualisierenden Indizes zu erhalten, und anschließend einen Scatter ausführte; ich habe das durch einen einzigen memcpy-Aufruf ersetzt. Auch die Kernel sind ineffizient. Jedes Mal, wenn solcher Compiler-Output gegen einen erfahrenen Assembly-Programmierer antritt, verliert der Compiler, meist um 30 % oder mehr. Es sieht so aus, als ließe sich das leicht lösen, aber wenn es in den letzten 50 Jahren niemand richtig geschafft hat, ist es offensichtlich nicht so einfach, wie es klingt.

  • Wenn man sich die Schach-Engine Stockfish ansieht: Bei der Brettbewertung hat sie jahrelang von Menschen geschriebene Heuristiken verworfen und durch ein kleines neuronales Netz ersetzt, und das funktioniert besser.
    Auch Compiler haben viele Heuristiken wie Inlining, Loop Unrolling und Vektorisierung; neuronale Netze könnten also helfen und womöglich leichter zu warten sein als die vielen von Menschen geschriebenen Heuristiken.
  • Du sagst, es sei derselbe Kompromiss wie bei klassischen Compilern; interessant ist dabei, dass diese klassischen Compiler unglaublich nützlich waren.
  • Das wirkt zu apodiktisch und verschlossen.
  • Genau. Warum sollte jemand gcc/clang benutzen, wenn man einfach jemanden einstellen kann, der Assembly von Hand schreibt?
  • Dieser Durchsatz ist der entscheidende Punkt. Man kann nicht für jede ML-Aufgabe einen Performance-Experten abstellen.
    Solche Optimierungen sind immer noch viel besser, als gar keine zu haben.

• Kann das jemand etwas realistischer erklären? Mich interessiert, wie der tatsächliche Stand heutiger ML-Compiler ist und was man in naher Zukunft erwarten kann.

  • Einer der einfachsten Ansätze ist torch.compile. Das ist die neueste Iteration des PyTorch-Compilers; frühere Ansätze waren TorchScript und FX Tracing.
    Man schreibt einfach model = torch.compile(model). „Über diese 163 Open-Source-Modelle hinweg funktionierte torch.compile in 93 % der Fälle und beschleunigte das Training auf NVIDIA-A100-GPUs um 43 %. Bei Float32-Präzision lag die durchschnittliche Beschleunigung bei 21 %, bei AMP-Präzision bei 51 %.“[1] Google scheint zu versuchen, mehr Leute in die Forschung und Entwicklung solcher Methoden einzubeziehen.
    [1] https://pytorch.org/get-started/pytorch-2.0/
  • Kurzfristig ist zu erwarten, dass man AMD, CUDA, TPU, CPU usw. nutzen kann, ohne explizite Vendor-Unterstützung für das Framework zu haben, in dem ein Modell entwickelt wurde.
    Die Realität ist kompliziert, aber stark vereinfacht gesagt kompiliert man den Berechnungsgraphen in eine Zwischenrepräsentation und implementiert dafür passende Backends. Relevante Projekte sind stableHLO, IREE und openXLA. Auch der jit-Compiler von Jax kann als eine Form eines solchen Compilers betrachtet werden. Er senkt die getraceten Operationen nach XLA ab, und XLA vollführt dann allerlei Magie, damit sie auf dem Backend laufen. Letztlich sind es nach unten hin immer weiter Transformationen und Abstraktionen.
  • Sieh dir torch.compile an.

• Zusammengefasst geht es darum, die Laufzeit-Performance-Prognose von Berechnungsgraphen mit Graph Neural Networks (GNNs) zu verbessern. Verwendet werden ein Embedding-Wörterbuch für den Opcode jedes Knotens sowie weitere Knotenmerkmale wie Shape, Bits und Window Size ([1]).
  In [2] wurde ein großer Graph-Datensatz mit verschiedenen XLA-Compiler-Einstellungen und der daraus resultierenden Performance auf TPUs veröffentlicht; in [3] verbessern sie mit einer Methode zur Graph-Partitionierung (METIS graph partition sehe ich zum ersten Mal) und verschiedenen Lernverfahren die Vorhersage für größere Graphen gegenüber früheren Ansätzen. Hier geht es darum, die Performance eines gegebenen Graphen vorherzusagen, nicht darum, neue äquivalente Graphen zu verbessern, vorzuschlagen oder zu verändern. Wie bei FunSearch kann ein Modell mit ordentlicher Vorhersagekraft zusammen mit evolutionärer Suche verwendet werden.
  [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

• Kann jemand erklären, wie Convolution in diesem Graphen funktioniert? Ein Tensor mit Shape [2,4,16] wird mit einem Kernel der Shape [4,16,8] gefaltet und es kommt ein Tensor [2,8] heraus – wie ist das möglich?

  • Ich weiß nicht, ob das hilft, aber beim Eingabetensor [2,4,16] kann 2 als Batch-Größe, 4 als Eingabe-Feature-Dimension und 16 als Eingabe-Channel-Dimension verstanden werden.
    Beim Kernel [4,16,8] ist 4 die Größe des Filterfensters, 16 entspricht der Eingabe-Channel-Dimension, und 8 ist die Ausgabe-Channel-Dimension. In der Ausgabe [2,8] bleibt 2 die Batch-Größe, und 8 entspricht der Ausgabe-Channel-Dimension des Kernels. Oberflächlich sehen die Dimensionen unpassend aus, aber Convolution in einem Graphen nutzt die Nachbarschaftsstruktur. Der Kernel bewegt sich über den Graphen, wendet Gewichte auf den aktuellen Knoten und die Nachbar-Features innerhalb eines bestimmten Radius an und sammelt die gewichteten Summen, um neue Features für jeden Ausgabe-Channel zu erzeugen. Auch Graphstruktur, Kantengewichte sowie Implementierungsdetails wie Padding und Stride können die Output-Shape beeinflussen.

• Wie ist der Stand bei Gemini?

  • Interessant ist, dass GPT-4 weiterhin vorne liegt: https://twitter.com/lmsysorg/status/1735729398672716114
    Allein spontan fallen mindestens fünf Foundation Models ein – Llama, Claude, Gemini, Falcon, Mistral –, die sich gegenseitig überholen, aber GPT ist weiterhin eine Stufe darüber, und das seit einem Jahr. Es hat sich gezeigt, dass Transformer-basierte große Sprachmodelle einfach genug sind, dass sie jeder bauen kann, der etwa eine Million Dollar für GPU-Zeit ausgeben kann, aber ganz zu OpenAI aufschließen konnten sie nicht. Was ist ihr besonderes Geheimnis?

• Wie sieht es mit Transformer selbst aus? Gibt es Hinweise darauf, in welchem Sinne sie optimal sind?

• Ich habe das Gefühl, dass im ersten Absatz der Kernpunkt vergraben wurde, aber der Rest ist großartig.

• Das aktuelle Tempo des ML-Fortschritts ist erstaunlich. Ich glaube nicht an die Singularität, aber es verändert Software und Gesellschaft auf eine Weise, die niemand vorhersagen kann.

  • Wenn man sich das hier und FunSearch ansieht, wirkt es, als stünde die Singularität kurz bevor.
    https://deepmind.google/discover/blog/funsearch-making-new-d...
  • Für mich sieht es nach einem weiteren Goldrausch nach Dotcom, Mobile, Cloud und VR aus.
  • Ich glaube nicht, dass Menschen in fünf Jahren noch so programmieren werden wie heute.
  • Ich möchte zuerst sehen, dass es Heilmittel für schwer behandelbare Krankheiten liefert. Die Singularität an sich ist bedeutungslos, wenn sie den Menschen keinen Nutzen bringt, und dieser Nutzen sollte vor allem in besserer Gesundheit und weniger Leid bestehen.