How To Scale Your Model, главы 5–6: как параллелить обучение трансформера

Здесь мы обсудим четыре главные схемы параллелизма, используемые при обучении LLM: параллелизм данных, полностью шардированный параллелизм данных (FSDP), тензорный параллелизм и конвейерный параллелизм. Для каждой мы посчитаем, в какой момент связь становится узким местом. — подзаголовок главы 5.

О переводе

Перевод главы 5 книги How To Scale Your Model (лицензия MIT, © 2025 Google LLC): весь основной текст, включая алгоритмы прямого и обратного прохода для DP, FSDP и TP (алгоритм для FSDP + TP из 12 шагов пересказан одной фразой), итоговые таблицы; условия задач перечислены в конце, приложение с выводом связи в обратном проходе не переведено. Глава 6 («Обучение LLaMA 3 на TPU») дана в кратком изложении в конце страницы. Предыдущая глава — «Вся математика трансформера», следующая — «Всё об инференсе трансформера». Те же расчёты для GPU — в главе 12.

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

Это теория к уроку «Распределённое обучение» (DDP, FSDP/ZeRO, тензорный и конвейерный параллелизм) и к практике «Четыре видеокарты: DDP и FSDP на практике». Главный результат главы — простое правило «параллелизм данных не упирается в связь, пока батч на устройство больше $C / W$» — в уроке 5.7.4 пересчитан для PCIe. Числа в книге — для TPU v5p ($C / W_\text{ici} \approx 2550$).

Что мы понимаем под масштабированием?

Цель «масштабирования модели» — иметь возможность увеличивать число чипов для обучения или инференса и получать при этом пропорциональный, линейный рост пропускной способности (это мы называем сильным масштабированием, strong scaling). Производительность одного чипа зависит от компромисса между пропускной способностью памяти и FLOPs, а производительность на уровне кластера — от того, удаётся ли скрыть связь между чипами, совмещая её с полезными FLOPs. Это нетривиально: с ростом числа чипов растёт нагрузка на связь, а объём вычислений на устройство, которым её можно скрыть, уменьшается. Как мы видели в главе 3, умножения шардированных матриц часто требуют дорогих AllGather или ReduceScatter, которые могут мешать TPU делать полезную работу. Цель этой главы — выяснить, когда они становятся слишком дорогими.

В этой главе мы обсудим четыре распространённые схемы параллелизма: (чистый) параллелизм данных, полностью шардированный параллелизм данных (FSDP / шардирование ZeRO), тензорный параллелизм (его ещё называют параллелизмом модели) и (кратко) конвейерный параллелизм. Для каждой покажем, какую стоимость связи она требует и с какого момента эта стоимость становится узким местом для вычислений¹. В этой главе можно думать только о стоимости связи между чипами: при достаточно большом батче на чип передача данных из HBM в MXU уже совмещена с вычислениями.

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

Обозначение	Смысл (параметры модели)
D	dmodel (скрытая размерность, размерность остаточного потока)
F	dff (размерность прямого слоя)
B	измерение батча (число токенов в батче; общее, а не на устройство)
T	длина последовательности
L	число слоёв модели

Обозначение	Смысл (характеристики железа)
C	FLOPs/s на чип
W	пропускная способность сети (двусторонняя, часто с индексом, например $W_{\text{ici}}$ или $W_{\text{dcn}}$)
X	число чипов вдоль оси сетки X
Y	число чипов вдоль другой оси сетки, Y
Z	число чипов вдоль третьей оси сетки, Z

Для простоты будем приближать трансформер стопкой блоков MLP: как мы видели в главе 4, у больших моделей внимание — сравнительно малая доля FLOPs. Гейтовое умножение тоже опустим, и у каждого слоя останется простая структура.

Рисунок в оригинале (transformer-layer.png): упрощённый слой трансформера. Каждый блок FFW — стопка из двух матриц: W_in: bf16[D, F] (повышающая проекция) и W_out: bf16[F, D] (понижающая проекция), вход — In: bf16[B, D].

Полный алгоритм нашего маленького трансформера без параллелизма

Прямой проход: нужно вычислить Loss[B]

1. Tmp[B, F] = In[B, D] *_D W_in[D, F]
2. Out[B, D] = Tmp[B, F] *_F W_out[F, D]
3. Loss[B] = …

Обратный проход: нужно вычислить dW_out[F, D], dW_in[D, F]

1. dOut[B, D] = …
2. dW_out[F, D] = Tmp[B, F] *_B dOut[B, D]
3. dTmp[B, F] = dOut[B, D] *_D W_out[F, D]
4. dW_in[D, F] = In[B, D] *_B dTmp[B, F]
5. dIn[B, D] = dTmp[B, F] *_F W_in[D, F] (нужно для предыдущих слоёв)

Он приведён для сравнения с алгоритмами, где добавлена связь.

Вот 4 схемы параллелизма, которые мы обсудим. Каждую можно считать однозначно заданной шардированием In, W_in, W_out и Out на схеме выше.

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

$$\text{In}[B_X, D] \cdot_D W_\text{in}[D, F] \cdot_F W_\text{out}[F, D] \rightarrow \text{Out}[B_X, D]$$

2. Полностью шардированный параллелизм данных (FSDP, или ZeRO-3): активации шардированы по батчу (как в чистом параллелизме данных), параметры шардированы по той же оси сетки и собираются AllGather «точно к сроку» перед использованием в прямом проходе. Состояние оптимизатора тоже шардировано по батчу. Сокращает дублирование памяти.

$$\text{In}[B_X, D] \cdot_D W_\text{in}[D_X, F] \cdot_F W_\text{out}[F, D_X] \rightarrow \text{Out}[B_X, D]$$

3. Тензорный параллелизм (его ещё называют шардированием Megatron или параллелизмом модели): активации шардированы по D ($d_\text{model}$), параметры — по F ($d_{ff}$). Активации собираются AllGather перед каждым блоком и разбрасываются ReduceScatter после. Совместим с FSDP.

$$\text{In}[B, D_Y] \cdot_D W_\text{in}[D, F_Y] \cdot_F W_\text{out}[F_Y, D] \rightarrow \text{Out}[B, D_Y]$$

4. Конвейерный параллелизм: веса шардированы по измерению слоёв, активации разбиты на микробатчи и «прокатываются» вдоль измерения слоёв. Связь между стадиями конвейера минимальна (активации просто передаются на один переход). Вольно обращаясь с нотацией:

$$\text{In}[L_Z, B, D][i] \cdot_D W_\text{in}[L_Z, D, F][i] \cdot_F W_\text{out}[L_Z, F, D][i] \rightarrow \text{Out}[L_Z, B, D][i]$$

