Ultra-Scale Playbook, часть 2: параллелизм по данным, ZeRO, тензорный параллелизм

Вторая часть перевода книги The Ultra-Scale Playbook. Начало — в части 1 (память на одной GPU и коллективные операции), продолжение — в части 3 и части 4.

Как читать вместе с курсом ChinaAI

Это расширенная версия урока 5.4.3 «Распределённое обучение». Перед чтением полезно знать коллективные операции AllReduce, AllGather и ReduceScatter — они объяснены в приложении A0 из части 1. Практику DDP и FSDP на ваших 4 GPU даёт урок 5.7.4 «Несколько GPU на практике». Код в оригинале спрятан под спойлерами «Picotron (click to expand)»; мы даём ссылки прямо на файлы репозитория picotron.

Параллелизм по данным

Идея параллелизма по данным (data parallelism, DP) — скопировать модель на несколько GPU (копии мы называем «экземплярами модели») и параллельно на каждой GPU выполнять прямой и обратный проходы на разных микробатчах данных; отсюда и название. Скорее всего, вы уже встречали параллелизм по данным в простых примерах обучения, но здесь мы копнём заметно глубже, так что не уходите, даже если общий подход вам знаком.

flowchart TB
D["Глобальный батч"] --> M1["Микробатч 1"]
D --> M2["Микробатч 2"]
D --> M3["Микробатч 3"]
M1 --> G1["GPU 1: копия модели, прямой и обратный проход"]
M2 --> G2["GPU 2: копия модели, прямой и обратный проход"]
M3 --> G3["GPU 3: копия модели, прямой и обратный проход"]
G1 --> AR["all-reduce: усреднить градиенты"]
G2 --> AR
G3 --> AR
AR --> O["Одинаковый шаг оптимизатора на каждой GPU"]

Схема переводчика: параллелизм по данным

Если вы не знакомы с шаблонами распределённой коммуникации вроде broadcast, gather и all-reduce, мы подготовили короткий ускоренный курс — приложение A0.

Раз у каждой GPU свой микробатч, градиенты на разных GPU получатся разными. Чтобы экземпляры модели на разных GPU оставались одинаковыми, мы усредняем градиенты всех экземпляров операцией «all-reduce». Она выполняется во время обратного прохода, до шага оптимизатора.

Это наш первый примитив «распределённой коммуникации» — all-reduce: он отвечает за синхронизацию и обмен между экземплярами на GPU и узлах.

Наивная реализация DP просто дождалась бы конца обратного прохода, чтобы получить все градиенты, а затем запустила бы all-reduce по всем рангам DP, чтобы синхронизировать градиенты. Но такие последовательные шаги — сначала вычисления, потом коммуникация — БОЛЬШОЕ «НЕТ»: мы не хотим, чтобы GPU простаивали, пока идёт обмен, как на рисунке оригинала.

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

Рассмотрим три оптимизации, которые позволяют сделать намного лучше наивной первой реализации.

Первая оптимизация: перекрыть синхронизацию градиентов с обратным проходом

Главный недостаток только что описанного наивного DP: после обратного прохода (вычисления) приходится ждать синхронизации градиентов (коммуникации), прежде чем обновить параметры. Можно ли перекрыть эту коммуникацию вычислениями? Да!

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

В PyTorch это делается так: к каждому параметру прикрепляется хук-функция all-reduce. Операция all-reduce запускается, как только градиент этого параметра готов, а градиенты других параметров ещё вычисляются. Такой подход перекрывает большинство операций all-reduce с вычислением градиентов и повышает эффективность. Вот простая функция, которая прикрепляет хук:

def register_backward_hook(self, hook):
    """
    Регистрирует хук обратного прохода для всех параметров модели,
    которым нужны градиенты.
    """
    for p in self.module.parameters():
        if p.requires_grad is True:
            # хук сработает, как только градиент параметра будет накоплен
            p.register_post_accumulate_grad_hook(hook)

Перекрытие вычислений и коммуникации сокращает время ожидания синхронизации градиентов по всей модели. Синхронизация градиентов может идти (хотя бы частично) параллельно с обратным проходом в том же шаге обучения, что заметно ускоряет параллелизм по данным. Полная реализация наивного DP с перекрытием синхронизации — в Picotron: data_parallel.py.

Это наш первый пример перекрытия вычислений и коммуникации — ключевого приёма для максимальной эффективности масштабирования, к которому мы ещё не раз вернёмся в этой книге. Но эффективность можно повысить ещё больше!

Вторая оптимизация: бакеты градиентов

Операции на GPU обычно эффективнее на больших тензорах, чем множество операций на маленьких. Это верно и для операций коммуникации. Поэтому выгодно группировать градиенты в «бакеты» (buckets, «корзины») и запускать один all-reduce для всех градиентов бакета, а не отдельный all-reduce для каждого градиента.

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

Реализация с бакетами в Picotron — bucket.py.

Третья оптимизация: взаимодействие с накоплением градиентов

Наконец, как мы видели, накопление градиентов работает так: несколько прямых и обратных проходов, затем обновление параметров через optimizer.step(). Сочетая накопление градиентов с параллелизмом по данным, нужно внимательно выбирать момент синхронизации градиентов.

В наивной версии all-reduce автоматически запускается после каждого обратного прохода во время накопления. Это неоптимально: один reduce после последнего шага дал бы тот же результат с меньшими накладными расходами.

В PyTorch это обычно решается декоратором model.no_sync(), который отключает синхронизацию градиентов на обратных проходах, где свёртка не нужна.

Заметка

