How To Scale Your Model, глава 4: вся математика трансформера

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

О переводе

Полный перевод главы 4 книги How To Scale Your Model (лицензия MIT, © 2025 Google LLC), включая все задачи и приложение о Flash Attention. Цветовую разметку формул оригинала (сворачиваемые измерения — красным, батчевые — синим) мы передаём словами. Предыдущая глава — «Шардированные матрицы», следующая — «Как параллелить обучение трансформера».

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

Это та глава, которую страница «Учимся у создателей» рекомендует к ступени 5. Она даёт строгий вывод формулы $C \approx 6ND$ из урока «Законы масштабирования», считает параметры блока из урока «Блок трансформера», объясняет, зачем GQA (урок «MQA и GQA») и как работает FlashAttention. Обозначения: $B$ — батч, $T$ — длина последовательности запросов, $S$ — длина последовательности ключей, $D$ — размерность модели ($d_\text{model}$), $F$ — размерность MLP ($d_\text{ff}$), $N$ — число голов запросов, $K$ — число голов ключей/значений, $H$ — размерность головы, $L$ — число слоёв, $V$ — размер словаря.

Считаем скалярные произведения

Начнём с векторов $x$, $y$ и матриц $A$, $B$ следующих форм:

Массив	Форма
$x$	[P]
$y$	[P]
$A$	[N P]
$B$	[P M]

• Скалярное произведение $x \cdot y$ требует $P$ сложений и умножений, всего $2P$ операций с плавающей точкой.
• Произведение матрицы на вектор $Ax$ — это $N$ скалярных произведений по строкам $A$, то есть $2NP$ FLOPs.
• Произведение матриц $AB$ — это произведение матрицы на вектор для каждого из $M$ столбцов $B$, всего $2NPM$ FLOPs.
• В общем случае, если есть два многомерных массива $C$ и $D$, у которых часть измерений — сворачиваемые (CONTRACTING), а часть — батчевые (BATCHING) (например, $C[GHIJKL], D[GHMNKL]$, где $G, H$ — батчевые, а $K, L$ — сворачиваемые), то стоимость этой свёртки в FLOPs — удвоенное произведение всех измерений $C$ и $D$, где батчевые и сворачиваемые измерения учитываются по одному разу (например, $2GHIJMNKL$). Заметьте: измерение батчевое, только если оно есть в обоих сомножителях. (И ещё: множитель 2 не применяется, если сворачиваемых измерений нет и это просто поэлементное произведение.)¹

$$\begin{array}{ccc} \textrm{операция} & \textrm{FLOPs} & \textrm{данные} \\ \hline x \cdot y & 2P & 2P \\ A x & 2NP & NP + P \\ AB & 2NPM & NP + PM \\ [c_0,...,c_N] \cdot [d_0,...,d_N] & 2 \prod c_i \times \prod_{\substack{d_j \notin \textrm{BATCH} \\ d_j \notin \textrm{CONTRACT}}} d_j & \prod c_i + \prod d_j \\ \hline \end{array}$$

Обратите внимание: при умножении матриц вычисления растут кубически, $O(N^3)$, а передача данных — лишь квадратично, $O(N^2)$. Значит, с ростом размера умножения проще выйти на предел вычислений. Это очень необычно и во многом объясняет, почему мы используем архитектуры, где доминирует умножение матриц, — они хорошо масштабируются!

Анимация в оригинале (matmul-flops.gif): как из скалярных произведений складывается стоимость умножения матриц.

FLOPs прямого и обратного прохода

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

Если представить, что B — лишь одна матрица в большой сети, а A — наши входные активации, и C = A B, то производная функции потерь L по B находится по правилу цепочки:

$$\frac{\partial L}{\partial B} = \frac{\partial L}{\partial C}\frac{\partial C}{\partial B} = A^T \left(\frac{\partial L}{\partial C}\right)$$

и требует $2NPM$ FLOPs (свёртка идёт по измерению $N$). Аналогично производная потерь по A —

$$\frac{\partial L}{\partial A} = \frac{\partial L}{\partial C}\frac{\partial C}{\partial A} = \left(\frac{\partial L}{\partial C}\right) B^T$$

снова $2NPM$ FLOPs, поскольку dL/dC — матрица размера $[N, M]$. Хотя эта величина — производная не по параметру, она нужна, чтобы вычислять производные для предыдущих слоёв сети (так же, как dL/dC используется выше, чтобы найти dL/dB).

