Ultra-Scale Playbook, часть 3: контекстный, конвейерный и экспертный параллелизм, 5D

Третья часть перевода книги The Ultra-Scale Playbook. Предыдущие части: 1 — память на одной GPU, 2 — DP, ZeRO, TP и SP. Следующая: 4 — внутри GPU: ядра и FlashAttention.

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

Контекстный параллелизм связан с уроком о длинном контексте (ступень 5.5), экспертный — с уроком о MoE (ступень 5.3), конвейерный и сводка «5D» — с уроком 5.4.3 «Распределённое обучение». Раздел «Поиск лучшей конфигурации» — практичный чек-лист: перечитайте его перед тем, как запускать обучение на нескольких GPU в уроке 5.7.4.

Контекстный параллелизм

С тензорным параллелизмом и параллелизмом по последовательности можно сильно уменьшить требования к памяти на одну GPU, потому что и веса модели, и активации распределены по GPU. Однако при обучении на всё более длинных последовательностях (например, при переходе к 128 тысячам токенов и более) мы всё ещё можем превысить память одного узла, потому что внутри области TP по-прежнему обрабатываем полную последовательность.

Более того, даже при полном перевычислении активаций (а оно даёт тяжёлые накладные расходы на вычисления, около 30 %) в памяти всё равно нужно держать часть активаций на границах слоёв, а они растут линейно с длиной последовательности. Посмотрим, как здесь поможет контекстный параллелизм (context parallelism, CP).

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

С контекстным параллелизмом, как и с параллелизмом по последовательности, мы режем вход по измерению последовательности, но теперь разрезание применяется ко всей модели, а не только к областям параллелизма по последовательности, как раньше в TP+SP.

Разрезание последовательности не влияет на большинство модулей, например MLP и LayerNorm, где каждый токен обрабатывается независимо. Оно и не требует дорогой коммуникации, как TP, потому что режутся только входы, а не матрицы весов. Как и в параллелизме по данным, после вычисления градиентов запускается all-reduce, чтобы синхронизировать градиенты в группе CP.

Есть, однако, важное исключение: особого внимания требуют блоки внимания (ха-ха, каламбур намеренный :D). В модуле внимания каждому токену нужен доступ к парам ключ/значение всех остальных токенов последовательности (или, при причинном внимании, хотя бы всех предыдущих токенов).

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

При наивной реализации это звучит очень дорого. Можно ли сделать дешевле и быстрее? К счастью, есть ключевой приём, позволяющий эффективно организовать обмен парами ключ/значение, — Ring Attention («кольцевое внимание»).

Заметка

У контекстного параллелизма есть концептуальное сходство с FlashAttention, о котором речь пойдёт позже: оба опираются на онлайн-вычисление softmax, чтобы сократить память. Но FlashAttention оптимизирует само вычисление внимания на одной GPU, а контекстный параллелизм экономит память, распределяя последовательность по нескольким GPU.

Ring Attention

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

Проиллюстрируем. Пусть у нас четыре GPU и вход из четырёх токенов. Сначала входная последовательность поровну режется по измерению последовательности, так что на каждой GPU оказывается один токен с его значениями Q/K/V. Пусть Q1, K1 и V1 — запрос, ключ и значение первого токена; они лежат на первой GPU. Вычисление внимания займёт четыре шага по времени. На каждом шаге каждая GPU выполняет три последовательные операции.

1. Неблокирующе отправить текущие ключи и значения следующей машине (на всех шагах, кроме последнего), чтобы начать следующую операцию, не дожидаясь окончания этой.
2. Локально вычислить оценки внимания на текущих ключах и значениях — обычно это $\text{Softmax}\left(\frac{QK^T}{\sqrt{d}}\right) \cdot V$.
3. Дождаться ключей и значений от предыдущей GPU и вернуться к шагу 1, где текущими ключами и значениями станут только что полученные.

Эти три шага повторяются четыре раза, и вычисление внимания завершено.

flowchart LR
G1["GPU 1: Q1, держит K/V блока 1"] -- "K/V" --> G2["GPU 2: Q2, держит K/V блока 2"]
G2 -- "K/V" --> G3["GPU 3: Q3, держит K/V блока 3"]
G3 -- "K/V" --> G4["GPU 4: Q4, держит K/V блока 4"]
G4 -- "K/V" --> G1

Схема переводчика: в Ring Attention запросы остаются на месте, а блоки ключей и значений ходят по кольцу

В оригинале весь процесс на четырёх GPU показан анимацией, и сразу понятно, почему авторы назвали подход кольцевым!

Есть, однако, большая проблема: наивная реализация Ring Attention приводит к сильному дисбалансу между GPU из-за формы матрицы причинного внимания. Посмотрим на вычисление softmax для матрицы оценок внимания с причинной маской.

Softmax вычисляется по строкам, то есть как только GPU получила все токены строки, её можно вычислить. GPU 1 может сделать это сразу: она начинает с токенов 1–4 и ни от кого ничего не ждёт. А GPU 2 придётся ждать второго раунда, чтобы получить токены 1–4 и иметь все значения для токенов 1–8. К тому же GPU 1, похоже, делает гораздо меньше работы, чем остальные.

