Ta strona została przetłumaczona przez Cloud Translation API.

Interfejs Compiler API

Wprowadzenie

Zakładamy, że czytelnicy znają co najmniej podstawy reprezentacji dzielenia, która opisuje, jak można zastosować dzielenie tensora 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 API kompilatora Shardy 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 NOT omawiamy zachowania propagacji ani jej projektowania.

Omówienie

Dzielenie na segmenty danych wejściowych i wyjściowych – dołącz do danych wejściowych lub wyjściowych funkcji głównej funkcję dzielenia na segmenty, aby wskazać, jak tensor danych wejściowych/wyjściowych powinien być dzielony na segmenty podczas przekazywania do funkcji lub zwracania z niej.
Ograniczenie podziału na fragmenty – dołącz podział na fragmenty do pośredniego tensora (np. wyniku matmul), aby wskazać, jak ten tensor lub podzbiór jego zastosowań powinien zostać podzielony na fragmenty.
Grupa dzielenia na części – możesz pogrupować wiele tensorów według identyfikatora, aby wskazać, że powinny być dzielone na części 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łączanie do pośredniego tensora w programie elementu podziału, który informuje partycjoner, 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łączony podział powinien być zastosowany do podziału tensora;
Have uses – oznacza to, że dołączone podziały to sposób, w jaki należy dzielić użycia operatora ograniczeń podziału, podczas gdy inne użycia 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 wymiarów, 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, gdy między 2 lub więcej tensorami nie ma zależności danych ani silnych zależności danych, a 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 dzielenia na części tensorów %1 i %2, więc zostaną one powielone. Jednak dołączając atrybut shard_group, który mówi, ż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"}]>, który jest propagowany z 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ć takie samo podział na części w wyjściu jak i w wejściu, co dałoby ten sam efekt, ponieważ wiemy już, jaki podział na części chcemy przypisać do danych wejściowych. W bardziej realistycznych przypadkach używamy jednak podziału na części, aby zachować podział na części wielu tensorów w zsynchronizowany sposób, niekoniecznie znając podział na części dowolnego z nich. Shardy zajmie się resztą i znajdzie najlepszy podział na części do przypisania do nich.

Obliczenia ręczne

Użytkownicy mogą chcieć mieć kontrolę nad tym, jak części ich obliczeń są dzielone i jakie zbiory zbiorcze są używane. Niektórzy użytkownicy chcą na przykład ręcznie stosować zbiorczą funkcję matmul (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 podobliczenia, 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ą.