Складывая, получаем, что при обучении всего 6NPM FLOPs против 2NPM при инференсе: 2NPM в прямом проходе и 4NPM в обратном. Поскольку PM — число параметров матрицы, это простейший вид знаменитого приближения $6 \cdot \text{число параметров} \cdot \text{число токенов}$ для FLOPs обучения трансформера: каждый токен требует $6 \cdot \text{число параметров}$ FLOPs. Более корректный вывод — ниже.

flowchart LR
F["Прямой проход: C = A·B — 2NPM"] --> G1["dL/dB = Aᵀ·dL/dC — 2NPM"]
F --> G2["dL/dA = dL/dC·Bᵀ — 2NPM"]
G1 --> T["Итого на обучение: 6NPM = 6 × параметры × токены"]
G2 --> T

Откуда шестёрка в 6ND (схема переводчика): N — токены (строки A), PM — параметры матрицы B

Бухгалтерия трансформера

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

Рисунок в оригинале (transformer-diagram.png): один слой стандартного трансформера-декодера, поток идёт сверху вниз. Формы и раскладки массивов обозначены однобуквенными символами; сворачиваемые измерения показаны красным, батчевые — синим. В каждой операции форма входа дана слева вверху, форма параметра — справа вверху, а получающаяся форма — снизу; например, BTD — форма входа для «гейтового» einsum, а DF — форма весов. Порядок: RMSNorm → проекции Q, K, V (веса $DNH$, $DKH$, $DKH$) → внимание → проекция O ($NHD$) → сложение с остаточной связью → RMSNorm → $W_\text{In1}$ и $W_\text{In2}$ ($DF$ каждая) → поэлементное произведение с активацией → $W_\text{Out}$ ($FD$) → сложение с остаточной связью.

Замечание [гейтовый einsum]. На схеме используется «гейтовый einsum»: матрица повышающей проекции разделена на две ($W_\text{In1}$ и $W_\text{In2}$ выше), выходы которых поэлементно перемножаются как своего рода «функция ворот» (gating). Так делают не все LLM, поэтому иногда вы увидите одну матрицу $W_\text{In}$ и общее число параметров MLP 2DF вместо 3DF. Обычно в этом случае D и F увеличивают, чтобы число параметров осталось тем же, что и в варианте с 3 матрицами. Впрочем, та или иная форма гейтового einsum используется в LLaMA, DeepSeek и многих других моделях.

Замечание 2 [MHA]. В самовнимании T и S совпадают, а в перекрёстном внимании могут различаться. В обычном многоголовом внимании (MHA) N и K совпадают, в Multi-Query Attention (MQA) K=1, а в групповом MQA (GMQA, оно же GQA) K должно лишь делить N.

Замечание 3 [pre-norm и post-norm]. На схеме показана архитектура «pre-norm», где нормализация стоит перед остаточной связью, обычно как x + attn(norm(x)). Так сегодня устроены модели вроде LLaMA-3. В оригинальной статье о трансформере использовалась архитектура «post-norm», где layernorm стоит после остаточной связи: norm(x + attn(x)).

Общий подсчёт FLOPs и параметров

Посчитаем FLOPs трансформера на один слой (чтобы не таскать везде множитель L). Заметьте: FLOPs обучения ниже почти всегда втрое больше FLOPs инференса, так что любую сумму можно разделить на 3, чтобы получить стоимость одного прямого прохода.

MLP

MLP трансформера обычно состоит из 2 входных умножений, результаты которых поэлементно объединяются, и одного выходного умножения (сворачиваемые измерения — D в первых двух строках и F в последней):

$$\begin{array}{ccc} \textrm{операция} & \textrm{FLOPs обучения} & \textrm{параметры} \\ \hline \\ A[B,T,D] \cdot W_{in1}[D, F] & 6BTDF & DF \\[10pt] A[B,T,D] \cdot W_{in2}[D, F] & 6BTDF & DF \\[10pt] \sigma\left(A_{in1}\right)[B,T, F] * A_{in2}[B,T, F] & O(BTF) \\[10pt] A[B,T,F] \cdot W_{out}[F, D] & 6BTDF & DF \\[10pt] \hline \\ & \approx 18BTDF & 3DF \end{array}$$

Внимание

Для общего случая группового внимания (grouped-query attention) с разным числом голов Q и KV будем считать размерность головы H одинаковой для проекций Q, K, V и оценим стоимость умножений QKVO:

$$\begin{array}{ccc} \textrm{операция} & \textrm{FLOPs обучения} & \textrm{параметры} \\ \hline \\ A[B,T,D] \cdot W_{Q}[D, N, H] & 6BTDNH & DNH \\[10pt] A[B,T,D] \cdot W_{K}[D, K, H] & 6BTDKH & DKH \\[10pt] A[B,T,D] \cdot W_{V}[D, K, H] & 6BTDKH & DKH \\[10pt] A[B,T,N, H] \cdot W_{O}[N, H, D] & 6BTDNH & DNH \\[10pt] \hline \\ & 12BTD(N+K)H & 2D(N+K)H \end{array}$$

Само скалярное внимание (dot-product attention) устроено тоньше: по сути, это умножение $TH \cdot HS$, батчевое по измерениям $B$ и $K$, затем softmax, затем умножение $TS \cdot SH$, снова батчевое по $B$ и $K$ (в первой строке сворачивается $H$, в третьей — $S$):

$$\begin{array}{cc} \textrm{операция} & \textrm{FLOPs обучения} \\ \hline \\[3pt] Q[B, T, K, G, H] \cdot K[B, S, K, H] & 6BTSKGH = 6BTSNH \\[3pt] \textrm{softmax}_S \;\; L[B, T, S, K, G] & O(BTSKG) = O(BTSN) \\[3pt] S[B, T, S, K, G] \cdot V[B, S, K, H] & 6BTSKGH = 6BTSNH \\[3pt] \hline \\ & \approx 12BTSNH = 12BT^2NH \\ \end{array}$$

(Здесь $G = N / K$ — число голов запросов на одну голову ключей.)

Замечание [каузальная маска]. Большинство современных трансформеров используют каузальную маску вместо полного двунаправленного внимания. Тогда полезные FLOPs скалярного внимания уменьшаются вдвое. Чтобы получить это сокращение на практике, нужно ядро внимания (attention kernel), а не наивный einsum.

Прочие операции

В трансформере есть и другие операции. Layernorm сравнительно дёшевы, и в оценках первого порядка ими можно пренебречь. Заметьте, что в каждом слое их обычно две (одна перед вниманием и одна перед MLP). Есть ещё последнее огромное (хотя и не на каждом слое) умножение на матрицу «разэмбеддинга» (unembedding) — выходного слоя в словарь.

$$\begin{array}{ccc} \textsf{операция} & \textsf{FLOPs обучения} & \textsf{параметры} \\ \hline \\ 2 \times \textrm{layernorm}_D \;\; A[B,T,D] & O\left(BTD\right) & 2D \\[10pt] A[B,T,D] \cdot W_{unembed}[D, V] & 6BTDV & DV \\ \end{array}$$

Общее практическое правило для FLOPs трансформера

Если пренебречь стоимостью скалярного внимания (это разумно при обучении с не слишком длинным контекстом), то общее число FLOPs по всем слоям

$$\begin{aligned} (18BTDF + 12BTD(N+K)H)L &= 6 \cdot BT \cdot (3DF + 2D(N+K)H)L \\ &= 6 \cdot \textrm{число токенов} \cdot \textrm{число параметров} \end{aligned}$$

Отсюда знаменитое практическое правило оценки FLOPs плотного трансформера без учёта FLOPs внимания. (Разэмбеддинг — ещё одно простое умножение с $6BTDV$ FLOPs и $DV$ параметрами — подчиняется тому же правилу.)

Пример с числами (примечание переводчика)

Модель вроде LLaMA-2 7B: $D = 4096$, $F = 11008$, $N = K = 32$, $H = 128$, $L = 32$, $V = 32000$. Параметры слоя: MLP $3DF \approx 135{,}3$ млн, внимание $2D(N+K)H = 4D^2 \approx 67{,}1$ млн, итого ≈ 202,4 млн; на 32 слоя — 6,48 млрд; плюс эмбеддинги и разэмбеддинг $2DV \approx 0{,}26$ млрд — всего ≈ 6,74 млрд (официально 6,74B). Обучение на 2 трлн токенов: $6 \cdot 6{,}74 \cdot 10^{9} \cdot 2 \cdot 10^{12} \approx 8{,}1 \cdot 10^{22}$ FLOPs.

Доля внимания при росте длины контекста

Если всё же учесть скалярное внимание и принять $F=4D$, $D=NH$ (как обычно) и $N=K$, то отношение FLOPs скалярного внимания к FLOPs всех умножений (включая проекции внимания):