Посмотрим, можно ли лучше сбалансировать вычисления.

Зигзагообразный Ring Attention — сбалансированная реализация

Нужен лучший способ распределить входные последовательности. Для этого токены раздают GPU не строго по порядку, а слегка перемешивают, чтобы на каждой GPU была хорошая смесь ранних и поздних токенов. Такой подход называется Zig-Zag Attention («зигзагообразное внимание»). При новой раскладке, если посчитать закрашенные клетки маски внимания, видно, что вычисления теперь поровну распределены между всеми GPU.

Здесь показан Zig-Zag Attention, который немного отличается от Striped Attention. Подробнее о различиях — в этом обсуждении на GitHub.

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

Пусть последовательность из 8 блоков и 4 GPU. Наивно: GPU 1 — блоки 1–2, GPU 2 — 3–4, GPU 3 — 5–6, GPU 4 — 7–8; у GPU 4 работы для причинного внимания в несколько раз больше, чем у GPU 1. Зигзагом: GPU 1 — блоки 1 и 8, GPU 2 — 2 и 7, GPU 3 — 3 и 6, GPU 4 — 4 и 5. У каждой пары «ранний + поздний» блок одинаковый суммарный объём работы.

Видно также, что для завершения всех строк каждой GPU нужна информация от всех остальных GPU.

Есть два общих способа перекрыть вычисления и коммуникацию: либо сделать общий all-gather, собрав все ключи и значения на каждой GPU одновременно (в духе ZeRO-3), либо собирать их с каждой GPU на каждую по мере надобности.

Ключевые различия двух реализаций — в шаблонах коммуникации и потреблении памяти.

1. Реализация через all-gather:

• все GPU одновременно собирают полные пары ключ/значение со всех остальных GPU;
• нужно больше временной памяти, потому что каждая GPU должна хранить все пары K/V сразу;
• коммуникация проходит за один шаг, но с большими накладными расходами по памяти.

2. Реализация all-to-all (кольцо):

• GPU обмениваются парами K/V по кольцу, по одному куску за раз;
• экономнее по памяти: каждой GPU нужно временно хранить только один дополнительный кусок;
• коммуникация растянута и перекрыта вычислениями, хотя несколько шагов коммуникации добавляют базовую задержку.

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

Реализация кольцевого внимания в Picotron — context_parallel.py.

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

Но мы по-прежнему знаем, что TP плохо масштабируется между узлами. Что делать, если веса модели не помещаются в один узел? На помощь приходит конвейерный параллелизм — наша четвёртая степень параллелизма!

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

В разделе «Тензорный параллелизм» мы видели, что попытка масштабировать TP дальше числа GPU в одном узле (обычно 4 или 8) вынуждает использовать сеть с меньшей пропускной способностью, а это сильно ухудшает производительность. Влияние межузловой коммуникации хорошо видно на бенчмарках all-reduce на нашем кластере из нескольких узлов (в каждом узле 8 GPU). В оригинале — график пропускной способности межузловой связи при разном числе узлов: медианы (линии) и диапазоны от 5-го до 95-го перцентиля (заливка) для all-reduce, all-gather и reduce-scatter.

Параллелизм по последовательности и контекстный параллелизм помогают на длинных последовательностях, но мало помогают, если корень проблем с памятью не в длине последовательности, а в размере самой модели. Для больших моделей (70B+ параметров) уже одни веса могут не поместиться в 4–8 GPU одного узла. Эту проблему решает ещё одно измерение параллелизма — конвейерный параллелизм (pipeline parallelism, PP).

Конвейерный параллелизм — простой, но мощный приём: мы раскладываем слои модели по нескольким GPU! Например, при 8 GPU можно положить слои 1–4 на GPU 1, слои 5–8 на GPU 2 и так далее. Тогда каждой GPU нужно хранить и обрабатывать только часть слоёв модели, что сильно снижает требования к памяти на одну GPU. В оригинале показано, как конвейерный параллелизм влияет на память модели с 8 млрд параметров.

Этот приём может напомнить обсуждение ZeRO-3, где мы делили параметры модели между GPU. Подробно мы сравним оба приёма ниже, в разделе «5D-параллелизм в двух словах».

На графике оригинала заметна интересная вещь: параметры модели хорошо распределяются по GPU, а память под активации на каждой GPU остаётся прежней! То есть активационной памяти этот подход не экономит.

Заметка

Так происходит потому, что каждая GPU должна выполнить PP прямых проходов, прежде чем начать первый обратный. Каждая GPU отвечает за $1/PP$ слоёв, но до первого обратного прохода обрабатывает PP микробатчей, поэтому в итоге хранит $PP \times (activs / PP) \approx activs$, то есть требования к памяти под активации остаются примерно такими же, как без конвейерного параллелизма.

Появляется новый шаблон коммуникации: вместо пересылки параметров, как в ZeRO-3 при параллелизме по данным, мы теперь последовательно передаём тензоры активаций между GPU по «конвейеру». Идея проста, но эффективно реализовать её довольно непросто. Перейдём к деталям!