flowchart LR
DP["DP: батч шардирован, веса копиями"] --> FSDP["FSDP: + веса и оптимизатор шардированы"]
FSDP --> MIX["FSDP + TP: веса шардированы по двум осям"]
TP["TP: активации по D, веса по F"] --> MIX
MIX --> PP["+ PP: слои по стадиям"]

Как схемы параллелизма надстраиваются друг над другом (схема переводчика)

Параллелизм данных

Запись: $\text{In}[B_X, D] \cdot_D W_\text{in}[D, F] \cdot_F W_\text{out}[F, D] \rightarrow \text{Out}[B_X, D]$

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

Рисунок в оригинале (data-parallelism.png): чистый параллелизм данных (прямой проход). Активации (слева) полностью шардированы по батчу, а веса полностью реплицированы, так что у каждого TPU одинаковая копия весов. Значит, общая память весов растёт в N раз, но в прямом проходе связь не нужна.

Полный алгоритм прямого и обратного прохода (для краткости dL/dOut пишем как dOut)

Алгоритм чистого параллелизма данных.

Прямой проход: нужно вычислить Loss[B_X]

1. Tmp[B_X, F] = In[B_X, D] *_D W_in[D, F]
2. Out[B_X, D] = Tmp[B_X, F] *_F W_out[F, D]
3. Loss[B_X] = …

Обратный проход: нужно вычислить dW_out[F, D], dW_in[D, F]

1. dOut[B_X, D] = …
2. dW_out[F, D] {U_X} = Tmp[B_X, F] *_B dOut[B_X, D]
3. dW_out[F, D] = AllReduce(dW_out[F, D] {U_X}) (не на критическом пути, можно асинхронно)
4. dTmp[B_X, F] = dOut[B_X, D] *_D W_out[F, D]
5. dW_in[D, F] {U_X} = In[B_X, D] *_B dTmp[B_X, F]
6. dW_in[D, F] = AllReduce(dW_in[D, F] {U_X}) (не на критическом пути, можно асинхронно)
7. dIn[B_X, D] = dTmp[B_X, F] *_F W_in[D, F] (нужно для предыдущих слоёв)

Детали функции потерь опущены, и $\text{Tmp} = W_\text{in} \cdot \text{In}$. Заметьте: хотя итоговые потери — среднее AllReduce(Loss[B_X]), AllReduce нужен только в обратном проходе — при усреднении градиентов весов.

Заметьте: в прямом проходе связи нет — вся она в обратном проходе! У обратного прохода есть ещё отличное свойство: AllReduce не лежат на «критическом пути», то есть каждый AllReduce можно выполнить, когда удобно, и он не блокирует следующие операции. Общая стоимость связи всё равно может стать узким местом, если превысит общую стоимость вычислений, но с точки зрения реализации это гораздо снисходительнее. Мы увидим, что у тензорного параллелизма этого свойства нет.

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

Почему бы так не делать? Чистый параллелизм данных никак не снижает давление на память со стороны параметров модели и состояния оптимизатора, поэтому для интересных моделей большого масштаба, где параметры + состояние оптимизатора не помещаются в один TPU, он редко полезен. Для масштаба: если обучать с параметрами в bf16 и состоянием оптимизатора Adam в fp32², самая большая модель, которая поместится, имеет $\text{память TPU} / 10$ параметров, например на чипе TPUv5p с 96 ГБ HBM при чистом параллелизме данных — около 9 млрд параметров.

Вывод

Самая большая модель, которую можно обучать с Adam и чистым параллелизмом данных, имеет $\text{число параметров} = \text{HBM на устройство} / 10$. Для TPU v5p это примерно 9 млрд параметров³.

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

Когда мы упираемся в связь? Как видно выше, на слой приходится два AllReduce размером $2DF$ каждый (для весов в bf16). Когда параллелизм данных упирается в связь?

Как в таблице выше, пусть $C$ — FLOPs на чип, $W_{\text{ici}}$ — двусторонняя пропускная способность сети, а $X$ — число шардов, на которые разбит батч⁴. Посчитаем время на нужные умножения, $T_\text{math}$, и время связи, $T_\text{comms}$. В прямом проходе эта схема связи не требует, так что считаем только обратный проход.

Время связи. Из прошлой главы мы знаем, что время AllReduce на одномерной сетке зависит только от общего числа байт массива и пропускной способности ICI $W_\text{ici}$, а именно — $2 \cdot \text{всего байт} / W_\text{ici}$. AllReduce нужен и для $W_\text{in}$, и для $W_\text{out}$, то есть 2 AllReduce на слой. Каждый — для матрицы весов из $DF$ параметров, или $2DF$ байт. Итого время AllReduce на один слой —

$$T_\text{comms} = \frac{2 \cdot 2 \cdot 2 \cdot D \cdot F}{W_\text{ici}}$$

Время умножений. В каждом слое два умножения в прямом проходе или четыре в обратном, каждое требует $2(B/X)DF$ FLOPs. Значит, для одного слоя в обратном проходе

$$T_\text{math} = \frac{2 \cdot 2 \cdot 2 \cdot B \cdot D \cdot F}{X \cdot C}$$

Поскольку мы их совмещаем, общее время на слой — максимум из двух величин:

$$T \approx \max\left(\frac{8 \cdot B \cdot D \cdot F}{X \cdot C}, \frac{8 \cdot D \cdot F}{W_\text{ici}}\right) = 8 \cdot D \cdot F \cdot \max\left(\frac{B}{X \cdot C}, \frac{1}{W_\text{ici}}\right)$$

Мы упираемся в вычисления, когда $T_\text{math}/T_\text{comms} > 1$, то есть когда

$$\frac{B}{X} > \frac{C}{W_\text{ici}}$$

Вывод: чтобы при параллелизме данных упираться в вычисления, батч на устройство $B / X$ должен превышать операционную интенсивность ICI, $C / W_\text{ici}$. В конечном счёте это следствие того, что время вычислений растёт с батчем на устройство, а время связи от него не зависит (передаются веса модели). Заметьте сходство условия $B/X > C/W_\text{ici}$ с правилом $B > 240$ для одного устройства: и там правило возникло потому, что время вычислений росло с размером батча, а объём передаваемых данных (при $B \ll F, D$) от него не зависел.

Подставим реальные числа, чтобы почувствовать масштаб. Для TPUv5p C=4.6e14 и W=2 * 9e10 при одномерном параллелизме данных по ICI, так что батч на чип должен быть не меньше 2550, чтобы не упираться в связь. Параллелизм данных можно вести по нескольким осям; если отдать все три оси пода TPUv5p чистому параллелизму данных, пропускная способность $W_\text{ici}$ утроится, и можно опуститься до батча 850 на TPU, или 7,6 млн токенов на батч на под (из 8960 чипов)! Значит, упереться в связь при чистом параллелизме данных довольно трудно!

Замечание [контекстный параллелизм]