При операциях коммуникации тензоры должны лежать в памяти непрерывно, иначе понадобятся лишние копирования. Чтобы делать это оптимально, часто заранее выделяют непрерывные буферы размером с активации или параметры модели специально для коммуникации. Это ускоряет обмен, но отчасти увеличивает пиковое потребление памяти при обучении.

Теперь посмотрим, что это значит для глобального размера батча.

Снова о глобальном размере батча

Обновим формулу размера батча с учётом новых параметров — параллелизма по данным и накопления градиентов:

$$bs = gbs = mbs \times grad\_acc \times dp$$

Здесь $grad\_acc$ — число шагов накопления градиентов, а $dp$ — число параллельных экземпляров в параллелизме по данным.

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

На практике стараются как можно сильнее увеличить размер параллелизма по данным ($dp$) в ущерб накоплению градиентов ($grad\_acc$): параллелизм по данным по природе параллелен, а накопление градиентов последовательно. Накопление градиентов добавляют поверх параллелизма по данным, чтобы достичь целевого глобального батча, когда одного масштабирования DP не хватает — GPU закончились.

Хороший материал для дальнейшего чтения о параллелизме по данным — https://siboehm.com/articles/22/data-parallel-training.

Возможность распределить обучение по разным примерам даёт первое измерение параллелизации — это 1D-параллелизм (постепенно мы разберём ещё четыре измерения).

Наш путь до сих пор

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

1. Сначала определяем лучший (глобальный) размер батча в токенах — по литературе или экспериментами со сходимостью модели.
2. Затем выбираем длину последовательности для обучения — тоже по литературе или экспериментами. Обычно 2–8 тысяч токенов надёжно хорошо работают для современных оценочных бенчмарков (рецепты обучения мы здесь не разбираем, но команды обычно увеличивают длину последовательности в конце обучения, добавляя в смесь примеры с длинным контекстом, чтобы получить сегодняшние длинные контексты).
3. Теперь мы знаем размер батча ($gbs$). Максимальный локальный размер батча ($mbs$) на одной GPU находим, увеличивая его, пока не закончится память.
4. Наконец, определяем число доступных GPU для целевого $dp$. Отношение $gbs$ к $dp$ даёт оставшееся число шагов накопления градиентов, нужное для желаемого $gbs$.

Например, модели DeepSeek и Llama на основной фазе предобучения учили с длиной последовательности 4 тысячи токенов.

Почему для предобучения хорошо работают 2–8 тысяч токенов? Более длинные документы в вебе очень редки. Подробный анализ — в посте Харма де Фриса.

Если отношение для накопления градиентов меньше 1 — то есть GPU слишком много и мы «богаты GPU» 🤑 (!), — можно не использовать все GPU, попробовать больший $gbs$ или проверить, ускорит ли обучение меньший $mbs$. В последнем случае мы ставим пропускную способность выше эффективности вычислений отдельной GPU и берём $mbs$ меньше возможного, чтобы ускорить обучение.

Конкретный пример. Пусть мы хотим обучить современную модель с $gbs$ = 4 млн токенов и длиной последовательности 4 тысячи. Тогда размер батча — 1024 примера (берём ближайшую степень двойки). Допустим, в одну GPU помещается только $mbs = 2$, а для обучения у нас 128 GPU. Значит, с 4 шагами накопления градиентов мы получим нужные 1024 примера, или 4 млн токенов на шаг обучения. А если вдруг стало доступно 512 GPU? Тот же $gbs$ получится при $mbs = 2$ и 1 шаге накопления градиентов — и обучение пойдёт быстрее!

Пример переводчика: ваша машина

У вас $dp = 4$. Если цель — батч 64 последовательности по 2048 токенов (≈131 тыс. токенов) и в одну RTX 4080 помещается $mbs = 4$, то $grad\_acc = 64 / (4 \cdot 4) = 4$.

Заметка

Учтите, что начиная с 512 GPU и больше, в зависимости от сети, операции коммуникации начнут упираться в кольцевую задержку (ring latency — время, за которое сигнал один раз обходит кольцо), и коммуникацию DP уже не удастся полностью перекрыть. Это снизит эффективность вычислений и ударит по пропускной способности. Тогда пора исследовать другие измерения параллелизации.

Параллелизм по данным хорошо перекрывает синхронизацию градиентов через all-reduce с обратными вычислениями и экономит время, но на больших масштабах это преимущество начинает исчезать. Почему? Когда мы добавляем всё больше GPU (сотни или тысячи), накладные расходы на их координацию сильно растут, а требования к сети начинают перевешивать выгоду. В результате каждый новый GPU делает систему всё менее эффективной.

Это видно на бенчмарках оригинала: начиная с некоторого предела, пропускная способность заметно падает, а память на одну GPU остаётся постоянной и не зависит от добавления рангов DP.

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

Внимательный читатель, однако, уже заметил: предполагается, что в память GPU помещается прямой проход хотя бы одного входного примера ($mbs = 1$). Так бывает не всегда! Большие модели часто не помещаются в одну GPU даже с перевычислением активаций.

Совет

Минимальную память под параметры модели можно быстро прикинуть, умножив число параметров на 2: например, 70B → 140 ГБ (= 133 ГиБ).

Мы также видели, что начиная с некоторого масштаба у параллелизма по данным появляются ограничивающие накладные расходы на коммуникацию. Есть ли другие варианты для больших моделей или больших батчей? К счастью, есть: либо переносить часть тензоров на CPU, либо разрезать тензоры весов, градиентов и состояний оптимизатора между GPU.

Есть два основных подхода к разрезанию: параллелизм (тензорный, контекстный или конвейерный) и шардирование (DeepSpeed ZeRO или PyTorch FSDP). Эти подходы в какой-то мере ортогональны и их можно сочетать!

Шардирование тесно связано с DP, поэтому начнём с него — с метода ZeRO.