$$\small{\frac{\textrm{FLOPs внимания}}{\textrm{FLOPs умножений}} = \frac{12BT^2NH}{18BTDF + 24BTDNH} = \frac{12BT^2D}{4 \cdot 18 BTD^2 + 24 BTD^2} = \frac{12BT^2D}{96 BTD^2} = \frac{T}{8D}}$$

Вывод: FLOPs скалярного внимания начинают доминировать при обучении, только когда T больше 8D. При D ~ 8 тысяч это ~64 тысячи токенов. Это логично: значит, чем больше MLP, тем менее критичны FLOPs внимания. Для больших моделей квадратичная стоимость внимания на самом деле не такое уж большое препятствие для обучения на длинном контексте. Но для моделей поменьше, например Gemma-27B с D=4608, внимание начинает доминировать примерно с длины 37 тысяч². Стоимость длинного контекста помогает снизить и Flash Attention, которую мы кратко обсуждаем в приложении А.

Разная математика

Разреженность и смесь экспертов

Было бы упущением не обсудить кратко модели со смесью экспертов (Mixture of Experts, MoE), которые заменяют единственный плотный блок MLP стандартного трансформера набором независимых MLP, между которыми можно динамически маршрутизировать. В первом приближении MoE — это обычная плотная модель с E блоками MLP на слой вместо одного. Каждый токен активирует $k$ из этих экспертов, обычно $k \ll E$. Отношение $E / k$ называют разреженностью (sparsity), обычно оно от 8 до 64 (например, у DeepSeek v3 фактически $k=8$, $E=256$). По сравнению с плотной версией это увеличивает число параметров в $O(E)$ раз, а общее число активируемых на токен параметров — в $k$ раз.

Рисунок в оригинале (moe.png): пример слоя MoE с $n$ экспертами. «Ворота» (gating) направляют каждый токен к $k$ из них, и выходы этих $k$ MLP суммируются. Число параметров — $n$ размеров эксперта, но на каждый токен используются только $k$.

По сравнению с плотной моделью MoE добавляет новую связь — прежде всего два AllToAll (один перед блоком MoE и один после), которые отправляют токены к нужному эксперту и возвращают их на «родное» устройство³. Однако, как мы видели в прошлой главе, каждый AllToAll стоит лишь 1/4 сравнимого AllGather по одной оси (для двунаправленного кольца).

Чекпоинтинг градиентов

Обратное распространение как алгоритм меняет вычисления на память. Вместо обратного прохода за $O(n_\text{layers}^2)$ FLOPs оно требует $O(n_\text{layers})$ памяти: сохраняются все промежуточные активации прямого прохода. Это лучше квадратичных вычислений, но невероятно дорого по памяти: модель с $B \cdot T=4M$ (4 млн токенов в батче), L=64 и D=8192, которая избегает всех лишних вычислений в обратном проходе, должна была бы сохранить примерно $2 \cdot 20 \cdot B \cdot T \cdot D \cdot L = 84$ ТБ активаций в bfloat16. Множитель 20 получается, если (примерно) посчитать каждый промежуточный узел на схеме трансформера выше, ведь, например,

$$f(x) = \exp(g(x))$$ $$\frac{df}{dx} = \exp(g(x)) \cdot \frac{dg}{dx}$$

так что, чтобы не пересчитывать, нужно сохранить из прямого прохода и $g(x)$, и $\exp(g(x))$. Чтобы не тратить столько памяти, можно сохранять лишь часть промежуточных активаций. Вот несколько стратегий, которые мы используем.

• Rematerialization по блокам (block remat): сохранять только вход каждого слоя. Это самый агрессивный метод из тех, что мы используем: всего 1 чекпоинт на слой, то есть в примере выше сохранилось бы лишь 4,2 ТБ. Зато в обратном проходе приходится повторить практически все FLOPs прямого, и FLOPs растут с $6 \cdot \text{параметры} \cdot \text{токены}$ примерно до $8 \cdot \text{параметры} \cdot \text{токены}$.
• Только большие умножения: ещё одна простая политика — сохранять только выходы больших умножений матриц. Тогда в обратном проходе не нужно пересчитывать большие умножения, но остальные функции активации и части внимания всё же пересчитываются. Это сокращает 20 на слой примерно до 7.

Список далеко не полный. В JAX этим обычно управляют jax.remat/jax.checkpoint (подробнее здесь).

То же в PyTorch (примечание переводчика)

Аналог block remat — torch.utils.checkpoint.checkpoint(block, x, use_reentrant=False) вокруг каждого блока трансформера или model.gradient_checkpointing_enable() в Hugging Face Transformers. Цена, как и в книге: примерно +33% вычислений (8 вместо 6 на параметр и токен) за многократное сокращение памяти активаций.

