Interfejs Compiler API

Tło

Zakładamy, że czytelnicy znają co najmniej podstawy reprezentacji dzielenia na części, która opisuje, jak można zastosować dzielenie na części w Shardy. Ten dokument pokazuje, jak można używać w programie reprezentowania podziału, np. do dołączenia podziału do konkretnego tensora w programie.

Propagowanie podziału na fragmenty to proces podejmowania decyzji o podziale na fragmenty każdego tensora w programie z uwzględnieniem ograniczeń podziału na fragmenty dla podzbioru tensorów. Interfejs kompilatora Shardy API udostępnia kilka sposobów na wpływanie na propagowanie dzielenia na fragmenty lub kontrolowanie tego procesu. Dodatkowo umożliwia użytkownikom wstawianie do swoich programów obliczeń podzielonych na fragmenty ręcznie.

Cel

Ten dokument opisuje projekt takich komponentów interfejsu API w Shardy oraz ich zachowanie i niezmienniki. Pamiętaj, że chociaż ten interfejs API służy do kontrolowania propagacji podziału, w tym dokumencie NIE omawiamy zachowania propagacji ani jej projektowania.

Omówienie

  • Dzielenie na częście wejścia i wyjścia – dołącz do wejścia lub wyjścia funkcji głównej funkcję dzielenia na częście, aby wskazać, jak tensor wejścia/wyjścia powinien być podzielony na częście podczas przekazywania do funkcji lub zwracania z niej.

  • Ograniczenie podziału na części – dołącz podział na części do pośredniego tensora (np. wyniku funkcji matmul), aby wskazać, jak ten tensor lub podzbiór jego zastosowań powinien być podzielony na części.

  • Grupa dzielenia na fragmenty – możesz pogrupować wiele tensorów według identyfikatora, aby wskazać, że należy je podzielić na fragmenty w ten sam sposób.

  • Ręczne obliczenia – otaczają podobliczenia, które są ręcznie dzielone za pomocą podzbioru osi siatki.W przypadku wszystkich wejść i wyjść określane są podziały na kawałki wzdłuż tych osi ręcznych.W ramach podobliczeń typy tensorów są lokalne w stosunku do tych podziałów.

Szczegółowy projekt

Dzielenie wejść i wyjść

Umożliwia użytkownikom określenie podziału na części dla danych wejściowych i wyjściowych funkcji głównej.

W MLIR atrybuty można dołączać do argumentów i wyników funkcji, dzięki czemu użytkownicy mogą w ten sposób dołączać atrybuty podziału do funkcji.

Na przykład:

@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"}]>}) {
  ...
}

Ograniczenie dotyczące podziału na fragmenty

Umożliwia użytkownikom dołączenie podziału do pośredniego tensora w programie, co informuje partycjonizera, jak należy podzielić tensor lub podzbiór jego zastosowań.

To operacja MLIR, która przyjmuje jako dane wejściowe tensor i ma atrybuty podziału. Operacja może:

  • nie mają żadnych zastosowań (są zawieszone) – oznacza to, że dołączone shardingi to sposób, w jaki tensor powinien być podzielony na części;
  • Have uses – oznacza, że dołączone podziały to sposób, w jaki należy dzielić elementy operacji ograniczeń podziału, podczas gdy inne elementy tensora wejściowego mogą mieć inny podział (jeśli tensor wejściowy nie ma innych zastosowań, zachowanie jest takie samo jak w przypadku braku zastosowań). Propagacja określi podział tensora i w razie potrzeby podzieli go ponownie.

Może zawierać otwarte podziały wymiaru, co oznacza, że operand może być dalej dzielony wzdłuż dostępnych osi.

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

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

Grupa fragmentacji

W przypadku braku zależności danych lub braku silnych zależności danych między co najmniej 2 tensorami, gdy użytkownicy wiedzą, że te tensory powinny być dzielone w taki sam lub podobny sposób, interfejs Shardy API umożliwia określenie tej relacji. Daje to użytkownikom swobodę w określaniu, że tensory powinny być podzielone na siebie.

