API компилятора

Фон

Мы предполагаем, что читатели знакомы хотя бы с основами представления шардинга , которые описывают, как шардинг тензора может быть выражен в Shardy. В этом документе показано, как представления сегментирования могут использоваться в программе, например, для присоединения сегментирования к определенному тензору программы.

Распространение сегментирования — это процесс принятия решения о сегментировании для каждого тензора в программе с учетом ограничений сегментирования для подмножества тензоров. API-интерфейс компилятора Shardy предоставляет несколько способов влияния/контроля распространения шардинга. Кроме того, он позволяет пользователям вставлять в свои программы сегментированные вручную вычисления.

Цель

Этот документ описывает дизайн таких компонентов API в Shardy и объясняет их поведение и инварианты. Обратите внимание: хотя этот API используется для управления распространением сегментирования, в этом документе НЕ обсуждается ничего ни о поведении распространения, ни о том, как оно устроено.

Обзор

  • Шардинги ввода/вывода — прикрепите сегментацию к входу или выходу основной функции, чтобы указать, что именно так должен сегментироваться тензор ввода/вывода при передаче в функцию или возврате из функции.

  • Ограничение сегментирования — прикрепите сегментирование к промежуточному тензору (например, результату matmul), чтобы указать, что именно так должен быть сегментирован этот тензор или подмножество его применений.

  • Группа сегментирования — группируйте несколько тензоров по идентификатору, чтобы указать, что их следует сегментировать одинаковым образом.

  • Ручное вычисление — включает в себя подвычисление, которое разделено вручную с использованием подмножества осей сетки, где сегменты вдоль этих ручных осей указаны для всех входов и выходов, а внутри подвычисления типы тензоров являются локальными относительно этих сегментов.

Детальный проект

Шардинги ввода/вывода

Позволяет пользователям указывать сегментирование входных и выходных данных основной функции.

В MLIR атрибуты могут быть прикреплены к аргументам и результатам функции, и поэтому пользователи могут таким образом прикреплять атрибуты сегментирования к функции.

Например:

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

Ограничение сегментирования

Позволяет пользователям прикреплять сегментирование к промежуточному тензору в своей программе, что сообщает секционатору, что именно так должен быть сегментирован этот тензор или подмножество его применений.

Это операция MLIR, которая принимает тензор в качестве входных данных и к нему прикреплен атрибут сегментирования. Операция может:

  • Не использовать (висячие) — это означает, что прикрепленный шардинг — это то, как должен быть сегментирован сам тензор.
  • Иметь использование - это означает, что прикрепленное сегментирование - это то, как должно быть сегментировано использование операции ограничения сегментирования, в то время как другие виды использования входного тензора могут иметь другое сегментирование (если входной тензор не имеет других применений, то поведение такое же, как и у входного тензора). нет варианта использования). Распространение определит шардирование самого тензора и при необходимости повторит его.

Он может иметь открытые сегменты измерений, что означает, что операнд может быть дополнительно сегментирован по доступным осям.

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

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

Шардинговая группа

В случаях, когда нет зависимостей данных или нет сильных зависимостей данных между двумя или более тензорами, хотя пользователи знают, что эти тензоры должны быть разделены одинаковым или похожим образом, API Shardy предлагает способ указать это отношение. Это дает пользователям свободу явно указывать, что тензоры должны быть разделены друг на друга.

Для достижения этой цели мы вводим понятие групп сегментов, где каждая группа содержит любое количество инструкций, связанных с одним и тем же идентификатором группы сегментов. Группы шардинга обеспечивают одинаковое шардинги внутри одной группы.

Например, в гипотетической пользовательской программе, такой как показанная ниже, мы хотим сегментировать выходные данные программы точно так же, как и входные данные программы, при этом между ними нет зависимостей данных.

Если мы запустим эту программу, распространение сегментирования не сможет сделать вывод о сегментировании тензоров %1 и %2 , и в конечном итоге они будут реплицированы. Однако, прикрепив атрибут shard_group , который говорит, что вход %0 и выход %2 находятся в одной и той же shard_group , мы разрешаем распространение сегментирования @mesh_xy, [{"x"},{"y"}]> из вход %0 на выход %2 и, в свою очередь, на остальную часть графика, которая здесь транслируется константой %1 . Мы можем присвоить значение группе с помощью операции 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>
  }
}

В этом простом примере выше, в качестве альтернативы, мы могли бы явно указать тот же сегмент на выходе, что и на входе, что позволило бы достичь того же эффекта, поскольку мы заранее знали, какой сегмент мы хотим назначить входу, но в В более реалистичных случаях мы используем сегментирование, чтобы синхронизировать сегментирование нескольких тензоров, не обязательно зная сегментирование для любого из них, в то время как Шарди позаботится обо всем остальном и найдет лучший сегмент для назначения им.

Ручной расчет

Пользователям может потребоваться явный контроль над тем, как распределяются части их вычислений и какие коллективы используются. Например, некоторые пользователи хотят применять коллективный matmul вручную (из интерфейсного API), а не откладывать выполнение компилятора. Мы предоставляем API ручных вычислений, который позволяет им это делать.

Это операция MLIR с одной областью для ручного подвычисления. Пользователи будут указывать сегментирование ввода/вывода для этого подвычисления, используя подмножество (включая, возможно, все) осей сетки. Подвычисления будут локальными/ручными относительно указанных осей сетки (также известных как ручные оси) и глобальными/неразделенными относительно неуказанных осей (также известных как свободные оси). Подвычисление может быть дополнительно сегментировано по свободным осям во время распространения таким же образом, как и вычисления вне этой операции.

Например:

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

Инварианты

  1. Все in_shardings , out_shardings и manual_axes должны относиться к одной и той же сетке. manual_axes сортируется по сетке.

  2. Параметр manual_axes должен явно использоваться во всех сегментах ввода/вывода, т. е. для каждого сегментирования все ручные оси должны либо сегментировать измерение, либо явно реплицироваться.

  3. Если свободная ось (любая ось сетки, не указанная в manual_axes ) существует в одном из входных/выходных сегментов, она должна быть второстепенной по отношению к любой ручной оси в том же сегменте измерения (в приведенном выше примере сегментирование измерения {"model", "data"} будет недействительным).

  4. Область/тело вычисления является локальным вычислением (например, включая определяемые пользователем коллективы). Оно должно быть локальным относительно сегментирования ввода/вывода по ручным осям (см. примечание выше).

Вложение ручных вычислений

Вы можете вкладывать несколько ручных вычислений друг в друга, если каждый из них работает со своим уникальным набором ручных осей.