ZeRO: оптимизатор без избыточности

В этом разделе мы познакомимся с DeepSpeed ZeRO (Zero Redundancy Optimizer) — технологией оптимизации памяти, созданной, чтобы убрать избыточность памяти при обучении LLM.

Параллелизм по данным — эффективный способ масштабирования, но наивное копирование состояний оптимизатора, градиентов и параметров на каждый ранг DP создаёт большую избыточность. ZeRO устраняет её, разрезая состояния оптимизатора, градиенты и параметры по измерению параллелизма по данным и при этом позволяя вычислять с полным набором параметров. Иногда для этого нужно больше обмена между рангами DP, и его не всегда удаётся полностью перекрыть, как мы увидим дальше!

В книге мы сосредоточимся на ZeRO-1 — ZeRO-3: это даст общее представление о том, как технология сокращает потребление памяти, и покажет компромиссы, которые надо учитывать. Другие варианты ZeRO описаны в документации DeepSpeed.

Подход делится на три возможные стадии оптимизации:

• ZeRO-1: разбиение состояний оптимизатора;
• ZeRO-2: разбиение состояний оптимизатора и градиентов;
• ZeRO-3: разбиение состояний оптимизатора, градиентов и параметров.

«Разбиение» здесь означает разбиение по оси DP, потому что ZeRO — метод параллелизма по данным. Позже увидим, что разбивать можно и по другим осям.

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

Посмотрим внимательнее, сколько экономит разбиение на каждой стадии ZeRO.

Снова о потреблении памяти

Раньше мы обсуждали, сколько памяти занимают состояния оптимизатора, градиенты и параметры при обычном обучении. Обозначим число параметров модели через $\Psi$ (раньше это было $N$, но здесь используем обозначения исходной статьи о ZeRO). При обучении со смешанной точностью (подробнее — в части 4) и оптимизаторе Adam каждый хранимый объект занимает:

• параметры модели (половинная точность, BF16/FP16): $2\Psi$;
• градиенты модели (половинная точность, BF16/FP16): $2\Psi$;
• параметры модели в FP32 и состояния оптимизатора: $4\Psi + (4\Psi + 4\Psi)$;
• градиенты модели в FP32: $4\Psi$ (необязательно — только если мы хотим накапливать градиенты в FP32).

Без накопления градиентов в FP32 общее потребление памяти равно $2\Psi + 2\Psi + 12\Psi$, а с ним — $2\Psi + 6\Psi + 12\Psi$. Для простоты пока рассмотрим случай без накопления градиентов в FP32.

Идея ZeRO — шардировать эти объекты по рангам DP так, чтобы каждый узел хранил только срез. Срезы собираются заново, когда и если они нужны; так потребление памяти делится на степень параллелизма по данным $N_d$.

Стадия	Память на одну GPU
Обычный DP	$(2 + 2 + k)\,\Psi$
ZeRO-1	$2\Psi + 2\Psi + \frac{k\Psi}{N_d}$
ZeRO-2	$2\Psi + \frac{2\Psi + k\Psi}{N_d}$
ZeRO-3	$\frac{2\Psi + 2\Psi + k\Psi}{N_d}$

Здесь $\Psi$ — число параметров, $k$ — множитель памяти для состояний оптимизатора ($k = 12$ для Adam, как мы только что видели), $N_d$ — степень DP. (В оригинале это рисунок; таблицу по нему составил переводчик.)

Если вы используете накопление градиентов в FP32 с ZeRO-2 или ZeRO-3, к слагаемому градиентов нужно добавить ещё $\frac{4\Psi}{N_d}$.

Пример переводчика

Модель $\Psi = 1$ млрд параметров, $k = 12$, ваши $N_d = 4$ GPU. Обычный DP: $16$ ГБ на каждую GPU. ZeRO-1: $2 + 2 + 12/4 = 7$ ГБ. ZeRO-2: $2 + 14/4 = 5{,}5$ ГБ. ZeRO-3: $16/4 = 4$ ГБ. Активации сюда не входят.

Объясним эти формулы, разобрав работу каждой стадии ZeRO. Начнём с ZeRO-1.

ZeRO-1: разбиение состояний оптимизатора

В обычном DP все ранги после обратного прохода собирают одинаковые градиенты и одновременно выполняют одинаковые шаги оптимизатора. Выглядит как много дублированной работы. Нельзя ли её избежать и заодно сократить память?

В ZeRO-1 состояния оптимизатора делятся на $N_d$ равных частей, где $N_d$ — степень DP. То есть каждая копия модели на своём ранге DP отвечает только за $\frac{1}{N_d}$ состояний оптимизатора, и на шаге оптимизации обновляется только $\frac{1}{N_d}$ весов в FP32.

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

Отсюда и формула памяти $2\Psi + 2\Psi + \frac{k\Psi}{N_d}$ из таблицы выше! Вот последовательность операций одного шага обучения:

1. Прямой проход с одинаковым полным набором параметров BF16 на каждой копии, но с разными микробатчами.
2. Обратный проход с одинаковым полным набором градиентов на каждой копии, но с разными микробатчами.
3. Reduce-scatter градиентов (ещё один примитив — скоро объясним).
4. Каждая копия выполняет шаг оптимизатора на своих локальных состояниях (только $\frac{1}{N_d}$ состояний оптимизатора) и получает $\frac{1}{N_d}$ обновлённых параметров в FP32, которые затем можно перевести в $\frac{1}{N_d}$ полного набора параметров BF16.
5. All-gather параметров BF16, чтобы вернуть недостающие срезы каждой копии. Это новая операция в ZeRO, в обычном DP её нет.

Примечание: reduce-scatter в два раза быстрее all-reduce! Ура, третий примитив коммуникации!