Во всей главе $B$ — общий размер батча в токенах. Но батч, очевидно, состоит из многих разных последовательностей — как это работает? С точки зрения MLP токен — это токен! Неважно, принадлежат ли токены одной последовательности или двум разным. Так что мы более-менее свободно можем вести параллелизм данных и по батчу, и по измерению последовательности: это называют контекстным (context) или последовательностным (sequence) параллелизмом, но можно считать его просто ещё одним видом параллелизма данных. С вниманием сложнее, чем с MLP, ведь там есть вычисления по всей последовательности, но это решается сбором KV или Q во время внимания и аккуратным совмещением FLOPs и связи (обычно с помощью так называемого «кольцевого внимания», ring attention). В этой главе мы просто игнорируем измерение последовательности и предполагаем какую-то долю параллелизма по батчу или последовательности.

О нескольких осях сетки. Коротко отметим, как несколько осей влияют на доступную пропускную способность. Когда для данной стратегии параллелизма используется несколько осей сетки, пропускная способность растёт.

• Определение: $M_X$ ($M_Y$, $M_Z$ и т. д.) — число аппаратных осей сетки, которые охватывает данная стратегия параллелизма.
• Эффект (в режиме, ограниченном пропускной способностью): $M$ осей дают примерно в $M$ раз большую суммарную пропускную способность каналов, так что время коллективной операции пропорционально $1/M_X$.

🖥 На вашей машине (примечание переводчика)

DDP в PyTorch — это ровно чистый параллелизм данных из этого раздела, и формула $B/X > C/W$ применима напрямую. Для 4×RTX 4080 через PCIe $C/W \approx 97{,}5 \cdot 10^{12} / 10^{10} \approx 9750$ токенов на карту при пиковых FLOPs и около 3900 при реалистичных 40% MFU — подробно в уроке «Четыре видеокарты: DDP и FSDP на практике». Предел по памяти: $32 / 10 \approx 3$ млрд параметров на карту при Adam в смешанной точности без учёта активаций (в уроке «Смешанная точность» счёт 16 байт на параметр, т. е. ~2 млрд, потому что там добавлены мастер-веса в FP32).

Полностью шардированный параллелизм данных (FSDP)

Запись: $\text{In}[B_X, D] \cdot_D W_\text{in}[D_X, F] \cdot_F W_\text{out}[F, D_X] \rightarrow \text{Out}[B_X, D]$

Полностью шардированный параллелизм данных (его часто называют FSDP или шардированием ZeRO) делит состояние оптимизатора и веса модели между шардами параллелизма данных и эффективно собирает и разбрасывает их по мере надобности. По сравнению с чистым параллелизмом данных FSDP резко сокращает расход памяти на устройство и экономит FLOPs обратного прохода при совсем небольших накладных расходах.

Рисунок в оригинале (fsdp.png): FSDP шардирует сворачиваемое измерение W_in и выходное измерение W_out по измерению данных. Это экономит память, но (см. главу 3) требует собрать веса W перед умножением. Заметьте: активации (слева) не шардированы по сворачиваемому измерению — именно это и вынуждает собирать. Состояние оптимизатора весов тоже шардировано по сворачиваемому измерению.

Вспомните (из главы 3): AllReduce раскладывается на AllGather и ReduceScatter. Значит, вместо полного AllReduce градиентов, как в обычном параллелизме данных, можно шардировать веса и состояние оптимизатора между чипами, собирать их AllGather на каждом слое в прямом проходе и разбрасывать ReduceScatter по весам в обратном — без дополнительных затрат.

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

Прямой проход: нужно вычислить Loss[B_X]

1. W_in[D, F] = AllGather(W_in[D_X, F]) (не на критическом пути, можно во время предыдущего слоя)
2. Tmp[B_X, F] = In[B_X, D] *_D W_in[D, F] (теперь W_in[D, F] можно выбросить)
3. W_out[F, D] = AllGather(W_out[F, D_X]) (не на критическом пути, можно во время предыдущего слоя)
4. Out[B_X, D] = Tmp[B_X, F] *_F W_out[F, D]
5. Loss[B_X] = …

Обратный проход: нужно вычислить dW_out[F, D_X], dW_in[D_X, F]

1. dOut[B_X, D] = …
2. dW_out[F, D] {U_X} = Tmp[B_X, F] *_B dOut[B_X, D]
3. dW_out[F, D_X] = ReduceScatter(dW_out[F, D] {U_X}) (не на критическом пути, можно асинхронно)
4. W_out[F, D] = AllGather(W_out[F, D_X]) (можно заранее)
5. dTmp[B_X, F] = dOut[B_X, D] *_D W_out[F, D] (здесь W_out[F, D] можно выбросить)
6. dW_in[D,F] {U_X} = In[B_X, D] *_B dTmp[B_X, F]
7. dW_in[D_X, F] = ReduceScatter(dW_in[D, F] {U_X}) (не на критическом пути, можно асинхронно)
8. W_in[D, F] = AllGather(W_in[D_X, F]) (можно заранее)
9. dIn[B_X, D] = dTmp[B_X, F] *_F W_in[D, F] (нужно для предыдущих слоёв; здесь W_in[D, F] можно выбросить)

Это называют ещё «шардированием ZeRO» — от «Zero Redundancy Optimizer» («оптимизатор без избыточности»), поскольку мы не выполняем лишних вычислений и не храним лишнего состояния. ZeRO-1, ZeRO-2 и ZeRO-3 обозначают шардирование таким способом соответственно состояния оптимизатора, градиентов и весов. Стоимость связи у всех одинакова⁵, так что практически всегда можно делать шардирование ZeRO-3, которое делит между устройствами параметры, градиенты и состояние оптимизатора.

Зачем это нужно? Обычный параллелизм данных выполняет много дублирующейся работы. Каждый TPU делает AllReduce полного градиента, затем обновляет полное состояние оптимизатора (одинаковая работа на всех TPU), затем обновляет параметры (снова полное дублирование). При шардировании ZeRO (шардирование градиентов и состояния оптимизатора) вместо AllReduce можно сделать ReduceScatter градиентов, обновить только свой шард состояния оптимизатора, обновить шард параметров, а затем собрать параметры AllGather, когда они понадобятся в прямом проходе.

Когда мы упираемся в связь? Соотношение FLOPs и связи в точности то же, что у чистого параллелизма данных: каждый AllReduce обратного прохода превратился в AllGather + ReduceScatter. Вспомните: AllReduce реализуется как AllGather и ReduceScatter, каждый — с половиной стоимости. Здесь моделируем прямой проход, поскольку отношение FLOPs к связи у него то же, что у обратного:

$$\begin{aligned} T_\text{math} &= \frac{2 \cdot 2 \cdot B \cdot D \cdot F}{X \cdot C} \\ T_\text{comms} &= \frac{2 \cdot 2 \cdot D \cdot F}{W_\text{ici}} \\ T &\approx 4 \cdot D \cdot F \cdot \max\left(\frac{B}{X \cdot C}, \frac{1}{W_\text{ici}}\right) \end{aligned}$$

