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, которая принимает тензор в качестве входных данных и имеет атрибут шардинга, прикрепленный к нему. Операция может:

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

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

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

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

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

Пользователи могут захотеть явного контроля над тем, как части их вычислений разбиваются на разделы и какие коллективы используются. Например, некоторые пользователи хотят применять коллективный 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. Регион/тело вычисления — это локальное вычисление (например, включая указанные пользователем коллективы). Оно должно быть локальным относительно сегментирования входа/выхода вдоль ручных осей (см. примечание выше).

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

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