Aby to osiągnąć, wprowadzamy pojęcie grup fragmentów, z których każda zawiera dowolną liczbę instrukcji powiązanych z tym samym identyfikatorem grupy fragmentów. Grupy fragmentacji powodują, że fragmentacje w ramach tej samej grupy są takie same.

Na przykład w przypadku hipotetycznego programu użytkownika, takiego jak ten poniżej, chcemy udostępnić dane wyjściowe programu w taki sam sposób jak dane wejściowe programu, przy czym nie ma między nimi żadnych zależności danych.

Jeśli uruchomimy ten program, propagacja dzielenia na części nie będzie mogła wywnioskować na podstawie podziału na części tensorów %1 i %2, więc zostaną one powielone. Jednak dołączając atrybut shard_group, który wskazuje, że dane wejściowe %0 i dane wyjściowe %2 znajdują się w tym samym shard_group, umożliwiamy podział danych @mesh_xy, [{"x"},{"y"}]> od danych wejściowych %0 do danych wyjściowych %2, a następnie do reszty grafu, który jest transmitowany jako stała %1 tutaj. Wartość możemy przypisać do grupy za pomocą operacji sdy.sharding_group.

@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>
  }
}

W tym prostym przykładzie moglibyśmy też jawnie określić ten sam podział na wyjściu i wejściu, co dałoby ten sam efekt, ponieważ wiemy już, jaki podział chcemy przypisać do wejścia. W bardziej realistycznych przypadkach używamy jednak podziału, aby zachować spójność podziału wielu tensorów bez konieczności znajomości podziału żadnego z nich. Shardy zajmie się resztą i znajdzie najlepszy podział do przypisania do nich.

Obliczenia ręczne

Użytkownicy mogą chcieć mieć pełną kontrolę nad tym, jak dzielone są części ich obliczeń i jakie zbiory zbiorcze są używane. Niektórzy użytkownicy chcą na przykład stosować zbiorczą funkcję matmul ręcznie (z interfejsu API) zamiast przekazywać ją do kompilatora. Udostępniamy interfejs Manual Computation API, który umożliwia im to.

Jest to operacja MLIR z jednym regionem do wykonania podzadań ręcznych. Użytkownicy określają dzielenie na części danych wejściowych i wyjściowych w ramach tego podzasobu obliczeń, używając podzbioru (w tym wszystkich) osi siatki. Podobliczenie byłoby lokalne/ręczne w stosunku do określonych osi siatki (czyli osi ręcznych) oraz globalne/niedzielone w stosunku do nieokreślonych osi (czyli osi swobodnych). Podobliczenie może być dalej dzielone na podzbiory wzdłuż osi swobodnych podczas propagacji w taki sam sposób, w jaki można dzielić obliczenia poza tą operacją.

Na przykład:

@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>

Niezmienniki

  1. Wszystkie parametry in_shardings, out_shardings i manual_axes muszą odnosić się do tego samego siatki. manual_axes jest posortowany względem siatki.

  2. Wartość manual_axes musi być używana w wszystkich podziale na podzbiory wejść i wyjść, czyli w przypadku każdego podziału wszystkie osi ręcznych muszą być albo podzielone na wymiary, albo wyraźnie powielone.

  3. Jeśli w jednym z podziałów we/out występuje wolna oś (dowolna oś siatki, która nie jest w grupie manual_axes), musi ona być podrzędna dowolnej osi ręcznej w tym samym podziale wymiaru (w powyższym przykładzie podział wymiaru {"model", "data"} byłby nieprawidłowy).

  4. Region/element obliczeń to obliczenia lokalne (np. obejmujące zbiory określone przez użytkownika). Musi być lokalny w stosunku do podziału wejść/wyjść wzdłuż osi ręcznych (patrz uwaga powyżej).

Zagnieżdżanie obliczeń ręcznych

Możesz zagnieżdżać dowolną liczbę obliczeń ręcznych, o ile każde z nich działa na własnym zestawie osi ręcznych.