Поэтому, как и при чистом параллелизме данных, мы упираемся в вычисления, когда $B / X > C / W_\text{ici}$, то есть когда батч на устройство $B/X$ превышает «операционную интенсивность ICI» $C/W_\text{ici}$ (4.59e14 / 1.8e11 = 2550 у v5p). Для нас это отличная новость: если батч на устройство достаточно велик, чтобы чистый параллелизм данных упирался в вычисления, можно — не боясь выйти из режима упора в вычисления — просто перейти на FSDP и сэкономить огромный объём памяти на параметрах и состоянии оптимизатора! Хотя в прямой проход пришлось добавить связь, эта стоимость несущественна: она просто перекрывается FLOPs прямого прохода.

Вывод

И FSDP, и чистый параллелизм данных на TPUv5 упираются в пропускную способность, когда батч на устройство меньше $2550 / M_X$, где $M_X$ — число осей сетки.

Например, DeepSeek-V2 (одна из немногих сильных недавних моделей, раскрывших размер батча при обучении) использовала батч около 40 млн токенов. Это позволило бы масштабироваться примерно до 47 000 чипов, или около 5 подов TPUv5, прежде чем упереться в пропускную способность.

Для LLaMA-3 70B, на обучение которой ушло примерно 6.3e24 (15e12 * 70e9 * 6) FLOPs, можно разделить батч в 16 млн токенов примерно на 16e6 / (2550 / 3) = 18,823 чипа (примерно 2 пода по 8960 чипов), каждый с 4.59e14 FLOPs при 50% использования пиковых FLOPs (это часто называют MFU), и обучить её примерно за 17 дней. Неплохо! Но посмотрим, как сделать лучше.

Замечание о критическом размере батча

Немного вопреки интуиции, при фиксированном числе чипов мы тем сильнее упираемся в связь, чем меньше общий размер батча. Параллелизм данных и FSDP позволяют масштабироваться на сколь угодно много чипов, пока можно увеличивать батч! Но на практике с ростом батча отдача от обучения обычно убывает: градиенты становятся почти бесшумными. Иногда возникает и неустойчивость обучения. Поэтому поиск оптимальной схемы шардирования в «режиме неограниченных вычислений» часто начинается с фиксированного размера батча, заданного законами масштабирования, и известного (большого) числа чипов, и цель — найти разбиение, которое позволит уместить этот маленький батч на столь многих чипах.

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

Запись: $\text{In}[B, D_Y] \cdot_D W_\text{in}[D, F_Y] \cdot_F W_\text{out}[F_Y, D] \rightarrow \text{Out}[B, D_Y]$ (мы используем $Y$, чтобы потом сочетать с FSDP)

При AllReduce полностью шардированного параллелизма данных мы перемещаем между чипами веса. Можно также шардировать измерение прямого слоя модели и перемещать во время слоя активации — это называют «одномерным параллелизмом модели» или шардированием Megatron. Это позволяет уменьшить эффективный размер батча на под.

Рисунок в оригинале (model-parallelism.png): пример базового тензорного параллелизма. Активации шардированы только по Y (в отличие от FSDP, где шардируем по X), поэтому по X они реплицированы. В нашей стандартной записи: A[B, D_Y] * B[D, F_Y] → C[B, F_Y]. Поскольку шардировано только одно из сворачиваемых измерений, обычно перед умножением активации A собирают AllGather.

Как отмечено, In[B, D_Y] *_D W_in[D, F_Y] *_F W_out[F_Y, D] → Out[B, D_Y] означает, что перед первым умножением нужно собрать активации. Это дешевле шардирования ZeRO, когда активации меньше весов. Обычно так бывает лишь при добавлении какой-то доли шардирования ZeRO (оно уменьшает размер сбора). Это одна из причин, почему ZeRO и тензорный параллелизм обычно смешивают.

Алгоритм тензорного параллелизма

Прямой проход: нужно вычислить Loss[B]

1. In[B, D] = AllGather(In[B, D_Y]) (на критическом пути)
2. Tmp[B, F_Y] = In[B, D] *_D W_in[D, F_Y] (сворачиваемое не шардировано — связи нет)
3. Out[B, D] {U_Y} = Tmp[B, F_Y] *_F W_out[F_Y, D]
4. Out[B, D_Y] = ReduceScatter(Out[B, D] {U_Y}) (на критическом пути)
5. Loss[B] = …

Обратный проход: нужно вычислить dW_out[F_Y, D], dW_in[D, F_Y]

1. dOut[B, D_Y] = …
2. dOut[B, D] = AllGather(dOut[B, D_Y]) (на критическом пути)
3. dW_out[F_Y, D] = Tmp[B, F_Y] *_B dOut[B, D]
4. dTmp[B, F_Y] = dOut[B, D] *_D W_out[F_Y, D] (здесь dOut[B, D] можно выбросить)
5. In[B, D] = AllGather(In[B, D_Y]) (можно пропустить, переиспользовав шаг 1 прямого прохода)
6. dW_in[D, F_Y] = In[B, D] *_B dTmp[B, F_Y]
7. dIn[B, D] {U_Y} = dTmp[B, F_Y] *_F W_in[D, F_Y] (нужно для предыдущих слоёв)
8. dIn[B, D_Y] = ReduceScatter(dIn[B, D] {U_Y}) (на критическом пути)

Приятная особенность тензорного параллелизма — он хорошо сочетается с двумя матрицами прямого прохода трансформера. Наивно мы делали бы AllReduce после каждой из двух матриц. Но здесь мы сначала делаем In[B, D_Y] * W_in[D, F_Y] → Tmp[B, F_Y], а затем Tmp[B, F_Y] * W_out[F_Y, D] → Out[B, D_Y]. То есть в начале делаем AllGather In, а в конце — ReduceScatter Out, вместо AllReduce.

Сколько это стоит? Смоделируем только прямой проход — обратный просто транспонирует каждую операцию. В одномерном тензорном параллелизме мы собираем активации AllGather перед первым умножением и разбрасываем ReduceScatter после второго, отправляя по два байта (bf16). Найдём, когда связь становится узким местом.

$$\begin{aligned} T_\text{math} &= \frac{4 \cdot B \cdot D \cdot F}{Y \cdot C} \\ T_\text{comms} &= \frac{2 \cdot 2 \cdot (B \cdot D)}{W_\text{ici}} \\ T &\approx \max \left(\frac{4 \cdot B \cdot D \cdot F}{Y \cdot C}, \frac{2 \cdot 2 \cdot (B \cdot D)}{W_\text{ici}}\right) \end{aligned}$$

Мы хотим, чтобы стоимость вычислений была больше стоимости связи:

$$\frac{4 \cdot B \cdot D \cdot F}{Y \cdot C} > \frac{2 \cdot 2 \cdot (B \cdot D)}{W_\text{ici}} \;\Leftrightarrow\; \frac{F}{Y \cdot C} > \frac{1}{W_\text{ici}} \;\Leftrightarrow\; F > Y \cdot \frac{C}{W_\text{ici}}$$

Например, для TPUv5p $C / W_{ici} = 2550$ в bf16, так что тензорный параллелизм возможен только до $Y < F / 2550$. При нескольких осях ICI $T_\text{comms}$ уменьшается в $M_Y$ раз, и получаем $Y < M_Y \cdot F / 2550$.

Вывод

Тензорный параллелизм упирается в связь, когда $Y > M_Y \cdot F / 2550$. Для большинства моделей это где-то между 8- и 16-путевым тензорным параллелизмом.

Заметьте: это не зависит от точности вычислений, поскольку, например, для int8 на TPUv5p $C_\text{int8} / W_{ici}$ равно 5100, а не 2550, но и объём связи вдвое меньше, так что двойки сокращаются.

Несколько примеров:

• На TPUv5p с LLaMA 3-70B ($D = 8192$, $F \approx 30\,000$) можно спокойно делать 8-путевой тензорный параллелизм, но 16-путевой упрётся в связь. Для 8-путевого шардирования модели нужно F не меньше 20 тысяч.
• У Gemma 7B $F \approx 50$ тысяч, так что в связь упрёмся при 19-путевом тензорном параллелизме. Значит, 16-путевой, скорее всего, ещё будет работать хорошо.

Сочетание FSDP и тензорного параллелизма

Запись: $\text{In}[B_X, D_Y] \cdot_D W_\text{in}[D_X, F_Y] \cdot_F W_\text{out}[F_Y, D_X] \rightarrow \text{Out}[B_X, D_Y]$

FSDP и тензорный параллелизм хороши тем, что их можно сочетать. Шардируя W_in и W_out по обеим осям, мы экономим и память, и вычисления. Поскольку B шардировано по X, уменьшаются AllGather параллелизма модели, а поскольку F шардировано по Y, уменьшаются накладные расходы связи FSDP. Значит, сочетание позволяет опуститься к ещё меньшему эффективному размеру батча, чем выше.

Рисунок в оригинале (mixed-fsdp-model-parallelism.png): сочетание FSDP и тензорного параллелизма. В отличие от других случаев, параметры модели нигде не дублируются. (Полный алгоритм из 12 шагов в оригинале — это объединение двух алгоритмов выше: AllGather_Y активаций и ReduceScatter_Y выходов на критическом пути, AllGather_X весов и ReduceScatter_X градиентов весов вне его. Связи много, но все AllGather и ReduceScatter меньше, потому что активации шардированы по батчу, а веса — тензорно.)

Какое сочетание FSDP и TP правильное? Простая, но ключевая максима: FSDP перемещает веса, а тензорный параллелизм перемещает активации. Значит, по мере уменьшения батча (особенно при большем параллелизме данных) тензорный параллелизм дешевеет, ведь активации на шард становятся меньше.

• Тензорный параллелизм выполняет $\mathbf{AllGather}_Y([B_X, D_Y])$, который уменьшается с ростом $X$.
• FSDP выполняет $\mathbf{AllGather}_X([D_X, F_Y])$, который уменьшается с ростом $Y$.

Так что, сочетая оба, можно ещё сильнее опустить минимальный батч на реплику. Оптимальные доли FSDP и TP можно посчитать так же, как выше.

Пусть $X$ — число чипов, отданных FSDP, $Y$ — тензорному параллелизму, $N = XY$ — общее число чипов в срезе, а $M_X$ и $M_Y$ — число осей сетки, по которым идут FSDP и TP (в сумме примерно 3). Будем моделировать только прямой проход, у которого больше всего связи на FLOP. Складывая связь в алгоритме, получаем

$$T_\text{связь FSDP}(B, X, Y) = \frac{2\cdot 2\cdot D \cdot F}{Y \cdot W_\text{ici} \cdot M_X}, \qquad T_\text{связь TP}(B, X, Y) = \frac{2 \cdot 2 \cdot B \cdot D}{X \cdot W_\text{ici} \cdot M_Y}$$

А общее время FLOPs —

$$T_\text{math} = \frac{2\cdot 2 \cdot B \cdot D \cdot F}{N \cdot C}$$

Для простоты сделаем два допущения: во-первых, разрешим $X$ и $Y$ быть нецелыми (лишь бы положительными и $XY=N$); во-вторых, будем считать, что связь по осям $X$ и $Y$ полностью перекрывается между собой. При втором допущении общее время связи —

$$T_\text{comms} = \max\left(T_\text{связь FSDP}, T_\text{связь TP}\right)$$

Прежде чем спрашивать, когда мы упираемся в вычисления, найдём $X$ и $Y$, минимизирующие общую связь. FLOPs от $X$ и $Y$ не зависят, так что оптимальны значения, просто минимизирующие связь. Запишем $T_\text{comms}$ через $X$ и $N$ (последнее фиксировано — это число чипов в системе):

$$T_\text{comms} (X) = \frac{4D}{W_\text{ici}} \max\left(\frac{F \cdot X}{N \cdot M_X}, \frac{B}{X \cdot M_Y}\right)$$

$T_\text{связь FSDP}$ монотонно растёт с $X$, а $T_\text{связь TP}$ монотонно убывает, так что максимум минимален при их равенстве, то есть при

$$\frac{FX_{opt}}{M_X} = \frac{BN}{X_{opt} M_Y} \;\rightarrow\; X_{opt} = \sqrt{\frac{B}{F} \frac{M_X}{M_Y} N}$$

Это очень полезно! Формула говорит, какая доля FSDP оптимальна при данных $B$, $F$ и $N$. Почувствуем масштаб: подставив реалистичные значения $N = 64$ (массив чипов 4×4×4), $B=48\,000$, $F=32768$, получаем примерно $X\approx 13{,}9$. Так что мы взяли бы $X$ = 16 и $Y$ = 4 — близко к рассчитанному оптимуму.

Вывод

В общем случае при обучении оптимальная доля FSDP — $X_{opt} = \sqrt{\frac{B}{F} \frac{M_X}{M_Y} N}$.

Вернёмся к вопросу, который мы задавали всем стратегиям: когда мы упираемся в вычисления? Поскольку FLOPs и связь совмещаются, мы упираемся в вычисления, когда

$$\max\left(T_\text{связь FSDP}, T_\text{связь TP}\right) < T_\text{math}$$

Обозначив $\alpha \equiv C / W_\text{ici}$ (арифметическую интенсивность ICI), упрощаем:

$$\max\left(\frac{F}{Y \cdot M_X}, \frac{B}{X \cdot M_Y}\right) < \frac{B \cdot F}{N \cdot \alpha}$$

Мы выбрали $X_{opt}$ так, чтобы обе части максимума были равны, так что можно подставить его в любую (учитывая $Y_{opt} = N/X_{opt}$). Упрощая, получаем

$$\sqrt{\frac{B\cdot F}{M_X \cdot M_Y \cdot N}} < \frac{B \cdot F}{N \cdot \alpha}$$

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

$$\frac{T_\text{math}}{T_\text{comms}} = \frac{\sqrt{BF}\sqrt{M_X M_Y}}{\alpha \sqrt{N}}$$

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

$$\frac{B}{N} > \frac{\alpha^2}{M_X M_Y F}$$

Для приближённых чисел снова подставим $F=32\,768$, $\alpha=2550$ и $M_X M_Y=2$ (так и должно быть для трёхмерной сетки). Получаем примерно $B/N > 99$. Это выигрыш примерно в восемь раз по сравнению с чистым параллелизмом данных (или FSDP), где для трёхмерной сетки $B/N$ должно превышать примерно 850.

Вывод

Сочетание тензорного параллелизма с FSDP позволяет опуститься до $B/N = 2550^2 / 2F$. Это позволяет работать с батчем всего около 100 на чип — примерно в восемь раз меньше, чем с одним FSDP.

Рисунок в оригинале (mixed-fsdp-comms-2.png): отношение времени FLOPs ко времени связи для оптимального сочетания FSDP/TP на срезе TPUv5p 4×4×4 при F=30 тысяч, в сравнении с чистым TP и чистым FSDP. Как и ожидалось, у тензорного параллелизма отношение от батча не зависит, у оптимального FSDP + TP растёт как $\sqrt{B}$, у FSDP — как $B$. Чистый FSDP выигрывает при очень больших батчах, но в промежуточном режиме, когда батч на число чипов примерно от 100 до 850, упираться в вычисления позволяет только FSDP + TP.

Ещё один пример — TPU v5p 16×16×16 (math-comms-time.png): время связи при разных схемах параллелизма в зависимости от батча; чёрная пунктирная линия — время FLOPs умножений, так что всё выше неё упирается в связь. Все стратегии упираются в связь ниже батча 6e5, что согласуется с ожидаемыми $4096 \cdot 2550^2 / (2 \cdot 8192 \cdot 4) = 4 \cdot 10^{5}$. В интерактивном графике оригинала можно двигать ползунок доли FSDP/TP: лучшее разбиение, FSDP=512 и TP=8, упирается в вычисления при батче больше ~4e5, а любое разбиение с большей долей тензорного параллелизма упирается в связь при любом батче.

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

Вы, вероятно, заметили, что до сих пор мы избегали разговора о конвейерах. Конвейер — доминирующая стратегия параллелизма на GPU, но на TPU он не так необходим. Коротко: конвейерное обучение делит слои модели между устройствами и передаёт активации между стадиями конвейера в прямом и обратном проходах. Алгоритм примерно такой:

1. Инициализировать данные на TPU 0, а веса шардировать по измерению слоёв ($W_\text{in}[L_Z, D_X, F_Y]$ для конвейера с FSDP и тензорным параллелизмом).
2. Выполнить первый слой на TPU 0, затем скопировать получившиеся активации на TPU 1 и повторять до последнего TPU.
3. Вычислить функцию потерь и её производную $\partial L / \partial x_L$.
4. На последней стадии конвейера вычислить производные $\partial L / \partial W_L$ и $\partial L / \partial x_{L-1}$, затем скопировать $\partial L / \partial x_{L-1}$ на предыдущую стадию и повторять до TPU 0.

(Рабочий) псевдокод на Python

Этот псевдокод должен работать на Cloud TPU VM. Он не очень эффективен и не реалистичен, но даёт представление о том, как данные распространяются между устройствами.

import jax
import jax.numpy as jnp

batch_size = 32
d_model = 128
d_ff = 4 * d_model

num_layers = len(jax.devices())

key = jax.random.PRNGKey(0)

# Представим, что каждый слой — просто одно умножение матриц.
x = jax.random.normal(key, (batch_size, d_model))
weights = jax.random.normal(key, (num_layers, d_model, d_model))

def layer_fn(x, weight):
  return x @ weight

# Считаем, что num_layers == num_pipeline_stages.
intermediates = [x]
for i in range(num_layers):
  x = layer_fn(x, weights[i])
  intermediates.append(x)

  if i != num_layers - 1:
    x = jax.device_put(x, jax.devices()[i+1])   # передаём активации следующей стадии

def loss_fn(batch):
  return jnp.mean(batch ** 2)  # выдуманная функция потерь

loss, dx = jax.value_and_grad(loss_fn)(x)

for i in range(num_layers - 1, -1, -1):
  _, f_vjp = jax.vjp(layer_fn, intermediates[i], weights[i])
  dx, dw = f_vjp(dx)  # вычисляем произведение dx на якобиан слоя (VJP)
  weights = weights.at[i].set(weights[i] - 0.01 * dw)  # обновляем веса

  if i != 0:
    dx = jax.device_put(dx, jax.devices()[i-1])  # передаём градиент предыдущей стадии

(Примечание переводчика: в оригинале обновление записано как weights[i] = ...; массивы JAX неизменяемы, поэтому здесь использован weights.at[i].set(...). Для запуска без TPU задайте несколько CPU-устройств, как в главе 3.)

Почему это хорошая идея? Конвейер хорош по многим причинам: стоимость связи между стадиями низкая, так что очень большие модели можно обучать даже при медленных межсоединениях. Это часто очень полезно на GPU, ведь они не связаны так плотно, как TPU по ICI.

Почему это трудно и неприятно? Вы, наверное, заметили в псевдокоде выше, что TPU 0 почти всё время простаивает! Он работает только на самом первом и последнем шаге конвейера. Период простоя называют пузырём конвейера (pipeline bubble), и бороться с ним очень неприятно. Обычно сначала пытаются смягчить его микробатчами: через конвейер пускают несколько маленьких батчей, и TPU 0 загружен хотя бы большую долю времени шага.

Второй подход — аккуратно совмещать прямое умножение $W_i @ x_i$, обратное умножение для $dx$ — $W_i @ \partial L / \partial x_{i+1}$ — и умножение для $dW$ — $\partial L / \partial x_{i+1} @ x_i$. Каждое из них требует FLOPs, так что их можно перекрыть и полностью скрыть пузырь. На рисунке из недавней статьи о DeepSeek v3 (deepseek-pipeline.png) — их расписание конвейера «без пузырей»: оранжевым — прямые умножения, зелёным — умножения dL/dx, синим — умножения dL/dW; отдавая приоритет обратным умножениям dL/dx, можно не «оставлять на мели» FLOPs.

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

