Это продолжение цикла статей о масштабировании тренировки и инференса LLM.

Предыдущая глава Следующая глава

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

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

Возьмем для примера матрицу A размера [I, J] и распределим ее на 4 ускорителя:

Заметим, что логическая или глобальная размерность матрицы как была так и осталась [|I|, |J|], но вот локальная размерность каждой подматрицы на каждом из устройств в данном случае сократилась в два раза, как по ширине, так и по высоте. Теперь давайте выведем общую нотацию для шардинга массивов данных.

Общая нотация для шардинга

Допустим у нас есть массив данных, он может быть одномерным (вектор), двумерным (матрица), трех и более мерным (тензор), в общем виде будем называть его тензором. Допустим у него есть оси, по одной на каждое его измерение, назовем их I, J, K, и т.д. Наша задача распределить этот тензор по массиву ускорителей, который также может быть многомерным, назовем его оси X, Y, Z, и т.д. Мы можем описать распределение тензора по массиву ускорителей, путем описания того, как каждая ось тензора соотносится с осями массива ускорителей.

Рассмотрим пример выше:

есть двумерный массив ускорителей, представляющий собой сетку 2х2, состоящий из 4 устройств с осями X и Y и размерами |X| и |Y| (в данном случае |X| = |Y| = 2), есть матрица с осями I, J и размерами |I| и |J|;
шардинг этой матрицы на данном массиве устройств может быть записан в формате, представленном ниже. Данная запись говорит нам о том, что первая ось I распределена вдоль оси X, а вторая ось J - вдоль оси Y:

То есть в данном случае на каждом устройстве хранится 1 / (|X| * |Y|) = 1/4 доля матрицы, и размеры этой доли [|I| // 2, |J| // 2]. Для простоты написания будем в дальнейшем размеры матрицы обозначать как [I, J], а размеры массива устройств - [X, Y].

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

$A[I_{XY}, J]$

В таком случае для массива ускорителей 2х2 на каждом из устройств будет храниться матрица размеров [I // 4, J]

Для простоты понимания давайте визуализируем распределение данных по устройствам для нашего случая. Вот собственно наша матрица и массив ускорителей с соответствующими осями:

Рассмотрим базовый сценарий, когда мы тупо копируем данные по всем устройствам, то есть на каждом из устройств находится полная реплика нашей матрицы:

Теперь шардируем ось I вдоль оси X (записывается как ):

Поскольку разбиение происходит только вдоль первой оси, каждая из подматриц на ускорителе имеет размеры 2х4 (т.е. [I//2, J]). Заметим, что если смотреть вдоль оси Y (вертикальной), мы будем видеть одни и те же массивы. Изменения происходят только если смотреть вдоль оси X (горизонтальной): матрица как бы смещается вдоль вертикальной оси I, если мы смещаемся вдоль оси X.

Ну и наконец рассмотрим случай шардирования обеих осей матрицы вдоль обеих осей массива ускорителей

А вот и все возможные комбинации:

Теперь рассмотрим поподробнее случаи вида $A[I_{XY}, J]$ и $A[I_{YX}, J]$ . В обоих случаях только первая ось матрицы "размазана" по обеим осям массива ускорителей, а вторая оставлена нетронутой, т.е. каждая подматрица на ускорителе имеет размер 1х4. А вот порядок размещения подматриц зависит от того, какая буква в нотации стоит первой, X или Y. Если $I_{XY},$ то следующий по порядку элемент встретится нам, если мы будем двигаться вдоль оси X, если $I_{YX}$ - соответственно вдоль оси Y. Вот как это выглядит на картинке:

А теперь о том, как шардировать не надо. Мы не можем шардировать обе оси матрицы вдоль одной оси устройства, например $A[I_{X}, J_{X}]$ . В этом нет никакого смысла, т.к. получится, что при движении вдоль оси Х мы одновременно смещаемся вдоль обеих осей матрицы и часть элементов будет просто потеряна. В случае с приведенным выше примером получится, что на первом устройстве будут элементы (00, 01), на втором - (12, 13), на третьем - (20, 21), на четвертом - (32, 33). То есть половина элементов тупо потеряется, что не есть хорошо. Поэтому как только мы использовали одну из осей массива ускорителей для шардинга, в дальнейшем мы ее уже не используем.

Умножаем шардированные матрицы

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

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

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

И тут уже все зависит от того, как именно шардированы матрицы на каждом из ускорителей. Рассмотрим например такое умножение:

$A[I_X, J] \cdot B[J, K_Y] \rightarrow C[I_X, K_Y]$

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

Но если мы хотим получить на выходе целую матрицу вместо "размазанной" по шардам, т.е.:

$A[I_X, J] \cdot B[J, K_Y] \rightarrow C[I, K]$

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

Вариант 1

Ни одна из смежных осей входных матриц А и В не шардирована.

$A[I_X, J] \cdot B[J, K_Y] \rightarrow C[I_X, K_Y]$

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

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

Другие варианты подобного умножения:

$\mathbf{A}[I, J] \cdot \mathbf{B}[J, K] \rightarrow \mathbf{C}[I, K] \\$ $\mathbf{A}[I_X, J] \cdot \mathbf{B}[J, K] \rightarrow \mathbf{C}[I_X, K] \\$ $\mathbf{A}[I, J] \cdot \mathbf{B}[J, K_Y] \rightarrow \mathbf{C}[I, K_Y] \\$

Вариант 2

Одна из смежных осей матриц шардирована. Пусть это будет вторая ось матрицы А:

$\mathbf{A}[I, J_X] \cdot \mathbf{B}[J, K] \rightarrow \mathbf{C}[I, K]$

В таком случае умножением на месте уже не отделаешься, придется по кусочкам собирать матрицу А вдоль оси J и только потом производить умножение матриц:

$\mathbf{AllGather}_X[I, J_X] \rightarrow \mathbf{A}[I, J] \\$ $\mathbf{A}[I, J] \cdot \mathbf{B}[J, K] \rightarrow \mathbf{C}[I, K]$

И тут мы впервые сталкиваемся с операцией AllGather. Что она делает? В нашем случае собирает кусочки матрицы А с каждого устройства и рассылает их по всем остальным устройствам так, что чтобы на каждом устройстве была полная копия матрицы А.

Собственно вот как это выглядит:

Сколько по времени занимает данная операция? Не вдаваясь в подробности просто приведу итоговую формулу:

$T_{total} = \dfrac{N}{W_{ici}}$

Где это количество байт, которое нужно передать (в нашем случае это размер матрицы А в байтах), а $W_{ici}$ - пропускная способность шины данных между ускорителями. То есть длительность этой операции по сути не зависит от количества ускорителей в нашем массиве, важно только, сколько байт нужно передать и с какой скоростью один ускоритель может передавать данные другому.

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

Вариант 3

Обе смежные оси обеих матриц шардированы, причем шардированы вокруг одной и той же оси массива ускорителей:

$\mathbf{A}[I, J_X] \cdot \mathbf{B}[J_X, K] \rightarrow \mathbf{C}[I, K]$

В таком случае нельзя провести прямое умножение строки матрицы А на столбец матрицы В, потому что только часть как строки, так и столбца находится на одном ускорителе. Но поскольку шардирование проведено вокруг одной и той же оси массива ускорителей, строки матрицы А и столбцы матрицы В на каждом устройстве имеют пусть и не полный, но хотя бы один и тот же набор индексов. То есть у нас есть I кусков строк матрицы А и K кусков столбцов матрицы В, поэтому после их перемножения получаем матрицу того же размера, что и матрица С. Записывается это все так:

$\mathbf{A}[I, J_X] \cdot_{\text{LOCAL}} \mathbf{B}[J_X, K] \rightarrow \mathbf{C}[I, K]\{U_X\}$

Нотация $\{U_X\}$ означает unreduced или в переводе на русский "частичная сумма" вдоль оси X. Теперь получившуюся частичную матрицу нужно поэлементно сложить с другими такими же матрицами на других устройствах и получим нужную нам матрицу С. Делается это при помощи операции AllReduce:

$\mathbf{AllReduce}_X C[I, K]\{U_X\} \rightarrow C[I, K]$

Давайте разберемся, что же такое операция AllReduce. По сути это последовательность двух операций: ReduceScatter и AllGather.

А что такое ReduceScatter? А это собственно когда мы собираем частичные суммы на каком-то из устройств в одну полную сумму. Выглядит это примерно так:

А далее мы копируем эти полные суммы по всем оставшимся устройствам при помощи уже знакомой нам операции AllGather.

Опять таки не вдаваясь в подробности просто скажем, что ReduceScatter занимает ровно столько же времени, сколько и AllGather, поэтому AllReduce занимает в два раза больше времени, чем AllGather.

В итоге вся последовательность операций выглядит так:

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

Вот схема для наглядности

Вариант 4

Обе несмежные оси обеих матриц шардированы, причем шардированы вокруг одной и той же оси массива ускорителей:

$A[I_X, J] \cdot B[J, K_X] \rightarrow C[I_X, K_X]$

Это некорректный вариант, по той же причине, по которой нельзя шардировать две оси матрицы вдоль одной оси массива устройств $A[I_{X}, J_{X}]$ - просто в итоге потеряем часть данных, поэтому тут для начала придется приводить к полной форме одну из двух матриц, а затем решать все через вариант 2.

AllToAll

А теперь давайте рассмотрим последнюю фундаментальную операцию, используемую для обмена данными между ускорителями - AllToAll. Она не используется при шардированном умножении матриц, однако широко применяется на практике для выполнения других операций над массивами данных.

По сути это это что-то вроде транспонирования: мы берем матрицу, которая шардирована вдоль одной из своих осей и меняем ее шардирование на другую ось:

$\mathbf{AllToAll}_{X,J} A[I_X, J] \rightarrow A[I, J_X]$

Схематично выглядит это примерно вот так:

AllToAll широко применяется в LLM с архитектурой Mixture-of-Experts (MoE), в которых вместо того, чтобы прогонять каждый токен через всю сеть целиком, на каждом слое активируется лишь небольшая часть параметров — тех самых "экспертов", которые наиболее релевантны для данного входа. Соответственно если LLM большая, и каждый эксперт живет на своем устройстве, то на каждом MoE-слое возникает необходимость доставить токены к нужным экспертам на произвольных устройствах и вернуть результаты обратно. Эта задача и решается при помощи AllToAll: сначала токены рассылаются экспертам (dispatch), а после обработки — возвращаются на исходные устройства (combine).

Опять таки, не вдаваясь в подробности, скажем, что сложность AllToAll равна 1/4 от сложности операции AllGather.

А теперь еще раз посмотрим на все основные операции обмена данными между ускорителями в сравнении:

Бонус. Как совместить передачу данных и вычисления при матричном умножении

Как мы уже говорили в первой части, обычно передачу данных между устройствами можно совместить с какой-нибудь полезной работой, если, конечно, сеть достаточно быстрая. Операции, описанные в этом разделе, в принципе можно запускать параллельно с самим матричным умножением, но на практике это не так просто, как звучит. Алгоритм, который для этого используется, называется collective matmul. Вот упрощённая анимация того, как он работает:

А на этом все, в следующей главе перейдем к чему-то более прикладному и поговорим о трансформерах и вычислительных операциях, необходимых для их работы. До встречи!

Масштабирование LLM: от одного чипа до ЦОДа. Глава 2. Шардинг