Возможно, вам интересно, что это за операция «reduce-scatter» и как всё это выглядит. В оригинале есть анимация, которая проходит все шаги цикла прямого и обратного прохода.

flowchart LR
F["Прямой проход: все веса BF16"] --> B["Обратный проход: все градиенты"]
B --> RS["reduce-scatter градиентов: каждая GPU получает свою 1/N_d"]
RS --> OPT["Шаг оптимизатора только для своей 1/N_d параметров"]
OPT --> AG["all-gather весов BF16: собрать полную модель"]
AG --> F

Схема переводчика: один шаг ZeRO-1

С точки зрения коммуникации ZeRO-1 по сравнению с обычным DP заменяет all-reduce градиентов на reduce-scatter и добавляет all-gather всех параметров после шага оптимизатора.

Если вы следили за изложением, то помните: в обычном DP all-reduce градиентов можно перекрыть с вычислениями обратного прохода. В ZeRO-1 можно также подумать, как эффективно перекрыть новый all-gather параметров BF16. Для этого есть две основные стратегии.

• Во время шага оптимизатора. Можно запускать all-gather сразу после того, как оптимизатор обновил первый срез параметров. Тогда коммуникация может перекрыться с обновлением остальных параметров.
• Во время прямого прохода. Можно перекрывать all-gather параметров каждого слоя с прямым проходом.

Заметка

К сожалению, эти приёмы непросто реализовать: нужны изощрённые хуки и бакеты. На практике можно просто взять встроенную в PyTorch реализацию ZeRO-3/FSDP и сделать FSDPUnit всей моделью (подробнее — ниже).

В ZeRO-1 состояния оптимизатора разбиты, то есть каждая копия обновляет только $\frac{1}{N_d}$ состояний. Внимательный читатель мог заметить, что тогда незачем держать все градиенты на всех рангах DP: для шага оптимизации нужна только их часть. Встречайте ZeRO-2!

ZeRO-2: добавляем разбиение градиентов

Поскольку каждой копии нужен только тот шард градиентов, который соответствует её шарду состояний оптимизатора, логично шардировать и градиенты, так же как состояния оптимизатора. Тогда на обратном проходе вместо all-reduce градиентов выполняется только reduce-scatter! В памяти хранится только нужная $\frac{1}{N_d}$ часть градиентов, и памяти экономится больше, чем в ZeRO-1.

При накоплении градиентов в FP32 нужно хранить только $\frac{1}{N_d}$ градиентов FP32, в которые накапливаются градиенты BF16, пришедшие из reduce-scatter. А на шаге оптимизатора эти $\frac{1}{N_d}$ градиентов FP32 обновляют локальный шард состояний оптимизатора.

Теперь легко видеть, что шардирование градиентов даёт $2\Psi + \frac{2\Psi + k\Psi}{N_d}$, и при росте $N_d$ памяти нужно до 8 раз меньше базового уровня. По коммуникации всё так же, как в ZeRO-1, с той лишь разницей, что мы обмениваемся и освобождаем память на лету: в обоих случаях нужен reduce-scatter градиентов и all-gather всех параметров. Поэтому и ZeRO-2 по объёму коммуникации эквивалентен обычному DP.

Примечание: у ZeRO-2 по сравнению с ZeRO-1 нет реальных накладных расходов, кроме сложности реализации, и обычно ZeRO-2 — лучший вариант.

Мы шардировали и градиенты. Всё, или можно продолжать улучшения? Встречайте ZeRO-3!

ZeRO-3: добавляем разбиение параметров (FSDP)

На стадии 3 мы распространяем шардирование состояний оптимизатора и градиентов по копиям DP на параметры модели.

Заметка

Встроенная реализация этой стадии в PyTorch называется FSDP (Fully Sharded Data Parallelism, полностью шардированный параллелизм по данным). В книге мы будем называть её ZeRO-3, но везде, где встречается ZeRO-3, можно думать о FSDP.

Как же выполнять прямой и обратный проходы, если параметры модели разбросаны по GPU? Очень просто: собираем их по требованию, когда они нужны. На прямом проходе, идя по слоям, мы получаем нужные параметры по требованию и сразу выбрасываем их из памяти, когда они больше не нужны. Обратный проход работает так же, только в обратном порядке; на нём мы получаем шарды градиентов.

flowchart LR
subgraph FWD["Прямой проход, слой n"]
  A1["all-gather весов слоя n"] --> C1["вычислить слой n"] --> D1["выбросить чужие шарды"]
end
subgraph BWD["Обратный проход, слой n"]
  A2["all-gather весов слоя n"] --> C2["градиенты слоя n"] --> R2["reduce-scatter градиентов"] --> D2["выбросить чужие шарды"]
end
FWD --> BWD

Схема переводчика: ZeRO-3 собирает веса каждого слоя по требованию

Другая проблема в том, что эти all-gather нужно делать постоянно в течение всего прямого и обратного прохода, то есть на шаге обучения добавляется $2 \cdot \text{num\_layers} - 1$ операций all-gather по сравнению с ZeRO-2. Каждая приносит небольшую базовую задержку, как видно на рисунке оригинала.

На прямом проходе мы делаем all-gather параметров, когда они нужны, — это «налог» $\Psi$ на коммуникацию. Поскольку параметры выбрасываются сразу после использования на прямом проходе, на обратном проходе нужен ещё один all-gather — ещё один налог $\Psi$. Наконец, для градиентов нужен тот же reduce-scatter, что и в ZeRO-2, — тоже $\Psi$ коммуникации. Итого стоимость коммуникации $3\Psi$ против $2\Psi$ у ZeRO-2.