Раскладываем слои по узлам: «все прямые, все обратные»

Для начала просто разложим слои по нескольким устройствам: первая GPU возьмёт несколько первых слоёв, вторая — вторую часть модели и так далее. Прямой проход через модель теперь сводится к последовательной передаче батча данных по модели и, значит, к поочерёдному использованию каждого вычислительного устройства.

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

Но вы, возможно, уже предчувствуете неприятности: «последовательно» и «поочерёдно»?!? В мире параллельных вычислений это звучит не очень эффективно, особенно после разговора о перекрытии вычислений и коммуникации.

Именно, читатель! Главная трудность конвейерного параллелизма — как эффективно обойти последовательную природу PP, чтобы GPU всё время были заняты и не было так, что одна GPU считает, а остальные ждут. На рисунке оригинала — загрузка GPU при наивном прямом и обратном проходе через модель с 16 слоями на 4 GPU (числа — номера слоёв).

Оставшееся время простоя показано серым и обычно называется «пузырём» (bubble). Вид этого пузыря, наверное, разбил вам сердце после всех усилий по оптимизации пропускной способности.

Насколько эффективна конвейерная схема, можно оценить по тому, сколько времени теряется из-за пузыря. Пусть $t_f$ и $t_b$ — время прямого и обратного проходов одного микробатча на одной стадии конвейера (часто просто считают $t_b \approx 2 \times t_f$, как на графике выше). При идеальной параллелизации общее время было бы $t_{id} = t_f + t_b$. Но в этом примере из-за пузыря добавляется время $t_{pb} = (p-1) \cdot (t_f + t_b)$, где $p$ — степень конвейерного параллелизма, то есть число GPU. Это время, которое каждая GPU ждёт, пока считают другие.

Отношение лишнего времени пузыря к идеальному времени:

$$r_{bubble} = \frac{(p-1) \cdot (t_f + t_b)}{t_f + t_b} = p - 1$$

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

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

Достанем первый инструмент: разрежем батч на маленькие порции, которые можно обрабатывать параллельно (или почти), как мы уже делали в DP. Теперь, пока вторая GPU занята микробатчем 1, первая уже может начать микробатч 2. В оригинале показано расписание с восемью микробатчами.

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

Такое расписание называют «все прямые, все обратные» (all forward, all backward, AFAB): сначала выполняются все прямые проходы, затем все обратные. Преимущество в том, что прямые и обратные шаги в целом остаются последовательными, и общая организация кода обучения сохраняется. Это одна из самых простых в реализации схем PP.

Полная реализация конвейера AFAB есть в Picotron — функция train_step_pipeline_afab в pipeline_parallel.py.

Оценим пузырь в этом примере. Отличие от первого примера в том, что идеальное время обработки $m$ микробатчей теперь $t_{id} = m \cdot (t_f + t_b)$:

$$r_{bubble} = \frac{(p-1) \cdot (t_f + t_b)}{m \cdot (t_f + t_b)} = \frac{p-1}{m}$$

Как видим, часть неэффективности стадий конвейера можно побороть, добавив больше микробатчей: пузырь уменьшается в $m$ раз.

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

На ваших 4 GPU ($p = 4$) с наивным конвейером $r_{bubble} = 3$: три четверти времени карты простаивают. С $m = 8$ микробатчами $r_{bubble} = 3/8 \approx 0{,}38$, с $m = 32$ — около 0,09.

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

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

«Один прямой, один обратный» и схема Llama 3.1

Это расписание называется «один прямой, один обратный» (one forward, one backward, 1F1B): в средней, установившейся фазе по очереди выполняются один прямой и один обратный проход. Общая идея — начинать обратный проход как можно раньше.

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

Если аккуратно посчитать, пузырь остался того же размера, так что эффективность обучения заметно не выросла. Зато хранить нужно активации только для $p$ микробатчей ($p$ — степень конвейерного параллелизма), а не для $m$ ($m$ — число микробатчей), и это сдерживает взрыв памяти под активации, который был в AFAB. Как следствие, можно добавить больше микробатчей, а это уже реально уменьшит пузырь.

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

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

Полная реализация 1F1B тоже есть в Picotron — функция train_step_pipeline_1f1b в pipeline_parallel.py.

Посмотрим, как расписание 1F1B масштабируется на практике, по бенчмаркам на нашем кластере (график в оригинале: пропускная способность конвейерного параллелизма при разном числе микробатчей).

Слева, когда число микробатчей не больше степени PP минус один ($m = p - 1$), видно, насколько вреден пузырь: производительность низкая и даже падает с ростом PP. Справа видно, что микробатчей намного больше степени PP ($m = 32 \gg p - 1$) улучшают производительность при малых степенях PP, хотя при очень больших степенях PP она всё равно ограничена. На практике нельзя бесконечно увеличивать число микробатчей, чтобы сохранить $m \gg p - 1$: нас ограничивает целевой глобальный размер батча. При максимально возможном числе микробатчей с ростом степени PP пузырь в итоге придётся увеличить в соответствии с $r_{bubble} = \frac{p - 1}{m}$.