Масштабирование между подами

Самый большой возможный срез TPU — суперпод TPU v5p из 8960 чипов (и 2240 хостов). Чтобы масштабироваться дальше, нужно пересечь границу сети дата-центра (DCN). У каждого хоста TPU одна или несколько сетевых карт (NIC), соединяющих хост с другими подами TPU v5p по Ethernet. Как отмечалось в главе о TPU, у каждого хоста около 200 Гбит/с (25 ГБ/с) полнодуплексной пропускной способности DCN, то есть около 6,25 ГБ/с исходящей пропускной способности на TPU.

Обычно при масштабировании за пределы одного пода внутри домена ICI делают какой-то параллелизм модели или FSDP, а между подами — чистый параллелизм данных. Пусть $N$ — число TPU, до которого мы хотим масштабироваться, а $M$ — число TPU в срезе, связанном ICI. Чтобы сделать AllReduce по DCN, можно выполнить кольцевую редукцию по набору подов, что даёт (в обратном проходе):

$$T_\text{math} = \frac{2 \cdot 2 \cdot 2 \cdot BDF}{N \cdot C}, \qquad T_\text{comms} = \frac{2 \cdot 2 \cdot 2 \cdot DF}{M \cdot W_\text{dcn}}$$

Пропускная способность связи растёт с $M$: в отличие от ICI, общая пропускная способность растёт по мере роста домена ICI — появляется больше сетевых карт. Упрощая, получаем, что $T_\text{math} > T_\text{comms}$ при

$$\frac{B}{\text{срез}} > \frac{C}{W_\text{dcn}}$$

Для TPU v5p $\frac{C}{W_\text{dcn}}$ — около 4.59e14 / 6.25e9 = 73,440. Значит, чтобы эффективно масштабироваться через DCN, нужен минимальный батч на домен ICI.

Насколько это проблема? Конкретный пример: мы хотим обучать LLaMA-3 70B на TPU v5p с батчем 2 млн токенов. У LLaMA-3 70B $F\approx 30\,000$. Из разделов выше мы знаем:

• тензорный параллелизм возможен до $Y = M_Y \cdot F / 2550 \approx 11 \cdot M_Y$;
• FSDP возможен, пока $B / N > 2550 / M_X$. Значит, при батче 2 млн и 3 осях параллелизма данных можно использовать не больше $\approx 2400$ чипов — примерно четверть пода TPU v5p;
• при сочетании FSDP + TP мы упираемся в связь при $B / N < 2550^2 / (2 \cdot 30000) = 108$, так что можно масштабироваться примерно до 18 тысяч чипов! Но максимальный размер пода TPU v5p — около 9 тысяч чипов, так что дальше придётся использовать DCN.

Коротко: есть хороший рецепт обучения с батчем 1 млн — примерно X (FSDP) = 1024 и Y (TP) = 8, — а с батчем 2 млн нужен DCN. Как отмечено выше, арифметическая интенсивность DCN — 73 440, так что нужно лишь, чтобы батч на домен ICI был больше. Для нас это тривиально: с 2 подами батч на под — 1 млн, а на TPU — 111, и это отлично (может, впритык, но теоретически верно).

Вывод

Масштабирование на несколько подов TPU с чистым параллелизмом данных довольно просто, если батч на под не меньше 73 тысяч токенов.

Главные выводы об обучении LLM на TPU

• Рост параллелизма или уменьшение батча обычно делают нас сильнее ограниченными связью, потому что уменьшают объём вычислений на чип.
• До разумной длины контекста (~32 тысячи) трансформер можно моделировать как стопку блоков MLP и определять каждую схему параллелизма тем, как она шардирует два-три главных умножения матриц на слой.
• При обучении рассматриваются 4 главные схемы параллелизма, у каждой свои требования к пропускной способности и вычислениям (параллелизм данных, FSDP, тензорный параллелизм и сочетание FSDP + тензорного параллелизма).

Стратегия	Описание
Параллелизм данных	Активации шардированы по батчу, всё остальное полностью реплицировано; в обратном проходе — AllReduce градиентов
FSDP	Активации, веса и оптимизатор шардированы по батчу; веса собираются прямо перед использованием, градиенты — ReduceScatter
Тензорный параллелизм (Megatron, параллелизм модели)	Активации шардированы по $d_\text{model}$, веса — по $d_{ff}$; активации собираются перед Win, результат разбрасывается ReduceScatter после Wout
FSDP + тензорный параллелизм	Оба вместе; FSDP собирает веса, шардированные по модели

«Формулы» каждого метода:

$$\small \begin{array}{cc} \text{Стратегия} & \text{Формула}\\ \hline \text{DP} & \text{In}[B_X, D] \cdot_D W_\text{in}[D, F] \cdot_F W_\text{out}[F, D] \rightarrow \text{Out}[B_X, D] \\ \text{FSDP} & \text{In}[B_X, D] \cdot_D W_\text{in}[D_X, F] \cdot_F W_\text{out}[F, D_X] \rightarrow \text{Out}[B_X, D] \\ \text{TP} & \text{In}[B, D_Y] \cdot_D W_\text{in}[D, F_Y] \cdot_F W_\text{out}[F_Y, D] \rightarrow \text{Out}[B, D_Y] \\ \text{TP + FSDP} & \text{In}[B_X, D_Y] \cdot_D W_\text{in}[D_X, F_Y] \cdot_F W_\text{out}[F_Y, D_X] \rightarrow \text{Out}[B_X, D_Y] \\ \hline \end{array}$$

• У каждой стратегии есть предел, после которого она упирается в сеть (связь), — он определяется вычислениями и связью на устройство. Вот вычисления и связь на слой ($X$ — FSDP, $Y$ — тензорный параллелизм):

$$\small \begin{array}{ccc} \text{Стратегия} & \text{Вычисления на слой} & \text{Связь на слой} \\ & \text{(без гейтового einsum)} & \text{(байт, прямой + обратный проход)}\\ \hline \text{DP} & 4BDF/X + 8BDF/X & 0 + 8DF \\ \text{FSDP} & 4BDF/X + 8BDF/X & 4DF + 8DF \\ \text{TP} & 4BDF/Y + 8BDF/Y & 4BD + 4BD \\ \text{FSDP + TP} & 4BDF/(XY) + 8BDF/(XY) & (4BD/X + 4DF/Y) + (8BD/X + 8DF/Y) \\ \hline \end{array}$$

