Tensor Seeks Layout: Formalizing Layout Selection for ML Compilers
Diese Arbeit präsentiert die erste formale Untersuchung der Layout-Auswahl in Machine-Learning-Compilern, indem sie diese als kombinatorisches Optimierungsproblem formuliert, deren rechnerische Komplexität nachweist und sowohl optimale Algorithmen für Graphen mit beschränktem Treewidth als auch eine gewichtete MaxSAT-Kodierung für allgemeine Instanzen vorschlägt, um zu demonstrieren, dass einfache Heuristiken die Leistung im Vergleich zu optimalen Lösungen um das bis zu 5-fache verschlechtern können.
Originalarbeit lizenziert unter CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). Dies ist eine KI-generierte Erklärung des untenstehenden Papers. Sie wurde nicht von den Autoren verfasst oder gebilligt. Für technische Genauigkeit konsultieren Sie das Originalpaper. Vollständigen Haftungsausschluss lesen
Moderne künstliche Intelligenz stützt sich auf massive mathematische Modelle, die riesige Mengen an Daten verarbeiten, um Sprache zu erkennen, Sprachen zu übersetzen oder Bilder zu generieren. Um diese Modelle schnell auszuführen, nutzen Ingenieure spezialisierte Computerchips, die speziell für diese Schwerstarbeit entwickelt wurden. Diese Chips führen jedoch nicht einfach nur Befehle aus; sie müssen auch Daten effizient bewegen. Die Geschwindigkeit, mit der ein Modell läuft, hängt oft weniger von der Rohleistung des Chips ab als vielmehr davon, wie die Daten in seinem Speicher angeordnet sind. Stellen Sie sich eine Bibliothek vor, in der Bücher in Regalen gelagert werden. Wenn ein Leser ein bestimmtes Set von Büchern finden möchte, hängt die Zeit, die er dafür benötigt, vollständig davon ab, ob diese Bücher über verschiedene Gänge verstreut oder ordentlich gruppiert in einem einzigen Regal stehen. In der Welt der Computerchips wird diese Anordnung als „Layout“ bezeichnet. Wenn ein Computerprogramm eine Berechnung durchführt, erwartet es die Daten in einer ganz bestimmten Anordnung, aber der vorherige Schritt im Programm könnte die Daten in einer anderen Anordnung hinterlassen haben. Wenn diese beiden nicht übereinstimmen, muss der Computer stoppen und die Daten neu anordnen, bevor er fortfahren kann – ein Prozess, der Zeit und Energie verschwendet.
Jahrelang verließ sich die Software, die diese Modelle für die Chips vorbereitet, auf eine Sammlung von groben Schätzungen und Faustregeln, um die Anordnung dieser Daten zu entscheiden. Diese Regeln funktionieren gut genug für einfache Aufgaben, aber da die Modelle komplexer wurden, begannen die Schätzungen zu versagen, was zu erheblichen Verlangsamungen führte. Ein Team von Forschern der Technischen Universität Wien und Amazon setzte sich zum Ziel, diesen Ansatz zu ändern. Anstatt sich auf Intuition zu verlassen, behandelten sie das Problem der Datenanordnung als ein formales mathematisches Rätsel. Sie entwickelten ein präzises Modell, das die exakten Kosten jeder möglichen Anordnung berechnet, einschließlich der Zeit, die für die Bewegung von Daten zwischen verschiedenen Formaten benötigt wird. Auf diese Weise konnten sie die eine beste Art und Weise bestimmen, die Daten für ein gegebenes Modell zu organisieren, anstatt darauf zu hoffen, dass ein Satz von Regeln „nah genug“ herankommt.
Die Forscher entdeckten, dass das Finden dieser perfekten Anordnung eine unglaublich schwierige Aufgabe ist. In der Sprache der Informatik ist das Problem so komplex, dass kein Computer es für jede mögliche Situation schnell lösen kann, insbesondere wenn die Modelle größer werden. Sie bewiesen, dass selbst für eine vereinfachte Version des Problems, die nur grundlegende Matrixberechnungen umfasst, die Anzahl der Möglichkeiten so gewaltig ist, dass ein Standardcomputer Schwierigkeiten hätte, die Antwort in einer angemessenen Zeit zu finden. Dieser Befund widerlegte die Idee, dass ein einziger, schneller, universeller Algorithmus das Problem für alle zukünftigen Modelle lösen könnte. Das Team fand jedoch auch einen Weg nach vorn. Sie zeigten, dass das Problem zwar im Allgemeinen schwierig ist, aber handhabbar wird, wenn die Struktur des Modells einem Baum mit begrenzter Verzweigung ähnelt. Für diese spezifischen Strukturen, die in vielen realen Anwendungen häufig vorkommen, entwickelten sie eine Methode, die die perfekte Lösung schnell findet. Für komplexere Strukturen, die nicht diesem Muster entsprechen, entwickelten sie einen Weg, das Problem in ein Format zu übertragen, das von bestehenden leistungsstarken Solvern verarbeitet werden kann, wodurch sie die beste Anordnung finden können, selbst wenn es keinen perfekten mathematischen Shortcut gibt.
Um ihre Ideen zu testen, implementierten die Forscher ihre neue Methode in einem realen Compiler, der für die Trainium-Chips von Amazon verwendet wird, welche darauf ausgelegt sind, Modelle der künstlichen Intelligenz auszuführen. Sie verglichen ihren neuen Ansatz mit den derzeit in der Industrie üblichen Standardmethoden, die auf den älteren Faustregeln basieren. Die Ergebnisse waren beeindruckend. Bei einigen komplexen Modellen, insbesondere jenen zur Bilderkennung, verursachten die alten Faustregeln, dass die Modelle bis zu fünfmal langsamer liefen als nötig. Dies geschah, weil die einfachen Regeln nicht das Gesamtbild sahen; sie ordneten Daten perfekt für einen Schritt an, schufen aber für den nächsten Schritt ein Chaos, was den Computer dazu zwang, ständig Zeit mit der Neuordnung der Daten zu verschwenden. Die neue Methode hingegen vermeidet diese kostspieligen Neuordnungen, indem sie die gesamte Sequenz der Schritte gleichzeitig betrachtet, und hält den Datenfluss reibungslos aufrecht.
Die Studie offenbarte jedoch auch eine entscheidende Einschränkung. Obwohl die neue Methode immer die mathematisch beste Anordnung gemäß ihren eigenen Berechnungen finden konnte, bedeutete dies nicht immer die schnellste Geschwindigkeit auf der tatsächlichen Hardware. In einigen Fällen lieferte die neue Methode ein Ergebnis, das theoretisch perfekt war, aber schlechter abschnitt als die älteren, einfacheren Regeln. Die Forscher führten diese Diskrepanz auf das Kostenmodell selbst zurück. Die Software, die vorhersagte, wie lange eine Aufgabe dauern würde, war nicht perfekt genau; sie unterschätzte die Zeit, die für bestimmte Arten der Datenbewegung benötigt wurde. Da die neue Methode so gut darin war, die niedrigsten Kosten gemäß ihrer eigenen fehlerhaften Vorhersagen zu finden, wählte sie manchmal eine Anordnung, die auf dem Papier günstig aussah, in der Realität jedoch teuer war. Dieser Befund deutet darauf an, dass die größte Hürde für zukünftige Verbesserungen nicht bessere Suchalgorithmen sind, sondern bessere Wege, wie lange Aufgaben tatsächlich dauern werden.
Die Arbeit bietet einen klaren Weg nach vorn. Sie zeigt, dass die Behandlung der Layout-Auswahl als formales Optimierungsproblem eine praktikable und leistungsstarke Strategie ist, die massive Geschwindigkeitssteigerungen liefern kann, wo einfache Regeln versagen. Sie stellt auch klar, dass die ultimative Leistungsgrenze nicht die Fähigkeit ist, die beste Lösung zu finden, sondern die Genauigkeit der Vorhersagen, die diese Suche leiten. Für Modelle mit regulären, vorhersehbaren Strukturen ist der neue Solver-basierte Ansatz bereits eine überlegene Wahl. Für chaotischere und komplexere Modelle muss sich der Fokus auf die Verfeinerung der Kostenmodelle verlagern, damit das mathematische Optimum mit der physischen Realität des Chips übereinstimmt. Indem sie das Problem der Suche nach der besten Lösung von dem Problem der Vorhersage der Kosten trennten, haben die Forscher den Compiler-Entwicklern ein neues Werkzeug gegeben, um ihren Fortschritt zu messen, sowie ein klares Ziel definiert, worauf sie ihre Bemühungen als Nächstes konzentrieren sollten.
Ertrinken Sie in Arbeiten in Ihrem Fachgebiet?
Erhalten Sie tägliche Digests der neuesten Arbeiten passend zu Ihren Forschungsbegriffen — mit technischen Zusammenfassungen, in Ihrer Sprache.