Diese Seite wurde von der Cloud Translation API übersetzt.

Compiler API

Hintergrund

Wir gehen davon aus, dass die Leser mit den Grundlagen der Sharding-Darstellung vertraut sind. Darin wird beschrieben, wie das Sharding eines Tensors in Shardy ausgedrückt werden kann. In diesem Dokument wird gezeigt, wie die Sharding-Darstellungen in einem Programm verwendet werden können, z.B. um einem bestimmten Tensor des Programms ein Sharding hinzuzufügen.

Bei der Sharding-Ausbreitung wird für jeden Tensor in einem Programm ein Sharding unter Berücksichtigung der Sharding-Einschränkungen für eine Teilmenge der Tensoren festgelegt. Die Compiler-API von Shardy bietet mehrere Möglichkeiten, die Sharding-Ausbreitung zu beeinflussen oder zu steuern. Außerdem können Nutzer manuell ge shardete Berechnungen in ihre Programme einfügen.

Ziel

In diesem Dokument wird das Design solcher API-Komponenten in Shardy beschrieben und ihr Verhalten und ihre Invarianten erläutert. Diese API wird zwar zur Steuerung der Sharding-Weitergabe verwendet, in diesem Dokument wird jedoch NICHT auf das Verhalten der Weitergabe oder ihre Funktionsweise eingegangen.

Übersicht

Eingabe-/Ausgabe-Sharding: Sie können einer Eingabe oder Ausgabe der Hauptfunktion ein Sharding zuweisen, um anzugeben, wie der Eingabe-/Ausgabetensor geSharded werden soll, wenn er an die Funktion übergeben oder von ihr zurückgegeben wird.
Sharding-Einschränkung: Hiermit wird einem Zwischentensor (z.B. dem Ergebnis einer Matrixmultiplikation) ein Sharding hinzugefügt, um anzugeben, wie dieser Tensor oder ein Teil seiner Verwendungen geSharded werden soll.
Sharding-Gruppe: Hiermit können Sie mehrere Tensoren anhand einer ID gruppieren, um anzugeben, dass sie auf dieselbe Weise gesplittet werden sollen.
Manuelle Berechnung: Umschließt eine Teilberechnung, die manuell mithilfe einer Teilmenge von Mesh-Achsen partitioniert wird.Die Shardings entlang dieser manuellen Achsen werden für alle Eingaben und Ausgaben angegeben.Innerhalb der Teilberechnung sind die Tensortypen relativ zu diesen Shardings lokal.

Detaillierte Konfiguration

Eingabe-/Ausgabe-Sharding

Ermöglicht Nutzern, ein Sharding für die Eingaben und Ausgaben der Hauptfunktion anzugeben.

In MLIR können Attribute an Funktionsargumente und ‑ergebnisse angehängt werden. So können Nutzer der Funktion sSharding-Attribute zuweisen.

Beispiel:

@mesh_xy = <["x"=2, "y"=2]>