KV-кэш

Как мы увидим в главе 7, у инференса LLM две ключевые части: prefill и генерация.

• Prefill обрабатывает длинный промпт и сохраняет его активации внимания в кэше ключей и значений (KV Cache) для использования при генерации — точнее, проекции ключей и значений в блоке внимания.
• Генерация объединяет в батч несколько таких KV-кэшей и сэмплирует токены из каждого.

Каждый KV-кэш — это, по сути, массив размера $[2, S, L, K, H]$, где 2 соответствует ключам и значениям. Это довольно много! Общий размер кэша ключей и значений в int8 — $2SLKH$. Для модели среднего размера с контекстом 8 тысяч, 64 слоями и $KH = NH = D = 8192$ это $2 \cdot 8192 \cdot 64 \cdot 8192 = 8\text{GiB}$. Видно, почему хочется использовать GMQA с $K \ll N$.

Что стоит вынести из этой главы?

• Общее число параметров и FLOPs трансформера посчитать довольно просто; сводка ниже для MHA (батч B, словарь V, последовательность длины T, D=d_model, F=d_ff):

Компонент	Параметров на слой	FLOPs обучения на слой
MLP	3DF	18BTDF
Внимание	4DNH	24BTDNH + 12BT2NH
Прочее	2D	BTD
Словарь	DV (всего, не на слой)	12BTDV

• Параметры блока MLP доминируют в общем числе параметров, и блок MLP доминирует в бюджете FLOPs, пока длина последовательности $T$ меньше $8D$.
• Общий бюджет FLOPs при обучении хорошо приближается величиной $6 \cdot \text{число параметров} \cdot \text{число токенов}$ при разумной длине контекста.
• При инференсе KV-кэш занимает примерно $2 \cdot S \cdot L \cdot K \cdot H$ на кэш (где K — число голов KV), хотя изменения архитектуры часто позволяют это уменьшить.

Примечание переводчика к таблице

В строке «Словарь» авторы указывают 12BTDV, хотя выше разэмбеддинг стоит 6BTDV; вероятно, учтены и входной эмбеддинг, и выходной слой при несвязанных весах (2DV параметров, как в задаче 1 ниже). В оценках первого порядка разница не важна.

Несколько задач

Вопрос 1. Сколько параметров у модели с $D=4096$, $F=4 \cdot D$, $V=32\,000$ и $L=64$? Какая доля из них — параметры внимания? Каков размер KV-кэша на токен? Считайте, что $N\cdot H=D$, внимание многоголовое, а KV хранятся в int8.

Ответ

1. Всего параметров примерно $L \cdot (3DF + 4DNH + 2D) + 2DV$ (с учётом двух layernorm на слой). Для данных чисел это $64 \cdot (3 \cdot 4e3 \cdot 16e3 + 4 \cdot 4e3 \cdot 4e3 + 2 \cdot 4e3) + 2 \cdot 4e3 \cdot 32e3 = 16e9$, то есть 16 млрд параметров.
2. Доля параметров внимания в общем случае — $4DNH / (4DNH + 3DF) = 4D^2 / (4D^2 + 12D^2) = 1/4$. То есть примерно 1/4 параметров приходится на внимание.
3. На токен KV-кэш занимает $2 \cdot L \cdot N \cdot H = 2 \cdot 64 \cdot 4096$ в int8, то есть 512 KiB / token.

Вопрос 2. Сколько всего FLOPs нужно, чтобы выполнить A[B_X, D_Y] *_D W[D_Y, F] на {'X': 4, 'Y': 8, 'Z': 4}? Сколько FLOPs выполняет каждый TPU?

Ответ

Общие «теоретические» FLOPs операции — $2 \cdot B \cdot D \cdot F$. Но вычисление не шардировано по измерению Z, так что на самом деле мы делаем в Z раз больше FLOPs, то есть всего $2 \cdot B \cdot D \cdot F \cdot Z$. Поскольку по остальным измерениям вычисление шардировано, на устройство приходится примерно $2 \cdot B \cdot D \cdot F / (X \cdot Y)$.

Вопрос 3. Сколько FLOPs в операции $A[I,J,K,L] * B[I,J,M,N,O] \rightarrow C[K,L,M,N,O]$?

Ответ

По правилу выше I и J — сворачиваемые измерения, а K, L, M, N и O — несворачиваемые. «Батчевых измерений» нет, так что это просто $2 \cdot I \cdot J \cdot K \cdot L \cdot M \cdot N \cdot O$ — произведение всех осей. Будь у нас общая ось, она учитывалась бы один раз.

