Shardy

Shardy ist ein MLIR-basiertes Tensor-Partitionierungssystem für alle Dialekte. Es wurde in Zusammenarbeit der Teams GSPMD und PartIR entwickelt und vereint das Beste aus beiden Systemen sowie die Erfahrungen der Teams und Nutzer.

Vorteile

  • Mehr Kontrolle und Vorhersehbarkeit für Nutzer durch die Kombination der GSPMD-Propagation mit der inkrementellen Partitionierung von PartIR.
  • Neue Funktionen, die auf gemeinsamen Erfahrungen basieren, z.B. eine neue Unterstützung für Neugestaltungen, die bekanntermaßen zusätzliche Kommunikation erfordern, es sei denn, die Nutzer wissen, wie sie damit umgehen können.
  • Bessere Nutzerfreundlichkeit und Fehlerbehebung zur Steigerung der Geschwindigkeit für Endnutzer, z.B. durch eine achsebasierte Sharding-Darstellung.
  • Eine einfache Open-Source-Codebasis mit MLIR und einer größeren Gruppe aktiver Mitwirkender (intern, extern und in verschiedenen Zeitzonen), die Nutzer unterstützen.

Komponenten

  • Sharding-Darstellung: Eine achsebasierte Sharding-Darstellung, die an ein bestimmtes logisches Mesh (aus potenziell mehreren Meshes) gebunden ist und unter anderem die Einschränkung von Dimensions-Shardings und Achsen, die Aufteilung von Achsen für Vorgänge wie die Neuformatierung und Prioritäten für die inkrementelle Partitionierung unterstützt.
  • Compiler-APIs: Eine Reihe von Compilerkomponenten, die zusammen mit der Sharding-Darstellung verwendet werden können, um die Sharding-Ausbreitung zu beeinflussen.
    • Eingabe-/Ausgabe-Sharding: Sie können einer Eingabe oder Ausgabe der Hauptfunktion ein Sharding zuweisen, um anzugeben, wie der Eingabe-/Ausgabetensor gesplittet werden soll, wenn er an die Funktion übergeben oder von ihr zurückgegeben wird.
    • Sharding-Einschränkung: Sie können einem Zwischentensor (z.B. dem Ergebnis einer Matrixmultiplikation) ein Sharding zuweisen, um anzugeben, wie dieser Tensor oder ein Teil seiner Verwendungen geSharded werden soll.
    • Shard As/Like: Hiermit können Sie mehrere Tensoren nach einer ID gruppieren, um anzugeben, dass sie auf dieselbe Weise gesplittet werden sollen.
    • Manuelle Berechnung: Umschließt eine Teilberechnung, die manuell mit einer Teilmenge von Mesh-Achsen partitioniert wird.Dabei werden die Shardings entlang dieser manuellen Achsen für alle Eingaben und Ausgaben angegeben.Innerhalb der Teilberechnung sind die Tensortypen relativ zu diesen Shardings lokal.
  • Sharding-Propagation: Ein Propagation-Algorithmus, der Nutzerprioritäten und ‑einschränkungen mit Compilerkostenmodellen und Heuristiken kombiniert:
    • Benutzerdefinierte Prioritäten, z.B. Batchparallelität und dann ZeRO
    • Operatorbasierte Prioritäten, z.B. zuerst elementweise Operationen und dann Matrizenmultiplikationen usw.
    • Detailliertere Heuristiken, z. B. Bevorzugung von Batchdimensionen
  • SPMD-Partitioner: Eine Komponente, die Entscheidungen zur Sharding-Ausbreitung vereinfacht, indem das Programm in ein SPMD-Programm partitioniert und dabei die erforderlichen Datenbewegungen/-formatierungen und kollektiven Vorgänge hinzugefügt werden.
    • Für die Erstimplementierung wird kurzfristig der aktuelle GSPMD-SPMD-Partitioner verwendet.
    • Langfristig planen wir, einen neuen MLIR-basierten SPMD-Partitionierer zu erstellen.

Code-Repository

Das Shardy-Projekt befindet sich in der aktiven Entwicklungsphase und wir suchen das Feedback der Open-Source-Community. Der Shardy-Code ist unter https://github.com/openxla/shardy verfügbar.