Звучит как большие накладные расходы, но на деле это не так страшно: коммуникацию параметров следующего слоя можно перекрыть с прямым проходом текущего слоя — это называется упреждающей загрузкой (prefetching). С ней мы делаем all-gather весов слоя n+1, пока идёт прямой проход слоя n, и аналогично all-gather весов слоя n−1, пока идёт обратный проход слоя n. Конечно, такое перекрытие работает, только пока мы не слишком раздули DP (эмпирическое правило: DP не должен превышать 512).

Примечание: «DP» у нас означает и сам приём параллелизма по данным, и число GPU, используемых для него (DP = размер DP = степень DP).

По памяти наша формула пришла к окончательному виду $\frac{2\Psi + 2\Psi + k\Psi}{N_d}$: теоретически потребление памяти можно уменьшать бесконечно, увеличивая размер DP, — по крайней мере для всего, что связано с параметрами модели. Обратите внимание: с промежуточными активациями это не помогает — для них, как мы видели, есть перевычисление активаций и накопление градиентов.

Подведём итог нашего пути по DP и ZeRO. С DP можно сильно увеличить пропускную способность обучения, просто добавляя копии модели. С ZeRO можно обучать даже модели, которые обычно не помещаются в одну GPU, шардируя параметры, градиенты и состояния оптимизатора по копиям DP ценой небольших затрат на коммуникацию.

Подробнее о FSDP1, FSDP2 и сложностях их реализации — в этом хорошем блоге.

Но и тут есть пределы: DP работает, только если слой модели помещается в одну GPU, а ZeRO разбивает лишь параметры, градиенты и состояния оптимизатора, но не память под активации! Вспомните из обсуждения памяти под активации, что эта часть памяти растёт с длиной последовательности и размером батча. Их можно просто ограничить, но на практике мы не хотим, чтобы железо вынуждало нас учить только на коротких последовательностях.

Чтобы обойти это, пора рассмотреть новую, ортогональную ось параллелизма — тензорный параллелизм (TP). В отличие от ZeRO-3, который опирается на интенсивный обмен параметрами, TP предлагает шардировать параметры, градиенты, состояния оптимизатора И активации между устройствами без пересылки параметров модели между GPU.

Что? Как такое вообще возможно?! Давайте вместе разберёмся в этом, казалось бы, волшебном подходе. 🙂

Тензорный параллелизм

Заметка

В оригинале и здесь можно послушать подкаст NotebookLM по следующим разделам.

Итак, мы шардировали параметры, градиенты и состояния оптимизатора модели с помощью ZeRO, но упёрлись в предел, когда память под активации превысила бюджет. Встречайте тензорный параллелизм (tensor parallelism, TP) — метод, который шардирует веса, градиенты и состояния оптимизатора, а также активации, и без необходимости собирать их все перед вычислением. Звучит как мечта! Сначала посмотрим, как TP работает на простых операциях умножения матриц (matmul).

Тензорный параллелизм использует математические свойства умножения матриц $A \times B$. Чтобы понять, как он работает, рассмотрим два основных равенства, которые делают такую параллелизацию возможной:

$$\begin{aligned} &\text{1.} \quad A\cdot B = A \cdot \begin{bmatrix} B_1 & B_2 & \cdots \end{bmatrix} = \begin{bmatrix} AB_1 & AB_2 & \cdots \end{bmatrix} \\ &\text{2.} \quad A\cdot B = \begin{bmatrix} A_1 & A_2 & \cdots \end{bmatrix} \begin{bmatrix} B_1 \\ B_2 \\ \vdots \end{bmatrix} = \sum_{i=1}^n A_i B_i \end{aligned}$$

Это значит, что произведение матриц можно вычислить, либо умножая на каждый столбец $B$ по отдельности, либо умножая каждую строку по отдельности и объединяя результаты. В нейросети умножение матриц чаще записывается как $X \times W$, где:

• $X$ — входы или значения активаций;
• $W$ — веса линейного слоя.

Посмотрим, как распараллелить эту операцию! В тензорном параллелизме тензоры разрезаются на $N$ шардов по какому-то измерению и распределяются по $N$ GPU. Матрицы можно резать по столбцам или по строкам — получается параллелизм по строкам или по столбцам. Как увидим дальше, шардирование по строкам и по столбцам требует разных примитивов коммуникации.

Первый вариант — шардирование по столбцам (column-wise, или column-linear): копируем полную входную матрицу на каждый воркер — для этого нужна операция broadcast — и разрезаем матрицу весов по столбцам. Входы умножаются на частичные матрицы весов, и в конце результаты объединяются операцией all-gather.

Реализация тензорного параллелизма по столбцам в Picotron — класс ColumnParallelLinear в tensor_parallel.py.

Второй вариант — шардирование по строкам (row-wise, или row-linear). Как догадается внимательный читатель, row-linear значит, что мы режем матрицу весов на блоки строк. Но тогда нужно разрезать и входы, поэтому вместо broadcast, как в column-linear, нужна операция scatter (наш четвёртый примитив распределённой коммуникации!). Результаты на каждом воркере уже нужной формы, но для итогового результата их надо сложить, поэтому здесь нужен ещё и all-reduce.

Реализация тензорного параллелизма по строкам — класс RowParallelLinear в том же файле tensor_parallel.py.

flowchart TB
subgraph COL["По столбцам: W = [W1 | W2]"]
  X1["X (копия на каждой GPU, broadcast)"] --> P1["GPU 1: X·W1"]
  X1 --> P2["GPU 2: X·W2"]
  P1 --> AG["all-gather: [X·W1 | X·W2]"]
  P2 --> AG