• Чистый параллелизм данных редко полезен, потому что модель с состоянием оптимизатора занимает в байтах 10 × число параметров. Значит, в память редко помещается больше нескольких миллиардов параметров.
• Параллелизм данных и FSDP упираются в связь, когда $\text{батч на шард} < C / W$ — арифметической интенсивности сети. Для ICI это 2550, для DCN — около 71 000. Больше параллельных осей — порог ниже.
• Тензорный параллелизм упирается в связь, когда $\lvert Y\rvert > F / 2550$. Для большинства моделей это около 8–16-путевого. От размера батча это не зависит.
• Сочетание FSDP + тензорного параллелизма позволяет опустить батч до $2550^2 / 2F \approx 100$. Это удивительно мало.
• Параллелизм данных между подами требует минимального батча на под около 71 000, иначе он упирается в DCN.
• По сути, если батч большой или модель маленькая, всё просто: можно делать параллелизм данных или FSDP + параллелизм данных через DCN. Самое интересное — посередине.

Задачи главы 5 (условия; ответы — в оригинале):

1. Сколько параметров у LLaMA-2 13B? (Учтите три большие матрицы FFW, как в главе 4.)
2. Обучение с батчем 16 млн токенов и Adam: сколько памяти занимают параметры, состояние оптимизатора и активации (параметры в bf16, оптимизатор в fp32, три чекпоинта активаций на слой)?
3. Обучение с длиной последовательности 32 тысячи и общим батчем 3 млн токенов на срезе TPUv5p 16×16×16: какое шардирование выбрать и сколько займёт обучение?

Глава 6 в кратком изложении: обучение LLaMA 3 на TPU

Пересказ с переводом выводов и ключевых расчётов главы «Training LLaMA 3 on TPUs». В оригинале это серия вопросов с ответами — попробуйте сначала посчитать сами.

Архитектура LLaMA 3-70B (из конфигурации модели): $L = 80$ слоёв, $D = 8192$, $F = 28\,672$, $N = 64$ головы запросов, $K = 8$ голов ключей/значений, $H = 128$, словарь $V = 128\,256$.

Число параметров. FFW: $3 \cdot D \cdot F \cdot L = 8192 \cdot 8192 \cdot 3{,}5 \cdot 3 \cdot 80 \approx 56{,}3$ млрд; словарь (вход и выход): $2 \cdot 128\,256 \cdot 8192 \approx 2{,}1$ млрд; внимание: $80 \cdot (2 \cdot 8192 \cdot 64 \cdot 128 + 2 \cdot 8192 \cdot 8 \cdot 128) \approx 12$ млрд. Итого ≈ 70,4 млрд — сходится. Три большие матрицы MLP настолько больше всех остальных массивов, что при рассуждениях о памяти и FLOPs остальные параметры часто можно почти игнорировать.

FLOPs. На токен — $6 \cdot 70 \cdot 10^{9} = 4{,}2 \cdot 10^{11}$ FLOPs (около половины терафлопа). На 15 трлн токенов — $6{,}3 \cdot 10^{24}$ FLOPs. На одном TPU v5p это 435 лет; на полном поде из 8960 чипов при 40% MFU — $6{,}3 \cdot 10^{24} / (8960 \cdot 4{,}59 \cdot 10^{14} \cdot 0{,}4) \approx 3{,}8 \cdot 10^{6}$ с, то есть около 44 дней.

Память при батче 4 млн токенов (bf16-параметры, fp32-оптимизатор, 2 чекпоинта активаций на слой): параметры ~140 ГБ, состояние оптимизатора ~560 ГБ, чекпоинты активаций $2 \cdot 8192 \cdot 4 \cdot 10^{6} \cdot 2 \cdot 80 \approx 10{,}5$ ТБ, итого ~11,2 ТБ. Активации резко доминируют. Это ~117 чипов по 96 ГБ — но на них обучение шло бы почти 10 лет. Большие кластеры нужны не из-за памяти, а ради FLOPs: на 8960 чипах это ~1,3 ГБ на чип.

Шардирование на поде из 8960 чипов при батче 4 млн токенов (1024 последовательности по 4096):

• только FSDP по батчу — нельзя: последовательностей всего 1024;
• FSDP по батчу и последовательности — батч на TPU 4 * 1024 * 1024 / 8960 = 468 токенов, меньше порога 850 для трёх осей, — упираемся в связь;
• FSDP + TP — порог $2550^2 / 2F = 113$, мы выше; оптимум $X_{opt} = \sqrt{2BN/F} = \sqrt{2 \cdot 4{,}19 \cdot 10^{6} \cdot 8960 / 28\,672} \approx 1618$, то есть примерно 2048-путевой FSDP и 4-путевой тензорный параллелизм.

Вывод главы 6

LLaMA-3 с батчем 4 млн токенов можно обучать на полном поде TPU v5p, сочетая параллелизм данных (1024-путевой), последовательностный параллелизм (2-путевой) и тензорный (4-путевой), не упираясь в связь. Чистый FSDP или FSDP + последовательностный параллелизм упрутся в связь. Уравнения из главы 5 очень практичны.

Словарик

English	Русский
strong scaling	сильное масштабирование
data parallelism (DP)	параллелизм данных
fully-sharded data parallelism (FSDP), ZeRO-1/2/3	полностью шардированный параллелизм данных, шардирование ZeRO
tensor (model) parallelism, Megatron sharding	тензорный параллелизм (параллелизм модели), шардирование Megatron
pipeline parallelism, pipeline bubble, microbatch	конвейерный параллелизм, пузырь конвейера, микробатч
context / sequence parallelism, ring attention	контекстный / последовательностный параллелизм, кольцевое внимание
critical path	критический путь
ICI operational intensity $C/W$	операционная интенсивность ICI
optimizer state	состояние оптимизатора
MFU (model FLOPs utilization)	доля использованных пиковых FLOPs
critical batch size	критический размер батча
DCN, NIC	сеть дата-центра, сетевая карта
activation checkpoint	чекпоинт активаций

Footnotes

1. Мы сосредоточимся на ограничениях связи: ограничения объёма памяти важны, но при использовании рематериализации (чекпоинтинга активаций) и очень большого числа чипов при предобучении они обычно нас не сдерживают. Экспертный параллелизм для MoE мы здесь тоже не обсуждаем — он существенно расширяет пространство решений, — только базовый случай плотного трансформера. ↩

2. Adam хранит параметры и накопители первого и второго порядка. Параметры в bfloat16, а состояние оптимизатора в float32, что даёт 2 + 8 = 10 байт на параметр. ↩

3. Заметьте, что сюда не входят чекпоинты для градиентов, так что на деле это бесполезно. Это абсолютная нижняя граница при батче в 1 токен. ↩

4. Предполагаем, что разбиение идёт по сетке ICI, так что нужная пропускная способность сети — $W_\text{ici}$. ↩

5. Строго говоря, FSDP добавляет связь в прямом проходе, которой у чистого DP нет, но в той же пропорции, что и в обратном, так что на roofline связи это не влияет. Главное здесь — ZeRO-3 превращает AllReduce обратного прохода в AllGather и ReduceScatter, у которых тот же общий объём связи. ↩