Вопрос 4. Какова арифметическая интенсивность группового внимания с несколькими запросами (без проекций Q/K/V/O)? Дайте ответ как функцию длин последовательностей запросов и KV — T и S — и группового множителя G. При какой длине контекста внимание упирается в FLOPs? С учётом пропускной способности HBM наших TPU постройте график относительной стоимости внимания по сравнению с блоком FFW при росте длины контекста. Подсказка: считайте, что реализация внимания эффективна и не делает лишних чтений и записей. Рассмотрите оба предела: T = S и T ≪ S.

Ответ

Самовнимание требует загрузить активации $Q$, $K$ и $V$, затем вычислить $\text{softmax}(Q \cdot K) \cdot V$, затем записать результат обратно в HBM. Это будет сделано с Flash Attention, так что у расчёта есть оговорки, но в целом в bf16 самовнимание выполняет

$$\text{Q[B,T,N,H]} \rightarrow_\text{reshape} \text{Q[B, T, K, G, H]} \cdot \text{K[B, S, K, H]} \rightarrow \text{O[B, T, S, K, G]}$$$$U=\text{softmax}_S(\text{O[B, T, S, K, G]})$$$$\text{U[B, T, S, K, G]} \cdot \text{V[B, S, K, H]} \rightarrow \text{X[B, T, K, G, H]}$$

Итого байт: $2 \cdot \text{sizeof}(Q) + 2 \cdot \text{sizeof(K или V)} = 4BTNH + 4BSKH = 4BHK \cdot (TG + S)$; всего FLOPs — $4BTSNH + O(BTSN)$; арифметическая интенсивность — $4BTSKGH / (4BHK \cdot (TG + S))$.

Так что во время prefill $S=T$, и интенсивность — $4BT^2KGH / 4BHKT \cdot (G+1) = TG/(G + 1) = O(T)$. Во время генерации $T=1$, и получаем $4BSKGH / (4BHK \cdot (G + S)) = SG / (G + S) \rightarrow G$ при очень большом $S$. В зависимости от трактовки вопроса, при prefill или обучении самовнимание упирается в вычисления при S=240 (если нет шардирования по последовательности). При генерации мы никогда не упираемся в вычисления, потому что $G$ мал. Тем не менее видно, что рост $G$ приближает нас к упору в вычисления.

Вопрос 5. При какой длине последовательности FLOPs самовнимания равны FLOPs проекций QKVO?

Ответ

Это просто вопрос, когда $24BTDNH = 12BT^2NH$. Упрощая, получаем $2D = T$, то есть, например, при $D=4096$ это $8192$. Значит, при большинстве разумных длин контекста FLOPs умножений матриц больше.

Вопрос 6. Допустим, при прямом проходе мы сохраняем только выходы 7 главных умножений слоя трансформера (Q, K, V, O и три матрицы FFW). Сколько дополнительных FLOPs нужно «рематериализовать» в обратном проходе?

Ответ

Если сохранять только выходы семи умножений (Q, K, V, O, W₁, W₂, W₃), то в обратном проходе нужно пересчитать два умножения внимания

$$QK^{\top} \quad\text{и}\quad \operatorname{softmax}(QK^{\top})V$$

чтобы получить $\frac{\partial L}{\partial W_\text{O}}$.

Каждое из них — умножение $T \times T$, батчевое по $B$ последовательностям и $N$ головам, так что дополнительные FLOPs —

$$4 \; B \, T^{2} \, N \, H.$$

Другие пересчитываемые операции:

1. $O(BTD)$ для $\frac{\partial L}{\partial W_\text{In1}}$ и $\frac{\partial L}{\partial W_\text{In2}}$.
2. И $O(BTF)$ для $\frac{\partial L}{\partial W_\text{Out}}$.

Вопрос 7. DeepSeek v3, по их словам, обучалась 2,79 млн часов H800 на 14,8 трлн токенов (источник). Учитывая, что у неё 37 млрд активируемых параметров, какой примерно загрузки железа они добились? Подсказка: они использовали FLOPs в FP8 без структурированной разреженности.

Ответ

По спецификации — 3026 TFLOPs/s в FP8 с разреженностью, или обычно вдвое меньше (1.513e15 FLOPs/s) без неё. 2,79 млн часов H800 — это 2.79e6 * 1.513e15 * 60 * 60 = 1.52e25 FLOPs всего. При 37 млрд активируемых параметров этот запуск обучения должен был потребовать около 6 * 37e9 * 14.8e12 = 3.3e24 FLOPs. Значит, использование FLOPs — около 3.3e24 / 1.52e25 = 21.7%.