Интересно, что при малом числе микробатчей производительность падает всего на 14 % при переходе от одного узла ($p = 8$) к двум ($p = 16$) — это гораздо лучшее масштабирование, чем у тензорного параллелизма, который в похожих межузловых сценариях обычно теряет около 43 %. Такое поведение при выходе в медленную межузловую сеть делает конвейерный параллелизм особенно привлекательным для распределённого обучения на нескольких узлах.

1F1B заметно сокращает память под активации, но на последнем графике видно, что пузырь конвейера остаётся главным узким местом эффективности. Пока размер пузыря пропорционален числу стадий конвейера, ценные вычислительные ресурсы GPU простаивают. Можно ли придумать ещё более умное расписание, чтобы минимизировать потерянное время?

Чередование стадий

Расписание 1F1B улучшило использование памяти, но почти не повлияло на размер пузыря простоя. Можно ли продвинуться дальше?

Оказывается, можно, если согласиться на несколько дополнительных операций коммуникации. Пора поговорить о чередующихся стадиях (interleaved stages)!

До сих пор мы наивно резали модель по глубине: например, слои 1–4 на первой GPU, слои 5–8 на второй. Но слои можно раскладывать и иначе, например нечётные слои (1, 3, 5, 7) на первую GPU, а чётные (2, 4, 6, 8) на вторую.

В целом это можно представить как «зацикленный конвейер»: микробатч во время прямого прохода через модель ходит по кругу от одной GPU к другой. В оригинале показан пример для модели, слои которой разложены по 4 GPU; числа по-прежнему — номера микробатчей, а первые и последние слои модели окрашены по-разному, чтобы было видно, как слои распределены по GPU.

Здесь нужна дополнительная коммуникация: модель проходит через каждую GPU несколько раз для того же вычисления, которое раньше занимало один проход. Зато каждый прямой и обратный проход делится на $v$, где $v$ — число стадий, или кусков модели, на одной GPU, потому что прямые и обратные проходы лучше чередуются:

$$\begin{aligned} &t_{pb} = \frac{(p-1) \cdot (t_f + t_b)}{v} \\ &r_{bubble} = \frac{1}{v} \cdot \frac{(p-1) \cdot (t_f + t_b)}{m \cdot (t_f + t_b)} = \frac{p-1}{v \cdot m} \end{aligned}$$

Значит, пузырь теперь можно уменьшать, добавляя микробатчи и чередующиеся стадии. Но учтите, что объём коммуникации тоже растёт в $v$ раз, так что это компромисс. На графике оригинала — несколько конфигураций PP с $p = 8$: частный случай $m = 1, v = 1$ — наивный конвейерный параллелизм, конфигурации с $v = 1$ — схемы AFAB или 1F1B, а случаи $v \neq 1$ — чередующиеся конфигурации.

Планирование здесь тоже усложняется: на каждой GPU в каждый момент нужно решать, что приоритетнее — ранние микробатчи, проходящие поздние слои, то есть как можно быстрее замыкать циклы прямого и обратного прохода (подход «сначала в глубину», depth-first, — как можно быстрее выводить батчи из модели), или поздние микробатчи, проходящие ранние слои (подход «сначала в ширину», breadth-first, — как можно плотнее заполнять конвейер). Этот выбор подробно разобран в статье «Breadth-First Pipeline Parallelism».

Теперь у вас есть всё, чтобы понять схему конвейерного параллелизма в Llama 3.1: она использует 1F1B с чередующимися стадиями и настраиваемым приоритетом между «сначала в глубину» и «сначала в ширину».

Однако на этом возможные расписания конвейера не кончаются: недавно предложены методы, уменьшающие пузырь почти до нуля! Их, например, использовали в реализации DeepSeek-V3/R1. Заинтриговали? Бросим последний быстрый взгляд на эти волшебные расписания, прежде чем покинуть мир конвейерного параллелизма!

Zero Bubble и DualPipe

Недавно предложены ещё более изощрённые способы уменьшить пузырь, приближающиеся к режиму «нулевого пузыря», например конвейерная реализация DeepSeek-V3/R1 под названием DualPipe. Секрет в том, чтобы разрезать операции на ещё более мелкие части и чередовать их максимально эффективно.

В техотчёте DeepSeek-V3 авторы пишут, что их схема достигла «почти нулевых накладных расходов на коммуникацию all-to-all».

Кратко посмотрим, как это может работать, на примере работы Sea AI Lab о «нулевом пузыре» (Zero Bubble) — предшественницы DualPipe. Основное наблюдение: обратный проход через умножение матриц на самом деле состоит из двух отдельных операций — обратной операции по входам (B) и обратной операции по весам (W).

Выход B, обратного прохода по входам, нужен для обратного прохода нижележащих слоёв, а обратный проход по весам, W, — нет: его обычно достаточно выполнить до шага оптимизатора. Это видно на схеме вычислительного графа из статьи Zero Bubble.

flowchart LR
F["F: прямой проход слоя"] --> B["B: градиент по входу"]
F --> W["W: градиент по весам"]
B --> PREV["B предыдущего слоя (ждёт этот результат)"]
W --> OPT["Шаг оптимизатора (ждёт только в конце)"]