end
subgraph ROW["По строкам: W = [W1 ; W2]"]
  X2["X = [X1 | X2] (scatter)"] --> Q1["GPU 1: X1·W1"]
  X2 --> Q2["GPU 2: X2·W2"]
  Q1 --> AR["all-reduce: X1·W1 + X2·W2"]
  Q2 --> AR
end

Схема переводчика: два способа разрезать линейный слой

Теперь, когда у нас есть базовые строительные блоки TP, посмотрим, как эффективно сочетать их внутри слоя трансформера!

Тензорный параллелизм в блоке трансформера

Чтобы выработать стратегию, перейдём от игрушечного примера к реальному строительному блоку модели. Модель-трансформер состоит из двух основных блоков: полносвязного многослойного перцептрона (MLP) и многоголового внимания (MHA). Тензорный параллелизм можно применить к обоим.

Полносвязную часть можно распараллелить так: сначала разрезание по столбцам, затем по строкам. На прямом проходе это даёт broadcast для копирования входа и all-reduce. Заметьте, что broadcast при реальном обучении не нужен: можно позаботиться, чтобы входы уже были синхронизированы между рангами TP. Такая схема эффективнее, чем сначала по строкам, затем по столбцам: можно пропустить промежуточный all-reduce между разрезанными операциями.

flowchart LR
X["X (одинаковый на всех рангах TP)"] --> A1["GPU 1: X·A1 → GELU"]
X --> A2["GPU 2: X·A2 → GELU"]
A1 --> B1["GPU 1: ·B1"]
A2 --> B2["GPU 2: ·B2"]
B1 --> AR["all-reduce: сумма"]
B2 --> AR
AR --> Y["Выход MLP"]

Схема переводчика: MLP с тензорным параллелизмом — A режется по столбцам, B по строкам, одна коммуникация в конце

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

Здесь в целом можно действовать так же: матрицы запросов (Q), ключей (K) и значений (V) разрезаются по столбцам, а выходную проекцию можно считать row-linear. У многоголового внимания разрезание по столбцам имеет очень естественную интерпретацию: каждая GPU вычисляет внимание для одной головы или подмножества голов. Тот же подход работает и для внимания с несколькими запросами (multi-query attention, MQA) и внимания с группами запросов (grouped query attention, GQA), где ключи и значения общие для нескольких запросов.

Тензорный параллелизм так эффективно применяется к блокам внимания и MLP, потому что у них есть естественно независимые измерения. Блок внимания можно распараллелить по измерению num_attention_heads, потому что каждая голова внимания работает независимо. Так же блок MLP можно распараллелить по измерению hidden_dim: операции в полносвязной сети независимы вдоль этого измерения.

Стоит, однако, заметить, что степень тензорного параллелизма не должна превышать число голов внимания, потому что проекция QKV шардируется по измерению num_attention_heads. При GQA у нас num_attention_heads голов запросов, но лишь num_kv_heads голов ключей и значений (num_attention_heads ≥ num_kv_heads). В этом случае всё равно можно взять TP = num_attention_heads, но придётся следить, чтобы головы K/V были правильно синхронизированы между GPU. Например, у Llama-3 8B 32 головы запросов, но только 8 голов ключей и значений, так что степень TP теоретически может доходить до 32, но потребуется аккуратная реализация, чтобы поддерживать синхронизацию голов K/V между воркерами тензорного параллелизма.

Заметим также, что тензорный параллелизм — не серебряная пуля. Мы добавили несколько примитивов распределённой коммуникации прямо в вычислительный путь модели, и их трудно полностью спрятать или перекрыть вычислениями (как мы делали в ZeRO). Поэтому итоговая производительность — компромисс между выигрышем в вычислениях и памяти и добавленными накладными расходами на коммуникацию.

Эту коммуникацию можно частично спрятать, выполняя умножение блочных матриц вместе с асинхронной коммуникацией и вычислениями.

Глядя на временную шкалу операций в MLP с тензорным параллелизмом (для MHA то же самое), лучше видны компромиссы. На прямом проходе каждого слоя декодера есть точка синхронизации — операция all-reduce, которую нельзя перекрыть вычислениями. Эти открытые накладные расходы на коммуникацию (exposed communication) необходимы, чтобы объединить частичные результаты рангов тензорного параллелизма, прежде чем можно будет применить финальный LayerNorm.

Например, Megatron-LM и Nanotron реализуют частичное перекрытие all-gather с вычислением полносвязного слоя (FC1): часть результата умножения матриц отправляется на другую GPU, пока оставшаяся часть ещё вычисляется.

Тензорный параллелизм помогает сократить память под активации для умножений матриц, потому что промежуточные активации шардированы между GPU. Однако для операций вроде LayerNorm полные активации всё равно нужно собирать, то есть мы получаем не весь возможный выигрыш в памяти. Кроме того, TP вносит большие требования к коммуникации, сильно зависящие от сетевой инфраструктуры. Поскольку этот all-reduce нельзя полностью спрятать за вычислениями, он прямо добавляется к критическому пути прямого распространения — последовательности операций, определяющей минимальное время прямого прохода.

Это активная область исследований; например, недавняя работа Domino исследует новые приёмы, позволяющие максимально увеличить такое перекрытие.

Посмотрим внимательнее на компромисс при росте степени TP. На графиках оригинала рост TP снижает пропускную способность на одну GPU (слева), но позволяет обрабатывать бо́льшие батчи (справа) — это иллюстрирует компромисс между эффективностью вычислений и доступной памятью в распределённом обучении.

На практике, как видно на левом графике, накладные расходы тензорного параллелизма особенно заметны при выходе за 8 GPU. Внутри одного узла тензорный параллелизм может опираться на быстрые соединения NVLink, а между узлами нужна более медленная сеть. Мы наблюдаем сильное падение при переходе от TP=8 к TP=16 и ещё более резкое — от TP=16 к TP=32. При высоких степенях параллелизма накладные расходы на коммуникацию становятся так велики, что быстро начинают доминировать над временем вычислений.

