Langfristig soll Shardy eine vollständig eigenständige Komponente werden, die mit jedem MLIR-Dialekt funktioniert. Derzeit hängt Shardy direkt von StableHLO ab, aber wir arbeiten daran, dies durch verschiedene Abstraktionsebenen und Schnittstellen zu ändern, um Shardy flexibler zu machen.
Regeln für das Sharding
Eine Sharding-Regel codiert, wie wir eine Änderung an einem Sharding-Cluster weitergeben. Da Shardy jetzt von StableHLO abhängt, werden Sharding-Regeln für jeden StableHLO-Vorgang definiert. Darüber hinaus bietet Shardy die Funktion ShardingRuleOpInterface, mit der Eigentümer von Dialects in ihren Vorgängen Sharding-Regeln für ihre eigenen Vorgänge definieren können. Solange ein Vorgang diese Schnittstelle implementiert, kann Shardy ihn weitergeben.
def ShardingRuleOpInterface : OpInterface<"ShardingRuleOpInterface"> {
let methods = [
InterfaceMethod<
/*desc=*/[{
Returns the sharding rule of the op.
}],
/*retType=*/"mlir::sdy::OpShardingRuleAttr",
/*methodName=*/"getShardingRule"
>,
];
}
Datenflussvorgänge
Für einige Vorgänge, z.B. regionbasierte Vorgänge, ist ein anderer Ansatz erforderlich, bei dem Sharding-Regeln, die nur die Übereinstimmung zwischen Dimensionen für alle Operanden und Ergebnisse beschreiben, nicht ausreichen. In diesen Fällen definiert Shardy eine ShardableDataFlowOpInterface, damit die Eigentümer des Dialects die Weiterleitung des Shardings über ihre Vorgänge beschreiben können. Diese Schnittstelle bietet Methoden, um die Quellen und Ziele jeder Datenflusskante über ihren Eigentümer abzurufen und die Shardings von Kanteneigentümern abzurufen und festzulegen.
def ShardableDataFlowOpInterface :
OpInterface<"ShardableDataFlowOpInterface"> {
(get|set)BlockArgumentEdgeOwnerShardings;
(get|set)OpResultEdgeOwnerShardings;
getBlockArgumentEdgeOwners;
getOpResultEdgeOwners;
getEdgeSources;
// ...
}
Weitere Informationen dazu, wie wir mit Datenfluss-Vorgängen umgehen, finden Sie unter Datenfluss-Vorgänge.
Noch nicht implementierte Oberflächen
In Zukunft werden weitere Schnittstellen und Merkmale hinzugefügt, um Shardy flexibler und dialektunabhängiger zu machen. Sie finden sie unten.
Konstante Aufteilung
Die meisten Tensorprogramme in MLIR haben eine Instanz einer Konstante, die von allen Operatoren wiederverwendet wird, die diesen Wert benötigen. Das ist sinnvoll, wenn die benötigte Konstante immer gleich ist. Für eine optimale Sharding-Verwaltung eines Programms möchten wir jedoch, dass jede Verwendung einer Konstante eine eigene Sharding-Verwaltung hat und nicht von der Verwendung dieser Konstante durch andere Vorgänge beeinflusst wird.
Wenn in der Abbildung unten beispielsweise add geSharded ist, sollte dies keinen Einfluss darauf haben, wie divide und subtract (in verschiedenen Teilen der Berechnung) geSharded werden.
Wir nennen das eine falsche Abhängigkeit: Da Konstanten kostengünstig sind, gibt es keine echte Abhängigkeit zwischen Operationen, die dieselbe Konstante verwenden. So können Nutzer das Sharding ihrer konstanten (und konstantenähnlichen) Vorgänge selbst festlegen. Jede Verwendung dieser Konstante kann dann ein anderes Sharding haben, das sich unabhängig auf eine eigene Kopie der konstanten Teilberechnung auswirken kann.
Dazu müssen Shardy-Nutzer Folgendes definieren:your_dialect.constantsdy.constantsdy::ConstantLikemlir::Elementwiseaddmultiplysdy::ConstantFoldable
Diese Operationen können technisch gesehen zur Kompilierungszeit berechnet werden, wenn alle Operanden/Ergebnisse Konstanten sind.
Op-Prioritäten
Bei GSPMD werden zuerst elementweise Vorgänge weitergegeben, gefolgt von Vorgängen wie matmul.
In Shardy möchten wir es Nutzern ermöglichen, ihre eigenen OP-Prioritäten festzulegen, da wir ihre Dialekte a priori nicht kennen. Daher bitten wir sie, eine Liste mit den Vorgängen in der Reihenfolge zu übergeben, in der sie von Shardy übernommen werden sollen.
Die folgende Abbildung zeigt, wie die Prioritäten in GSPMD verwendet werden, um Vorgänge in der richtigen Reihenfolge zu übertragen.
Im GSPMD-Papier wird erläutert, warum operative Prioritäten wichtig sind.
Dialektunabhängig
Solange Sie die vorherigen Schnittstellen, Merkmale und Ausweise implementieren, kann Shardy für Ihren Dialekt verwendet werden. Wir arbeiten daran, Shardy flexibler und dialektunabhängig zu machen. Wir halten dich auf dem Laufenden.