Масштабирование LLM: от одного чипа до ЦОДа. Глава 4. Тренировка трансформера

Предыдущая глава ^[1]

Наконец-то мы добрались непосредственно до того, как тренировать трансформер, и не просто тренировать, а делать это эффективно и масштабируемо.

Как мы уже знаем из прошлых глав, трансформер штука тяжелая и на один ускоритель обычно не влезает, поэтому цель масштабирования состоит в том, чтобы распихать тренировку модели по нескольким ускорителям, и желательно при этом, чтобы производительность и пропускная способность такой системы росли пропорционально количеству ускорителей в ней. А этого добиться довольно сложно, просто потому что чем больше ускорителей в системе, тем сложнее и накладнее передавать данные между частями системы и все это дело синхронизировать. Как мы видели в одной из предыдущих глав про шардинг матричных операций ^[2], распределенное матричное умножение требует разных дополнительных операций вроде AllGather или ReduceScatter, которые занимают шину данных и тратят процессорное время. В общем наша задача не просто масштабировать систему, а еще и понять, когда остановиться, потому что накладные расходы становятся совсем неподъемными.

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

Теперь кратко перечислим условные обозначения

Модель

D – размерность входных эмбеддингов

F – размерность скрытого MLP слоя (как я писал в предыдущей главе, большая часть параметров и вычислений трансформера приходится именно на такие большие MLP слои)

B – размер батча

T – длина последовательности

L – количество слоев модели

Оборудование

C – производительность ускорителя FLOPS/с

W – суммарная пропускная способность шины данных (по умолчанию будем считать шину двунаправленной)

X – количество ускорителей вдоль оси Х

Y – количество ускорителей вдоль оси Y

Z – количество ускорителей вдоль оси Z

Опять таки для простоты будем считать что трансформер это просто последовательность MLP блоков, так как именно они кушают больше всего вычислений. Так что слой трансформера выглядит в нашем представлении примерно вот так:

Ну а теперь рассмотрим все виды параллелизма (всего их 4).

1. Параллелизм данных (Data Parallelism)

Самый простой и очевидный вариант. Формула выглядит так:

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

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

Когда применять?

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

Так почему же данный метод масштабирования не применяют везде и всюду? А потому что, как я уже писал выше, трансформер штука тяжелая, а трансформер, который тренируется, тяжелеет минимум раз в 10 (точнее от 10 до 20 раз, как я уже писал в предыдущей статье ^[1] цикла). То есть если хочешь запихнуть 3B модель в один ускоритель, и успешно ее тренировать – готовь ускоритель с минимум 30 GB памяти ^[3], а лучше все 60.

Но если все таки условия для данного вида параллелизма выполняются – ни о чем больше не думайте, просто используйте его, так как масштабируется он практически линейно. Как написано в оригинальной статье, для достижения вычислительно-ограниченного (computation bound) состояния (то есть время, затрачиваемое на вычисления, больше времени, затрачиваемого на передачу данных – а значит ускоритель не простаивает), достаточно разместить на каждом устройстве батч размером от тысячи до нескольких тысяч токенов – копейки.

2. Fully-Sharded Data Parallelism (FSDP)

Или, как его еще называют ZeRO шардинг. Суть его в том, что веса, градиенты и состояния оптимизатора модели шардятся по ускорителям (такой вариант называется ZeRO-3, потому что шардятся все трое; если шардить например только градиенты и состояния оптимизатора – будет ZeRO-2, только состояния оптимизатора – ZeRO-1).

Формула:

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

Или вот FSDP для входной матрицы с точки зрения ^[4] схемы соединения нейронов:

Алгоритм FSDP

Прямое распространение

• собираем входную матрицу с разных ускорителей при помощи операции AllGather, получаем матрицу W_in[D, F]

• умножаем эту матрицу на матрицу активаций In[B_x, D] (ее собирать необязательно, можно просто оставить шардированной), получим временную матрицу Tmp[B_x, F]

• собираем через AllGather выходную матрицу W_out[F, D], умножаем на матрицу Tmp, получаем матрицу Out[B_x, D]

Обратное распространение

• во время обратного распространения ситуация для обеих матриц примерно одинаковая: нас интересуют только веса матрицы, находящиеся на конкретном шарде, больше они нигде не нужны, поэтому мы собираем только относящиеся к конкретному шарду градиенты/состояния оптимизатора при помощи операции ReduceScatter и обновляем веса данного шарда

• дальше раскидываем обновленные веса по всем ускорителям при помощи AllGather, дабы в итоге на всех ускорителях была одна и та же реплика модели

Полный алгоритм, если кого интересует:

Скрытый текст

То есть в случае с FSDP каждое устройство считает полный и окончательный набор выходов нейронов, никакого суммирования между устройствами не требуется. Но чтобы посчитать их, нужно собрать целиком все необходимые матрицы, отсюда AllGather весов – иначе просто не получится умножить. Поэтому данный вид шардирования и называется ZeRO, от “Zero Redundancy Optimizer” – потому что не надо хранить ничего лишнего, только все самое необходимое. И ничего лишнего пересылать тоже не надо.

Преимущества

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

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

Когда применять?

Условия достижения computation bound для FSDP примерно такие же как и для Data Parallelism. То есть если размер батча уже чуть больше чем ничего – значит все пучком, можем масштабироваться.