Примечание переводчика: ваша машина

У вас «быстрой» связи нет даже внутри машины: 4 карты общаются через PCIe 4.0. Поэтому тензорный параллелизм на вашей машине обычно проигрывает FSDP/ZeRO по скорости, и его стоит включать, только когда без него модель не помещается в память.

Тем не менее тензорный параллелизм даёт важный выигрыш по памяти, распределяя параметры модели, градиенты, состояния оптимизатора и (в какой-то мере) активации между GPU. На примере модели на 70 млрд параметров в оригинале видно: с ростом степени тензорного параллелизма память под параметры, градиенты и состояния оптимизатора на каждой GPU уменьшается настолько, что большую модель уже можно начать размещать на одном узле из 8 GPU.

Можно ли выжать из этого приёма ещё больше? Нормализация слоя (LayerNorm) и dropout всё ещё требуют собирать полные активации на каждой GPU, отчасти сводя на нет экономию памяти. Можно сделать лучше, распараллелив и эти оставшиеся операции.

Заметка

Интересная деталь про нормализацию слоя при тензорном параллелизме: после all-gather каждый ранг TP видит одни и те же активации, поэтому градиенты весов LayerNorm не нужно синхронизировать через all-reduce после обратного прохода — они естественно остаются одинаковыми на всех рангах. А вот для dropout нужно синхронизировать зерно генератора случайных чисел между рангами TP, чтобы поведение было детерминированным.

Дальше разберём небольшое естественное расширение тензорного параллелизма — параллелизм по последовательности (sequence parallelism), который делает ровно это.

Параллелизм по последовательности

Параллелизм по последовательности (sequence parallelism, SP) разрезает активации и вычисления для тех частей модели, которые не охвачены тензорным параллелизмом, — dropout и LayerNorm, — но по измерению входной последовательности, а не по скрытому измерению.

Заметка

Термин «параллелизм по последовательности» перегружен. Тот SP, о котором речь в этом разделе, тесно связан с тензорным параллелизмом и применяется к dropout и нормализации слоя. Но при переходе к длинным последовательностям узким местом станет вычисление внимания, а для него нужны приёмы вроде Ring Attention. Их тоже иногда называют параллелизмом по последовательности, но мы, чтобы различать два подхода, будем называть их контекстным параллелизмом. Поэтому, когда в книге встречается «параллелизм по последовательности», помните: он используется вместе с тензорным параллелизмом (в отличие от контекстного, который можно применять независимо).

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

$$\text{LayerNorm}(x) = \gamma \cdot \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta$$

где $\mu = \text{mean}(x)$ и $\sigma^2 = \text{var}(x)$ вычисляются по скрытому измерению $h$.

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

На схеме оригинала показано, как мы переходим между областями тензорного параллелизма и параллелизма по последовательности с помощью разных коллективных операций (обозначенных f и g). Подпись к схеме:

• на прямом проходе: f = ничего не делать (no-op); f* = all-reduce; g = all-gather; g* = reduce-scatter;
• на обратном проходе: f = all-reduce; f* = no-op; g = reduce-scatter; g* = all-gather;
• области SP нужно полное измерение hidden_dim.

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

При тензорном параллелизме на прямом проходе:

• f — no-op (пустая операция), потому что активации и так продублированы на всех рангах;
• f* — all-reduce, чтобы синхронизировать активации и обеспечить корректность.

А на обратном проходе:

• f* — no-op, потому что градиенты уже продублированы на всех рангах;
• f — all-reduce, чтобы синхронизировать градиенты.

Операции f и f* называют сопряжённой парой (conjugate pair): они дополняют друг друга — на каждом проходе, когда одна из них no-op, другая — all-reduce, а на другом проходе наоборот.

Для параллелизма по последовательности используются другие операции, g и g*. В частности, мы избегаем all-reduce в областях SP: он потребовал бы собрать полные активации и увеличил бы пиковое потребление памяти, сводя на нет смысл SP.

Что же здесь на самом деле происходит? Как сказала бы одна известная LLM, разберёмся шаг за шагом.

Начальный слой LayerNorm (область SP)

• Входные тензоры X1* и X2* формы (b, s/2, h) поступают уже разрезанными по измерению последовательности.
• Каждая GPU независимо вычисляет LayerNorm на своём куске последовательности и получает Y1* и Y2*.

Первый переход (SP → TP)

• Операция g (all-gather) объединяет Y1 и Y2 обратно в полную длину последовательности.
• Восстанавливается Y формы (b, s, h), потому что слоям column-linear нужно полное скрытое измерение h.

Первый линейный слой (область TP)

• A1 и A2 — слои column-linear, поэтому они разрезают Y по скрытому измерению.
• GELU применяется на каждой GPU независимо.
• Z1* и Z2* имеют форму (b, s, h/2).

Второй линейный слой (область TP)

• B1 и B2 — слои row-linear, поэтому они восстанавливают скрытое измерение.
• W1 и W2 формы (b, s, h) нужно сложить.

Последний переход (TP → SP)

• Операция g* (reduce-scatter) выполняет свёртку, нужную для корректности предыдущего row-linear, и одновременно разрезает по измерению последовательности.
• W1* и W2* имеют форму (b, s/2, h).

Ключевое преимущество параллелизма по последовательности — он уменьшает максимальный размер активаций, которые нужно хранить. С одним тензорным параллелизмом нам приходилось в разных местах хранить активации формы (b, s, h). С SP максимальный размер активаций уменьшается до $\frac{b \cdot s \cdot h}{tp}$, потому что мы всегда режем либо по последовательности, либо по скрытому измерению.