Вопрос 8. У моделей со смесью экспертов (MoE) $E$ копий стандартного плотного блока MLP, и каждый токен активирует $k$ из этих экспертов. Какой батч в токенах нужен, чтобы MoE с весами в int8 упиралась в вычисления на TPU v5e? Каково это число для DeepSeek с 256 (маршрутизируемыми) экспертами и $k=8$?

Ответ

Поскольку у нас $E$ копий каждого эксперта, в int8 для каждой матрицы весов нужно загрузить $E \cdot D \cdot F$ байт. Поскольку каждый токен активирует $k$ экспертов, для каждой матрицы весов получаем $2\cdot k \cdot B \cdot D \cdot F$ FLOPs. Чтобы упираться в вычисления с весами int8 и FLOPs в bfloat16, арифметическая интенсивность (FLOPs на загруженный байт) должна превышать ~240 FLOPs/байт у TPU, то есть $(2\cdot k \cdot BDF) / EDF > 240$, или $k \cdot B / E > 120$.

Значит, нужно $B > 120 \cdot E / k$. Для DeepSeek это $B > 120 \cdot 256 / 8 = 3840$. Это удивительно большой батч для этапа генерации.

Это всё для главы 4. Глава 5 — о масштабировании обучения трансформера — здесь.

Приложение

Приложение А: как работает Flash Attention?

Традиционное возражение против масштабирования трансформеров на очень длинный контекст — FLOPs и память внимания растут квадратично с длиной контекста. Действительно, произведение QK во внимании имеет форму $[B, T, S, N]$, где B — размер батча, T и S — длины последовательностей Q и K, а N — число голов, но у этого утверждения серьёзные оговорки:

1. Как мы отмечали, хотя рост квадратичный, FLOPs внимания доминируют, только когда $T > 8 \cdot D$, а при обучении память одной матрицы внимания мала по сравнению со всеми весами и чекпоинтами активаций в памяти, особенно при шардировании.
2. Чтобы вычислить внимание, не нужно материализовывать всю матрицу внимания! Можно вычислять локальные суммы и максимумы и никогда не держать в памяти больше небольшого куска массива. Общие FLOPs всё равно квадратичны, но давление на память резко снижается.

Второе наблюдение впервые сделали Rabe et al. 2021, а затем — в статье о Flash Attention (Dao et al. 2022). Основная идея — вычислять внимание кусками по K/V: вычисляется локальный softmax и вспомогательная статистика, которые передаются следующему куску, а он объединяет их со своим локальным. А именно, вычисляются:

1. M: текущий максимум $q \cdot k$ по измерению последовательности;
2. O: текущий полный результат softmax-внимания по измерению последовательности;
3. L: текущий знаменатель $\sum_i \exp(q \cdot k_i - \text{текущий максимум})$.

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

$$\text{Attn}(Q, K, V) = \sum_i \frac{\exp(Q \cdot K_i - \max_j Q \cdot K_j) V_i}{\sum_l \exp(Q \cdot K_l - \max_j Q \cdot K_j)}$$

Максимум вычитается ради численной устойчивости, и его можно вычесть без влияния на результат, поскольку $\sum_i \exp(a_i + b) = \exp(b) \sum \exp(a)$. Посмотрим только на знаменатель: представим, что у нас два соседних куска векторов ключей, $K^1$ и $K^2$, и для каждого мы посчитали локальные суммы softmax $L^1$ и $L^2$:

$$L^1 = \sum_i \exp(Q \cdot K_i^1 - \max_j Q \cdot K_j^1)$$ $$L^2 = \sum_i \exp(Q \cdot K_i^2 - \max_j Q \cdot K_j^2)$$

Тогда их можно объединить в полную сумму softmax для двух кусков вместе:

$$L^\text{combined} = \exp(M^1 - \max(M^1, M^2)) \cdot L^1 + \exp(M^2 - \max(M^1, M^2)) \cdot L^2$$

где

$$M^1 = \max_j Q \cdot K_j^1 \text{ и } M^2 = \max_j Q \cdot K_j^2$$

То же можно сделать и для полного softmax — так получается способ накапливать сколь угодно большие суммы softmax. Полный алгоритм из статьи о Flash Attention — на рисунке flash-algo.png в оригинале.