3. Тензорный параллелизм (Tensor Parallelism, TP)

Также он называемый 1D model parallelism или Megatron sharding (по имени модели, при тренировке которой этот метод впервые применили). Если в FSDP мы перекидывали между ускорителями веса, то в данном типе параллелизма ускорители обмениваются активациями.

Вот формула:

Ну то есть тупо взяли и поменяли шардируемые оси матриц

Схема разбиения входной матрицы:

Схемы соединения нейронов:

В отличие от FSDP, в котором шардинг входных данных (активаций) происходил вдоль оси батча, поэтому никаких дополнительных действий с ними делать не требовалось, в случае с TP активации шардятся по другой оси, а значит перед умножением их на входную матрицу их придется собрать со всех ускорителей. Это вычислительно менее затратно только в том случае, когда количество активаций меньше количества весов. Такого можно добиться только в случае если у нас уже применен FSDP, поэтому связку FSDP + TP можно встретить довольно часто.

Алгоритм TP

Прямое распространение

• Собрали входные данные через AllGather, получили матрицу активаций In[B, D]

• Умножаем матрицу активаций на входную матрицу W_in[D, F_y], получаем промежуточную матрицу Tmp[B, F_y]

• Дальше умножаем эту промежуточную матрицу на выходную матрицу W_out[F_y, D] и применяем операцию ReduceScatter, чтобы собрать нужные данные на нужных ускорителях, получаем в итоге матрицу Out[B, D_y]

Обратное распространение

• собираем градиенты для обновления весов выходной матрицы через AllGather, обновляем

• для входной матрицы ничего собирать особо не надо, по крайней мере если сохранили в кэше матрицу In[B, D], полученную во время прямого распространения

• в отличие от ZeRO, у TP во время обратного распространения будут встречаться операции обмена данных, которые нельзя отложить или выполнить заранее, и эта особенность ограничивает применение Tensor Parallelism

Полный алгоритм тут.

Скрытый текст

Когда применять?

В оригинальной статье приводится расчет на эту тему, но вас я грузить подробностями не хочу, поэтому просто скажу, что для тренировки типичной LLM оптимальное значение количества шардов при работе с TP примерно от 8 до 16. Но лучше конечно не применять TP в одиночку, а комбинировать с FSDP.

FSDP + TP

Формула:

На диаграмме выглядит так:

Схемы соединения нейронов:

Как уже говорилось выше, TP работает хорошо, когда размер входных данных меньше размера весов, а добиться этого можно применив FSDP, который как раз и уменьшает размер входного батча данных в пересчете на ускоритель. Таким образом, опуская долгие расчеты, наша система становится computation bound (а значит и масштабируемой), когда размер батча больше чем примерно 100 токенов (то есть практически в любой ситуации).

Вот собственно график сравнения TP, FSDP и FSDP + TP. Из него видно, что комбинация FSDP + TP является computation bound практически всегда.

4. Конвейерный параллелизм (Pipelining)

Pipelining это основная стратегия обучения ^[5] моделей на GPU. Идея проста: разбиваем модель по слоям и каждый слой или несколько слоев обучаем на отдельном GPU.

Алгоритм Pipelining:

Прямое распространение

• Инициализируем веса первого слоя модели на GPU 0 (или несколько слоев на GPU, или один слой на нескольких GPU, если использованы FSDP и/или TP).

• Прогоняем входные данные через первый слой на GPU 0, копируем активации на GPU 1 и так далее до тех пор, пока не дойдем до последнего GPU.

Обратное распространение

• Вычисляем функцию потерь Loss и ее градиенты dLoss/dx_L

• Для последнего слоя модели на последнем GPU вычисляем dLoss/dW_L и dLoss/dx_L-1, затем копируем dLoss/dx_L-1 на предпоследний GPU и так далее до тех пор, пока не дойдем до GPU 0.

Преимущества

Pipelining не требует передачи большого количества данных – только активации и градиенты, а это значит, что можно тренировать очень большие модели, даже если шина данных между ускорителями низкоскоростная. Собственно поэтому он часто применяется при работе с GPU, так как их топология подключения и скорость их шины данных зачастую уступает тем же TPU.

Недостатки

Вот значит приняли мы данные на GPU 0, обработали, отправили. А дальше мы будем сидеть и ждать, когда данные дойдут до последнего GPU в цепочке, посчитается функция потерь, пойдут обратно градиенты и т.д. И все это время GPU 0 будет простаивать, как и все остальные GPU в цепочке. Этот промежуток простоя называется “конвейерный пузырь” (pipeline bubble), и разбираться с ним довольно геморно.

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

Второй способ это попробовать совместить вычисление выходных активаций W_i @ x_i и градиентов dx = W_i @ dLoss / dx_i+1 и dW = dLoss / dx_i+1 @ x_i. Поскольку каждая из этих операций требует некоторого времени, мы можем наложить их друг на друга, таким образом избавившись от “пузыря”. Вот например график “беспузырного” конвейера из статьи про DeepSeek v3:

Когда применять?

Собственно когда скорость шины данных между ускорителями относительно невелика или же когда топология их соединения не позволяет быстро обмениваться данными. Такое часто бывает при работе с GPU.

Заключение

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

Ну а на этом все, подписывайтесь, чтобы не пропустить следующие статьи, до скорого!