Схема переводчика: B стоит на критическом пути, а W можно отложить и заполнить им пузыри

Значит, W можно гибко планировать в любой момент после соответствующего B той же стадии. Это позволяет стратегически размещать W, заполняя пузыри конвейера. В оригинале приведены рисунки из статьи Zero Bubble: сверху (рисунок 2 статьи) — классическое расписание 1F1B, где прямые и обратные проходы чередуются, но обратный проход остаётся крупным блоком; снизу (рисунок 3 статьи) — два вручную построенных расписания, где обратный проход разделён на более мелкие операции B и W. Нижнее расписание (ZB-H2) — пример (теоретического) расписания с нулевым пузырём, использующего такое мелкое разложение.

DualPipe от DeepSeek, представленный в техотчёте V3, расширяет этот подход с разложением на случай двух потоков, распространяющихся с обоих концов измерения PP; эти потоки чередуются, чтобы ещё сильнее сократить простои GPU. Расписание показано в оригинале — оно ещё сложнее предыдущих.

В общем случае полная оптимизация таких сложных расписаний требует тщательно измерить длительность разных мелких операций и решить задачу целочисленного линейного программирования (Integer Linear Programming, ILP), минимизирующую итоговое время пузыря. (Эвристики и алгоритмы такого планирования обсуждаются, например, в статье Zero Bubble.) Поэтому расписания Zero Bubble и DualPipe слишком сложны, чтобы приводить здесь фрагменты кода, но общее представление о концепциях у вас должно сложиться.

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

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

Экспертный параллелизм

Это последний метод параллелизма, который мы обсудим. Если вы не сталкивались с моделями «смесь экспертов» (Mixture of Experts, MoE), стоит сначала прочитать о них в гораздо более коротком посте, который мы публиковали раньше: он поможет лучше понять архитектуру MoE в целом.

Парадигма MoE недавно стала популярной и заметной благодаря моделям вроде GPT-4, Mixtral и DeepSeek-V3/R1. Основная идея: вместо одного полносвязного модуля в слое можно иметь несколько параллельных модулей и направлять токены через них, чтобы обрабатывать их по-разному. (В оригинале — иллюстрация слоя MoE из статьи Switch Transformers.)

Устройство слоёв MoE позволяет легко реализовать параллелизм по измерению экспертов — это называется экспертным параллелизмом (expert parallelism, EP). Поскольку полносвязные слои полностью независимы, можно просто положить полносвязный слой каждого эксперта на свой воркер. По сравнению с TP это гораздо легче: не нужно разрезать умножение матриц, достаточно направить скрытые состояния токена нужному эксперту.

На практике EP обычно используют вместе с другими формами параллелизма, например с параллелизмом по данным. Дело в том, что EP затрагивает только слои MoE и не шардирует входные токены (в отличие от контекстного параллелизма, который шардирует токены по длине последовательности). Значит, если использовать только EP, GPU делали бы избыточные вычисления во всех блоках, кроме MoE. Сочетая EP с DP, можно эффективно шардировать и экспертов, и входные батчи между GPU (в оригинале — упрощённая схема из обзора «A Survey on Mixture of Experts»).

flowchart LR
subgraph G1["GPU 1 (свой микробатч)"]
  A1["Внимание (копия)"] --> R1["Маршрутизатор"]
  E1["Эксперты 1–2"]
end
subgraph G2["GPU 2 (свой микробатч)"]
  A2["Внимание (копия)"] --> R2["Маршрутизатор"]
  E2["Эксперты 3–4"]
end
R1 -- "all-to-all: токены к своим экспертам" --> E2
R2 -- "all-to-all" --> E1
R1 --> E1
R2 --> E2

Схема переводчика: EP + DP — внимание продублировано, эксперты разложены по GPU, токены пересылаются через all-to-all

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

На практике есть несколько приёмов, чтобы EP работал эффективно, и они тесно связаны с устройством модели. Например, DeepSeek-V3 накладывает ограничение в маршрутизаторе: каждый токен отправляется не более чем на M узлов (у них 4), чтобы токены оставались ближе друг к другу и накладные расходы на коммуникацию были меньше. Экспертный параллелизм существует давно, но только сейчас снова набирает популярность вместе с архитектурой MoE.

Мы планируем скоро добавить более полный пример EP в Picotron/Nanotron — следите за новостями!

5D-параллелизм в двух словах

Поздравляем, читатель! Вы познакомились со всеми пятью стратегиями параллелизма для масштабирования обучения:

1. параллелизм по данным (DP) — по измерению батча;
2. тензорный параллелизм (TP) — по скрытому измерению;
3. параллелизм по последовательности и контекстный (SP/CP) — по измерению последовательности;
4. конвейерный параллелизм (PP) — по слоям модели;
5. экспертный параллелизм (EP) — по экспертам модели;

а также с тремя стратегиями ZeRO, которые сочетаются с параллелизмом по данным для экономии памяти:

1. ZeRO-1 — шардирование состояний оптимизатора между копиями DP;
2. ZeRO-2 — шардирование состояний оптимизатора и градиентов между копиями DP;
3. ZeRO-3 — шардирование состояний оптимизатора, градиентов и параметров между копиями DP.