С точки зрения железа это позволяет поместить кусок Q в VMEM (в алгоритме выше она называется SRAM на кристалле), так что на каждой итерации загружаются только куски KV, и арифметическая интенсивность растёт. В VMEM можно держать и текущую статистику.

# Примечание переводчика: онлайн-softmax по кускам ключей — суть Flash Attention для одного запроса q.
import numpy as np

rng = np.random.default_rng(0)
q, K, V = rng.normal(size=8), rng.normal(size=(1000, 8)), rng.normal(size=(1000, 4))

m, l, o = -np.inf, 0.0, np.zeros(4)            # текущий максимум, знаменатель, ненормированный выход
for start in range(0, 1000, 128):              # куски по 128 ключей
    s = K[start:start + 128] @ q               # оценки q·k для куска
    m_new = max(m, s.max())
    scale = np.exp(m - m_new)                  # пересчёт старых сумм к новому максимуму
    p = np.exp(s - m_new)
    l = l * scale + p.sum()
    o = o * scale + p @ V[start:start + 128]
    m = m_new

w = np.exp(K @ q - (K @ q).max()); ref = (w / w.sum()) @ V
print(np.allclose(o / l, ref))                 # True: результат совпадает с обычным softmax

Последний тонкий момент, который стоит подчеркнуть, — свойство softmax во внимании, благодаря которому вычисление VJP (производной обратного режима) Flash Attention становится практичным для обучения. Определим промежуточный массив softmax:

$$S_{ij} = \frac{e^{\tau q_i \cdot k_j}}{\sum_l e^{\tau q_i \cdot k_l}}$$

Во внимании мы получаем dS из массивов обратного режима dO и V:

$$dS_{ij} = dO_{id} \cdot_d V_{jd} = \sum_d dO_{id} V_{jd}$$

При обратном распространении этого градиента к Q и K

$$d(q_i \cdot k_j) = (dS_{ij} - S_{ij} \cdot_j dS_{ij}) S_{ij}$$

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

$$\begin{aligned} S_{ij} \cdot_j dS_{ij} &= \sum_j \frac{e^{\tau q_i \cdot k_j}}{\sum_k e^{\tau q_i \cdot k_k}} \sum_d dO_{id} V_{jd} \\ &= \sum_d dO_{id} \sum_j \frac{e^{\tau q_i \cdot k_j}}{\sum_k e^{\tau q_i \cdot k_k}} V_{jd} \\ &= \sum_d dO_{id} O_{id} \\ &= dO_{id} \cdot_d O_{id} \end{aligned}$$

Эта замена критически важна, чтобы реализовать для VJP локальное вычисление по блокам последовательности, и она открывает путь к другим хитрым схемам шардирования вроде кольцевого внимания (ring attention).

Словарик

English	Русский
contracting / batching dimension	сворачиваемое / батчевое измерение
forward / backward (reverse-mode) FLOPs	FLOPs прямого / обратного прохода
gating einsum (GLU)	гейтовый einsum
MHA, MQA, GQA (GMQA)	многоголовое внимание, внимание с одним ключом, групповое внимание
pre-norm / post-norm	нормализация до / после остаточной связи
dot-product attention	скалярное внимание
causal mask	каузальная маска
unembedding	разэмбеддинг (выходной слой в словарь)
Mixture of Experts, sparsity, routing	смесь экспертов, разреженность, маршрутизация
gradient checkpointing, rematerialization	чекпоинтинг градиентов, рематериализация (пересчёт)
KV cache, prefill, generation	KV-кэш, prefill (обработка промпта), генерация
online softmax, running max	онлайн-softmax, текущий максимум
VJP (vector-Jacobian product)	произведение вектора на якобиан (обратный режим)
ring attention	кольцевое внимание

Footnotes

1. Сворачиваемые измерения — оси, по которым в операции идёт суммирование (они есть в обоих входах, но не в выходе), как внутреннее измерение при умножении матриц. Батчевые измерения — общие оси, которые есть в обоих входах и без изменений переходят в выход; они нумеруют независимые подзадачи и при подсчёте FLOPs не перемножаются между собой. В терминах einsum: метки, которые есть в обоих входах и в выходе, — батчевые; метки, которые есть в обоих входах, но отсутствуют в выходе, — сворачиваемые. ↩

2. Заметьте, что некоторые современные открытые модели вводят локальное внимание или другие оптимизации, которые снижают стоимость внимания и меняют этот roofline. ↩

3. Строго говоря, это происходит, только если данные или последовательность шардированы по той же оси, что и эксперты. ↩