// The 1st input has a sharding specified, but the 2nd input doesn't.
// The output has a sharding specified.
func @main(%arg0: tensor<8x8xf32>
            {sdy.sharding = #sdy.sharding<@mesh_xy, [{"x"}, {}]>},
            %arg1: tensor<8x16xf32>)
    -> (tensor<8x16xf32> {sdy.sharding = #sdy.sharding<@mesh_xy, [{}, {"y"}]>}) {
  ...
}

Sharding-Einschränkung

Damit können Nutzer einem Zwischentensor in ihrem Programm ein Sharding zuweisen, das dem Partitionierer mitteilt, wie dieser Tensor oder ein Teil seiner Verwendungen geSharded werden soll.

Dies ist eine MLIR-Operation, die den Tensor als Eingabe nimmt und mit einem Sharding-Attribut verknüpft ist. Der Vorgang kann entweder:

Sie werden nicht verwendet (dangling). Das bedeutet, dass der Tensor selbst gemäß dem angehängten Sharding gesplittet werden sollte.
Es gibt Anwendungsfälle. Das bedeutet, dass die Sharding-Beschränkungsoperationen mit der angehängten Sharding-Funktion gesplittet werden sollten, während andere Verwendungen des Eingabetensors eine andere Sharding-Funktion haben können. Wenn der Eingabetensor keine anderen Verwendungen hat, entspricht das Verhalten dem Fall ohne Verwendungen. Bei der Propagation wird das Sharding des Tensors selbst bestimmt und bei Bedarf neu gesplittet.

Es kann offene Dimensionsshardings haben, was bedeutet, dass der Operand entlang der verfügbaren Achsen weiter gesplittet werden kann.

@mesh_xy = <["x"=2, "y"=2]>

%0 = ... : tensor<8x8xf32>
%1 = sdy.sharding_constraint %0 <@mesh_xy, [{"x"}, {?}]> : tensor<8x8xf32>

Sharding-Gruppe

Wenn zwischen zwei oder mehr Tensoren keine Datenabhängigkeiten oder keine starken Datenabhängigkeiten bestehen, die Nutzer aber wissen, dass diese Tensoren auf dieselbe oder eine ähnliche Weise partitioniert werden sollten, bietet die Shardy API eine Möglichkeit, diese Beziehung anzugeben. So können Nutzer explizit angeben, dass Tensoren gleich partitioniert werden sollen.

Dazu führen wir Shard-Gruppen ein. Jede Gruppe enthält eine beliebige Anzahl von Anweisungen, die derselben Shard-Gruppen-ID zugeordnet sind. Bei Sharding-Gruppen müssen die Shardings innerhalb derselben Gruppe identisch sein.

In einem hypothetischen Nutzerprogramm wie dem unten gezeigten möchten wir beispielsweise die Ausgabe des Programms genau wie die Eingabe des Programms schardern, ohne dass es Datenabhängigkeiten zwischen den beiden gibt.

Wenn wir dieses Programm ausführen, kann die Sharding-Übertragung die Sharding-Informationen für die Tensoren %1 und %2 nicht ableiten und sie werden repliziert. Wenn wir jedoch ein shard_group-Attribut anhängen, das angibt, dass sich die Eingabe %0 und die Ausgabe %2 im selben shard_group befinden, können wir das Sharding @mesh_xy, [{"x"},{"y"}]> von der Eingabe %0 an die Ausgabe %2 und dann an den Rest des Graphen weitergeben, der hier als konstante %1 gesendet wird. Mit dem Vorgang sdy.sharding_group können wir einer Gruppe einen Wert zuweisen.

@mesh_xy = <["x"=2, "y"=2]>

module @"jit_zeros_like" {
  func.func @main(%arg0: tensor<8x2xi64> {sdy.sharding = #sdy.sharding<@mesh_xy, [{"x"},{"y"}]>} }) -> (tensor<8x2xi64>) {
    %0 = sdy.sharding_group %arg0, id=0 : tensor<8x2xi64>
    %1 = stablehlo.constant dense<0> : tensor<8x2xi64>
    %2 = sdy.sharding_group %1, id=0 : tensor<8x2xi64>
    return %2 : tensor<8x2xi64>
  }
}

In diesem einfachen Beispiel oben hätten wir alternativ die gleiche Sharding-Methode für die Ausgabe wie für die Eingabe angeben können, was denselben Effekt hätte, da wir bereits im Voraus wissen, welchen Shard wir der Eingabe zuweisen möchten. In realistischeren Fällen verwenden wir „shard“, um die Sharding-Methode mehrerer Tensoren synchron zu halten, ohne unbedingt die Sharding-Methode für jeden von ihnen zu kennen. Shardy kümmert sich um den Rest und findet die beste Sharding-Methode, die ihnen zugewiesen werden kann.

Manuelle Berechnung

Nutzer möchten möglicherweise genau festlegen, wie Teile ihrer Berechnungen partitioniert werden und welche Kollektive verwendet werden. Einige Nutzer möchten beispielsweise die kollektive Matrixmultiplikation manuell (über die Frontend-API) anwenden, anstatt sie an den Compiler zu übergeben. Wir stellen eine API für manuelle Berechnungen bereit, mit der dies möglich ist.

Dies ist der MLIR-Vorgang mit einer einzelnen Region für die manuelle Teilberechnung. Nutzer geben für diese Teilberechnung Eingabe-/Ausgabe-Shardings mit einer Teilmenge (einschließlich aller) der Mesh-Achsen an. Die Teilberechnung ist lokal/manuell im Hinblick auf die angegebenen Mesh-Achsen (sogenannte manuelle Achsen) und global/nicht partitioniert im Hinblick auf nicht angegebene Achsen (sogenannte kostenlose Achsen). Die Teilberechnung kann während der Übertragung auf die kostenlosen Achsen aufgeteilt werden, genau wie die Berechnung außerhalb dieses Vorgangs.

Beispiel:

@mesh_name = <["data"=2, "model"=2]>

%0 = ... : tensor<16x32xf32>
%1 = sdy.manual_computation(%0)
    in_shardings=[<@mesh_name, [{"data"}, {"model",?}]>]
    out_shardings=[<@mesh_name, [{"data"}, {?}]>]
    manual_axes={"data"}
    (%arg1: tensor<8x32xf32>) {
  // body
  return %42 : tensor<8x32xf32>
} : (tensor<16x32xf32>) -> tensor<16x32xf32>

Invarianten

Alle in_shardings, out_shardings und manual_axes müssen sich auf dasselbe Mesh beziehen. manual_axes ist nach dem Mesh sortiert.
Die manual_axes muss in allen s Shardings explizit verwendet werden. Das bedeutet, dass für jedes Sharding alle manuellen Achsen entweder eine Dimension sharden oder explizit repliziert werden müssen.
Wenn eine kostenlose Achse (eine beliebige Mesh-Achse, die nicht in manual_axes enthalten ist) in einer der Sharding-Ebenen vorhanden ist, muss sie einer manuellen Achse in derselben Dimensions-Sharding-Ebene untergeordnet sein. Im obigen Beispiel wäre eine Dimensions-Sharding-Ebene {"model", "data"} ungültig.
Die Region/das Objekt der Berechnung ist die lokale Berechnung (z.B. einschließlich vom Nutzer angegebener Gruppen). Sie muss lokal in Bezug auf das In-/Out-Sharding entlang manueller Achsen sein (siehe Hinweis oben).

Manuelle Berechnungen verschachteln

Sie können mehrere manuelle Berechnungen verschachteln, solange für jede Berechnung eine eigene Gruppe von manuellen Achsen verwendet wird.