Сейчас вам, наверное, интересно, как все эти стратегии параллелизма и ZeRO соотносятся и взаимодействуют друг с другом. Иначе говоря, какие из них можно эффективно сочетать, а какие лучше держать отдельно?

Посмотрим на сходства и взаимодействия. Начнём со сравнения конвейерного параллелизма и ZeRO-3: у них очень много общего, но есть и важные различия.

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

Однако между PP и ZeRO-3 есть несколько важных различий.

Для простоты мы говорим «слой», но реальная единица шардирования модели может быть разной — несколько слоёв, один слой или даже часть слоя, в зависимости от реализации.

	ZeRO-3	Конвейерный параллелизм
Каждое вычислительное устройство хранит…	лишь часть слоя	целый слой
Коммуникация передаёт…	веса	активации
Оркестровка	не зависит от модели	не зависит от модели
Трудности реализации	сложно организовать разбиение модели и коммуникацию	сложно организовать эффективные расписания PP
Масштабирование	предпочитает большие $mbs$ и $seq\_len$, чтобы спрятать коммуникацию	предпочитает большой $grad\_acc$, чтобы спрятать пузырь

Как видите, ZeRO-3 и PP решают одну задачу разными способами, и выбор между ними зависит от того, что вы предпочитаете пересылать — веса или активации. Их можно сочетать, но на практике это делают нечасто: для этого нужно заметно увеличить глобальный размер батча, чтобы окупить затраты на коммуникацию, и возникает компромисс между глобальным батчем, размером модели, пропускной способностью сети и эффективностью обучения. Если всё же сочетать, ZeRO-3 нужно настроить так, чтобы веса оставались в памяти на протяжении серии микробатчей PP, — так лишняя коммуникация будет минимальной.

С другой стороны, ZeRO-1 и ZeRO-2, касающиеся состояний оптимизатора и градиентов, легко сочетаются с конвейерным параллелизмом и дополняют его. Никаких новых трудностей такие сочетания не создают. Например, DeepSeek-V3 обучали с PP в сочетании с ZeRO-1.

Тензорный параллелизм (с параллелизмом по последовательности) естественно дополняет и конвейерный параллелизм, и ZeRO-3 и сочетается с ними, потому что опирается на дистрибутивность умножения матриц: веса и активации можно шардировать и вычислять независимо, а потом объединить.

Главная причина, по которой не хочется использовать один только TP, — у него на практике два ограничения (мы их уже обсуждали). Во-первых, его операции коммуникации лежат на критическом пути вычислений, поэтому после некоторого предела он плохо масштабируется: начинают доминировать накладные расходы на коммуникацию. Во-вторых, в отличие от ZeRO и PP, которые не зависят от модели, TP требует аккуратно обращаться с шардированием активаций — то по скрытому измерению (в области TP), то по измерению последовательности (в области SP), — поэтому его сложнее правильно реализовать и нужно знать устройство конкретной модели, чтобы шардирование везде было верным.

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

Контекстный и экспертный параллелизм тоже помогают шардировать активации и могут рассматриваться как дополнение к TP. CP справляется с длинными последовательностями, EP делает возможным распределённое обучение MoE, и их можно сочетать без особых проблем.

CP нацелен именно на обучение с очень длинными последовательностями: он шардирует активации по измерению последовательности между GPU. Большинство модулей, таких как MLP и LayerNorm, могут обрабатывать шардированные последовательности независимо, но блокам внимания нужна коммуникация, потому что каждому токену нужен доступ к ключам и значениям всей последовательности. Как мы видели в разделе о CP, это эффективно решается шаблонами Ring Attention, перекрывающими вычисления и коммуникацию. CP особенно ценен при масштабировании до экстремальных длин последовательности (128 тысяч токенов и больше), когда даже при полном перевычислении активаций требования внимания к памяти были бы неподъёмны для одной GPU.

Экспертный параллелизм нацелен именно на обучение моделей MoE: он шардирует специализированных «экспертов» между GPU и динамически направляет токены к нужным экспертам во время вычислений. Ключевые операции коммуникации в EP — «all-to-all»: они направляют токены к назначенным экспертам и собирают результаты обратно. Это добавляет накладные расходы на коммуникацию, но позволяет сильно нарастить ёмкость модели, потому что каждый токен при инференсе (и обучении) обрабатывается гораздо меньшей частью всех параметров. С точки зрения распределённого обучения и инференса разбиение экспертов между GPU становится важным, когда экспертов в модели много.

Например, у DeepSeek-V3 256 экспертов.

Заметка

Из-за этого сходства EP и DP в обработке входов некоторые реализации считают экспертный параллелизм частным случаем параллелизма по данным; ключевое отличие в том, что EP использует специализированную маршрутизацию по экспертам, а не пропускает входы на всех GPU через одинаковые копии модели.

Область действия

Кратко подытожим, на какие части модели сильнее всего влияют разные стратегии параллелизма:

• тензорный параллелизм (и параллелизм по последовательности) затрагивает вычисления во всей модели, шардируя и веса, и активации;
• контекстный параллелизм в основном затрагивает слои внимания, потому что именно там нужна коммуникация по последовательности, а остальные слои работают с шардированными последовательностями независимо;
• экспертный параллелизм в основном затрагивает слои MoE (которые заменяют обычные блоки MLP), не меняя слои внимания и другие компоненты;
• конвейерный параллелизм и ZeRO не привязаны к конкретным подмодулям, кроме того, что при конвейерном параллелизме модули и слои нужно сбалансировать (поэтому первые и последние слои часто обрабатывают особо из-за дополнительных слоёв эмбеддингов).

Тензорный + по последовательности	Контекстный	Экспертный
Шардирует веса и активации по скрытому измерению / последовательности	Шардирует активации по последовательности	Шардирует веса и активации специализированных экспертов
Коммуникация для умножений матриц (column/row linear)	Коммуникация для ключей и значений внимания	Коммуникация для маршрутизации токенов к экспертам
Нужна реализация под конкретную модель	Не зависит от модели, кроме внимания	Не зависит от модели, кроме слоёв MoE
Предпочитает быструю связь внутри узла	Предпочитает большую длину последовательности	Требует слоёв MoE

Сводим всё вместе

А что если собрать все рассмотренные приёмы на одной схеме? Да, мы готовы принять вызов!

На итоговой схеме оригинала изображены активации и модули одного слоя трансформера в варианте MoE, а также все направления параллелизма и операции коммуникации из предыдущих разделов. Там же дан полный обзор экономии памяти для каждой стратегии при разных длинах последовательности, с выборочным (сверху) и полным (снизу) перевычислением, — чтобы было видно, как каждая стратегия влияет на активации.

Завершим раздел общим обзором всех приёмов, их главных идей и основных узких мест:

Метод	На чём экономит память	Измерение параллелизма / шардирования	Недостаток
DP	Активации (уменьшает локальный батч)	Батч	Ограничен максимальным размером батча
PP	Параметры модели	Слои модели	Пузырь простоя и сложные расписания
TP+SP	Параметры модели и активации	Скрытое измерение / длина последовательности	Нужна связь с высокой пропускной способностью
CP	Активации	Длина последовательности	Добавляет коммуникацию в модулях внимания
EP	Параметры экспертов	Измерение экспертов	Нужны слои MoE, добавляет коммуникацию для маршрутизации
ZeRO-1	Состояния оптимизатора	Шардирование между копиями DP	Накладные расходы на передачу параметров
ZeRO-2	Состояния оптимизатора и градиенты	Шардирование между копиями DP	Накладные расходы на передачу параметров
ZeRO-3	Состояния оптимизатора, градиенты и параметры модели	Шардирование между копиями DP	Накладные расходы на передачу параметров

Очевидно, ни один из приёмов не серебряная пуля для волшебного масштабирования, и их почти всегда придётся как-то сочетать. Можно ли сформулировать несколько правил, чтобы найти хорошую отправную точку для выбора (и сочетания)? Об этом — следующий раздел.

Поиск лучшей конфигурации обучения

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

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

Шаг 1. Уместить шаг обучения в память

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

1. Много GPU 🤑 — когда GPU в избытке:

• для моделей меньше 10B параметров можно использовать одну технику параллелизма, например тензорный параллелизм или ZeRO-3/DP с полным перевычислением на 8 GPU;
• для моделей от 10B до 100B параметров, которым нужно больше 8 GPU, есть несколько вариантов:
  • тензорный параллелизм (TP=8) в сочетании с конвейерным;
  • тензорный параллелизм (TP=8) в сочетании с параллелизмом по данным (ZeRO-3);
  • только ZeRO-3 (то есть чистый параллелизм по данным);
• от 512 GPU чистый параллелизм по данным/ZeRO-3 становится неэффективным из-за затрат на коммуникацию — тогда лучше сочетать DP с тензорным или конвейерным параллелизмом;
• от 1024 GPU рекомендуемая схема — тензорный параллелизм (TP=8) с параллелизмом по данным (ZeRO-2) и конвейерным параллелизмом.

Пока мы сосредоточены на том, чтобы уместить один экземпляр модели. Хотя для этого мы можем использовать DP вместе с ZeRO, здесь нас интересует только экономия памяти под параметры модели, которую даёт ZeRO-3.

Особые случаи:

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

2. Мало GPU 😭 — когда ресурсов GPU не хватает:

• можно включить полное перевычисление активаций, обменяв часть вычислений на память (и обучать чуть медленнее);
• можно увеличить накопление градиентов, чтобы обрабатывать большие батчи при ограниченной памяти.

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

4× RTX 4080 по 32 ГБ через PCIe — это случай «мало GPU». Для полного обучения моделей до ~1–2B обычно хватает DDP или ZeRO-2/FSDP с перевычислением и накоплением градиентов; для LoRA больших моделей — FSDP или просто размещение модели по картам. Тензорный параллелизм без NVLink редко окупается.

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

Шаг 2. Достичь целевого глобального размера батча

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

Чтобы увеличить текущий глобальный размер батча:

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