Уследить за всеми частями, которые по-разному шардированы в TP и TP+SP, непросто — поверьте, нам самим трудно это держать в голове, поэтому мы сделали небольшую таблицу: как меняется форма активаций (hidden_states) по скрытому измерению h и по последовательности s на прямом проходе.

Область	Только TP	TP с SP
Вход в TP (column-linear)	h: шардировано (weight_out шардирован); s: полное	h: шардировано (weight_out шардирован); s: all-gather до полного
Область TP	h: шардировано; s: полное	h: шардировано; s: полное
Выход из TP (row-linear)	h: полное (weight_out полный + all-reduce для корректности); s: полное	h: полное (weight_out полный + reduce-scatter для корректности); s: reduce-scatter до шардированного
Область SP	h: полное; s: полное	h: полное; s: шардировано

И для слоя эмбеддингов:

Область	Обычный TP	TP с SP
Слой эмбеддингов (row-linear, шардирован по словарю)	h: полное (weight_out полный + all-reduce для корректности); s: полное	h: полное (weight_out полный + reduce-scatter для корректности); s: reduce-scatter до шардированного

С параллелизмом по последовательности можно сэкономить ещё больше памяти под активации и увеличить размер батча и длину последовательности сильнее, чем с одним тензорным параллелизмом. Для того же примера модели 70B в оригинале видно: максимальное потребление памяти на GPU снова сильно сократилось, и с TP+SP=16 помещаются последовательности длиной 16 тысяч токенов — улучшение по сравнению с обычным TP! (TP=16 всё ещё многовато, как мы видели в предыдущем разделе, но в следующем разделе увидим, как это улучшить.)

Возможно, вы спросите: не добавляет ли TP+SP больше накладных расходов на коммуникацию, чем обычный TP? И да, и нет. На прямом проходе обычного TP было две операции all-reduce на блок трансформера, а в SP — две all-gather и две reduce-scatter на блок. То есть SP делает вдвое больше операций коммуникации, чем TP. Но поскольку all-reduce можно разложить на all-gather и reduce-scatter (см. раздел о кольцевом AllReduce в приложении), по стоимости коммуникации они на самом деле эквивалентны. То же рассуждение верно и для обратного прохода: мы просто используем сопряжённую операцию (no-op ↔ all-reduce и all-gather ↔ reduce-scatter).

Если вы читали внимательно, то заметили, что речь о четырёх операциях коммуникации в каждом слое (две для внимания и две для MLP). На профиле MLP с TP+SP в оригинале это хорошо видно.

Как и обычный TP, TP+SP трудно перекрыть с вычислениями, поэтому пропускная способность сильно зависит от пропускной способности связи. И здесь, как и в обычном TP, TP+SP обычно используют только внутри узла (степень TP не больше числа GPU в узле, например TP ≤ 8).

Можно измерить, как эти накладные расходы становятся всё более проблемными при росте тензорного параллелизма. В оригинале измерены пропускная способность и память при масштабировании TP с SP для модели на 3 млрд параметров с длиной последовательности 4096. И снова компромисс между эффективностью вычислений (слева) и объёмом памяти (справа): более высокие степени параллелизма позволяют обрабатывать заметно бо́льшие батчи, уменьшая память под активации, но снижают пропускную способность на GPU — особенно после порога, равного числу GPU в узле.

Подытожим наблюдения:

• для обоих методов самое сильное падение производительности — при переходе от TP=8 к TP=16, потому что именно тогда мы переходим от коммуникации внутри одного узла (NVLink) к коммуникации между узлами (EFA);
• экономия памяти под активации при TP с SP позволяет уместить гораздо бо́льшие батчи, чем с одним TP.

Мы увидели, как TP помогает шардировать активации между несколькими GPU, разрезая операции внимания и полносвязной сети по скрытому измерению, и как SP естественно дополняет его для оставшихся операций, разрезая по измерению последовательности.

Заметка

Поскольку слои LayerNorm в области SP работают на разных частях последовательности, их градиенты на разных рангах TP различаются. Чтобы веса оставались синхронизированными, на обратном проходе их градиенты нужно свернуть через all-reduce — так же, как DP обеспечивает синхронность весов. Впрочем, накладные расходы малы: параметров у LayerNorm сравнительно немного.

И всё же у TP+SP два ограничения: если увеличивать длину последовательности, память под активации всё равно «взорвётся» в области TP, а если модель слишком велика, чтобы поместиться при TP=8, начнётся сильное замедление из-за связи между узлами.

Первую проблему решает контекстный параллелизм (context parallelism), вторую — конвейерный параллелизм (pipeline parallelism). Начнём с контекстного параллелизма — в части 3.

Словарик

EN	RU
data parallelism (DP), DP rank, DP degree	параллелизм по данным, ранг DP, степень DP
model instance / replica	экземпляр (копия) модели
overlap communication and computation	перекрывать коммуникацию и вычисления
bucketing	группировка градиентов в бакеты
sharding / partitioning	шардирование / разбиение
ZeRO (Zero Redundancy Optimizer)	ZeRO, оптимизатор без избыточности
FSDP (Fully Sharded Data Parallel)	полностью шардированный параллелизм по данным
prefetching	упреждающая загрузка
tensor parallelism (TP), column-linear / row-linear	тензорный параллелизм, разрезание по столбцам / по строкам
exposed communication, critical path	открытая (неперекрытая) коммуникация, критический путь
sequence parallelism (SP)	параллелизм по последовательности
conjugate pair	сопряжённая пара операций
intra-node / inter-node	внутри узла / между узлами