Чтобы уменьшить текущий глобальный размер батча:

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

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

Шаг 3. Оптимизировать пропускную способность обучения

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

• увеличивать тензорный параллелизм (используя быструю связь внутри узла), пока его степень не приблизится к размеру узла, чтобы уменьшить другие формы параллелизма;
• увеличивать параллелизм по данным с ZeRO-3, сохраняя целевой размер батча;
• когда коммуникация параллелизма по данным становится узким местом, переходить к конвейерному параллелизму;
• пробовать увеличивать разные виды параллелизма по одному;
• экспериментировать с размером микробатча ($mbs$), добиваясь оптимального баланса между максимальным глобальным батчем, размером модели, вычислениями и коммуникацией.

Тысячи конфигураций на бенчмарке

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

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

Мы сами прогнали бенчмарки на нескольких тысячах распределённых конфигураций, охватив все обсуждаемые здесь размеры моделей и очень много конфигураций кластера (а именно 1–64 узла по 8×H100), чтобы получить результаты, приведённые в этой книге.

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

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

Все следующие бенчмарки проведены с длиной последовательности 4096 и глобальным размером батча 1 млн токенов. Мы собрали лучшие конфигурации для каждого размера модели и кластера и нарисовали тепловые карты. В оригинале — тепловая карта оптимальных конфигураций обучения для разных размеров моделей и числа вычислительных узлов (по 8 GPU в узле). Для каждой комбинации указаны параллелизм по данным (DP), тензорный (TP), конвейерный (PP), число шагов накопления градиентов (GAS), размер микробатча (MBS) и стадия ZeRO. Интенсивность цвета показывает утилизацию FLOPs модели (MFU): чем ярче, тем эффективнее.

Из этой общей картины можно сделать несколько важных выводов.

• Во-первых, с ростом числа узлов (больше параллелизма) эффективность снижается. Особенно сильно это заметно на маленьких моделях, у которых меньше отношение вычислений к размеру модели. Обычно маленький размер модели можно компенсировать большим батчем, но нас ограничивает глобальный батч в 1 млн токенов.
• Во-вторых, большие модели ставят другую задачу. С ростом модели требования к памяти растут очень сильно. При малом числе узлов возникают два сценария: модель либо вообще не помещается, либо помещается, но работает неэффективно, потому что почти упирается в предел памяти GPU (см., например, обучение модели на 80 млрд параметров на 4 узлах).
• Наконец, бенчмарки показывают, как сильно производительность зависит от качества реализации. Когда мы впервые реализовали обе стратегии параллелизма, тензорный параллелизм обгонял конвейерный. После оптимизации кода PP быстрее стал он. Сейчас мы улучшаем перекрытие коммуникации в реализации TP и ожидаем, что он вернёт лидерство.

Уроки бенчмаркинга

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

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

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

Чтобы провести все эксперименты за конечное время, понадобилась дополнительная инженерная работа, и мы потратили много времени на такие вещи, как:

• сокращение времени перезапуска кластера и оптимизация простоев;
• анализ подробных отладочных логов NCCL;
• понимание шаблонов использования памяти и поведения аллокатора памяти CUDA;
• улучшение производительности конвейерного параллелизма на нескольких узлах.

Эти трудности дали нам ценные уроки о сложности инфраструктуры распределённого обучения. То, что в теории выглядит просто, на практике требует внимания ко множеству движущихся частей.

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

---

На этом наше очень глубокое погружение в методы распределения 5D-параллелизма заканчивается.

Если отступить на шаг, наше обсуждение до сих пор часто опиралось на важное допущение: что вычисления и коммуникацию можно эффективно перекрывать на GPU без ущерба для пропускной способности вычислений. В реальности всё тоньше. При использовании распространённых примитивов коммуникации вроде NCCL send/recv возникает скрытая конкуренция между ресурсами вычислений и коммуникации: ядра коммуникации обычно используют те же потоковые мультипроцессоры GPU (о них — в следующем разделе), что и вычисления. Поэтому при перекрытии коммуникации с вычислениями пропускная способность падает. Чтобы по-настоящему оптимизировать распределённое обучение, нужно глубже погрузиться в архитектуру самой GPU.

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

Пора выключить свет и включить режим CUDA! Продолжение — в части 4.

Словарик

EN	RU
context parallelism (CP)	контекстный параллелизм
Ring Attention, Zig-Zag Attention	кольцевое внимание, зигзагообразное внимание
causal mask	причинная маска
all-to-all	обмен «все со всеми»
pipeline parallelism (PP), stage	конвейерный параллелизм, стадия конвейера
bubble	пузырь (простой) конвейера
AFAB / 1F1B	«все прямые, все обратные» / «один прямой, один обратный»
interleaved stages	чередующиеся стадии
depth-first / breadth-first	сначала в глубину / сначала в ширину
Zero Bubble, DualPipe	расписания с (почти) нулевым пузырём
expert parallelism (EP), router	экспертный параллелизм, маршрутизатор
GPU-rich / GPU-poor	«богатые GPU» / «бедные GPU»
MFU	утилизация FLOPs модели