How To Scale Your Model, глава 1: всё о roofline

Когда мы запускаем алгоритмы на железе, нас ограничивают три вещи: как быстро компьютер считает (операций в секунду), какая у него пропускная способность для перемещения данных (байт в секунду) и сколько всего памяти есть для хранения данных (байт). Эти ограничения — «потолки» (roofline) — позволяют оценить время вычисления сверху и снизу. — подзаголовок главы.

О переводе

Книга How To Scale Your Model («Как масштабировать свою модель») опубликована в репозитории jax-ml/scaling-book под лицензией MIT (© 2025 Google LLC). Это полный перевод главы 1, включая задачи и ответы. Сноски авторов — внизу страницы, примечания переводчика оформлены отдельными блоками. Рисунки описаны словами со ссылками на оригинал. Путеводитель по всей книге и список переведённых глав — на странице «How To Scale Your Model: обзор». Следующая глава — «Как устроены TPU».

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

Глава — основа всей книги и прямое продолжение урока «Как устроена видеокарта и куда уходит время», где та же арифметика сделана на числах RTX 4080. Понадобятся: умножение матриц и его стоимость $2BDF$ (урок «Законы масштабирования»), понятие bfloat16 (урок «Смешанная точность»). Вывод главы «матмул упирается в вычисления, когда в батче больше ~240 токенов» объясняет, почему decode в инференсе упирается в память (урок «Prefill и decode, узкие места»), а задачи про int8 — к уроку «Квантование».

Куда уходит время?

Начнём с предельно простого вопроса: почему алгоритм работает 50 мс, а не 50 с или 5 мс? Что на самом деле происходит внутри модели и занимает заметное время, и сколько времени это должно занимать?

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

$$T_\text{math} = \frac{\text{FLOPs вычисления}}{\text{FLOPs/s ускорителя}}$$

Например, NVIDIA H100 выполняет около $9{,}89 \cdot 10^{14}$ FLOPs/s в bfloat16¹, а TPU v6e — $9{,}1 \cdot 10^{14}$ FLOPs/s². Значит, $10^{12}$ FLOPs на H100 займут (примерно) 1e12 / 9.89e14 = 1.01ms, а на TPU v6e — 1e12 / 9.1e14 = 1.1ms³.

Связь внутри чипа. Внутри ускорителя тензоры нужно передавать между памятью ускорителя (HBM) и вычислительными ядрами. Пропускную способность этого канала называют «пропускной способностью HBM» (HBM bandwidth)⁴. У H100 это около 3,35 ТБ/с, у TPU v6e — около 1,6 ТБ/с.

Связь между чипами. Когда модель распределена по нескольким ускорителям, тензоры часто приходится передавать между ними. Обычно для этого у железа есть несколько вариантов (ICI, DCN и PCIe), каждый со своей пропускной способностью.

Внутри чипа или между чипами — мы измеряем связь в байтах в секунду и оцениваем общее время связи так:

$$T_\text{comms} = \frac{\text{байт передачи}}{\text{пропускная способность сети или памяти, байт/с}}$$

Обычно (но не всегда) вычисления внутри чипа можно совместить по времени со связью внутри чипа и между чипами. Значит, время обучения и инференса можно оценить снизу максимумом из времени вычислений и времени связи. А сверху — их суммой. На практике мы оптимизируем по максимуму: алгебра проще, и обычно можно приблизиться к этой границе, совмещая связь и вычисления. Если оптимизировать с прицелом на максимум, нижняя и верхняя оценки различаются не более чем вдвое, поскольку $T_\text{math} + T_\text{comms} \leq 2 \cdot \max(T_\text{math}, T_\text{comms})$. Дальше точность можно повысить, моделируя «области перекрытия» и накладные расходы, — для этого стоит профилировать конкретную модель и целевую систему.

$$T_\text{lower}=\max(T_\text{math}, T_\text{comms})$$ $$T_\text{upper} = T_\text{math} + T_\text{comms}$$

Если считать, что связь и вычисления совмещаются идеально, то при $T_\text{math} > T_\text{comms}$ железо загружено полностью. Это называют «упором в вычисления» (compute-bound). При $T_\text{comms} > T_\text{math}$ мы, как правило, «упираемся в связь» (communication-bound)⁵, и хотя бы часть FLOPs/s ускорителя пропадает впустую в ожидании передачи данных. Понять, во что упрётся операция — в вычисления или в связь, — можно по её арифметической интенсивности (arithmetic intensity, или operational intensity).

Определение: арифметическая интенсивность алгоритма — отношение общего числа FLOPs, которые он выполняет, к числу байт, которые ему нужно передать — внутри чипа или между чипами.

$$\text{арифметическая интенсивность} = \frac{\text{FLOPs вычисления}}{\text{байт передачи}}$$

Арифметическая интенсивность измеряет «FLOPs на байт» данной операции. В первом приближении, когда арифметическая интенсивность высока, $T_\text{math}$ велико по сравнению с $T_\text{comms}$, и мы обычно используем большую часть доступных FLOPs. Когда наоборот, мы тратим больше времени на связь, а FLOPs пропадают. Точка, где одно сменяется другим, — «пиковая арифметическая интенсивность» железа: отношение пиковых FLOPs/s ускорителя к его пропускной способности.

$$\begin{aligned} T_\text{math} > T_\text{comms} &\Leftrightarrow \frac{\text{FLOPs вычисления}}{\text{FLOPs/s ускорителя}} > \frac{\text{байт передачи}}{\text{байт/с канала}} \\[0.5em] &\Leftrightarrow \frac{\text{FLOPs вычисления}}{\text{байт передачи}} > \frac{\text{FLOPs/s ускорителя}}{\text{байт/с канала}} \\[0.5em] &\Leftrightarrow \text{интенсивность}(\text{вычисления}) > \text{интенсивность}(\text{ускоритель}) \end{aligned}$$

Величина $\text{интенсивность}(\text{ускоритель})$ — это арифметическая интенсивность, при которой ускоритель выходит на свои пиковые FLOPs/s. Для MXU в TPU v5e это около 240 FLOPs/байт: TPU выполняет 1.97e14 FLOPs/s и загружает из HBM 8.2e11 байт/с⁶. Значит, если арифметическая интенсивность алгоритма ниже 240 FLOPs/байт, он упрётся в загрузку байт и железо будет использоваться плохо⁷. Рассмотрим один такой пример.

Пример (скалярное произведение). Чтобы вычислить скалярное произведение двух векторов в bfloat16, x • y: bf16[N], bf16[N] → bf16[1], нужно загрузить из памяти $x$ и $y$ — по $2 \cdot N = 2N$ байт каждый, выполнить $N$ умножений и $N-1$ сложений и записать в HBM 2 байта:

$$\text{интенсивность}(\text{скалярное произведение}) = \frac{\text{всего FLOPs}}{\text{всего байт}} = \frac{N + N - 1}{2N + 2N + 2} = \frac{2N - 1}{4N + 2} \rightarrow \frac{1}{2}$$

при $N\rightarrow\infty$. То есть арифметическая интенсивность скалярного произведения — $\frac{1}{2}$: оно делает 0,5 операции с плавающей точкой на каждый загруженный байт. Это ниже интенсивности железа, и мы упрёмся в связь⁸.

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

Пусть $N = 10^6$. Загрузить нужно $4 \cdot 10^6$ байт, операций — около $2 \cdot 10^6$. На TPU v5e загрузка займёт $4 \cdot 10^6 / 8{,}2 \cdot 10^{11} \approx 4{,}9$ мкс, а сами вычисления на MXU — $2 \cdot 10^6 / 1{,}97 \cdot 10^{14} \approx 0{,}01$ мкс. Почти всё время чип ждёт данные: это и есть «упор в память».

Как выглядит roofline

Компромисс между памятью и вычислениями можно показать на roofline-графике: по вертикали — пиковые достижимые FLOPs/s алгоритма на нашем железе (пропускная способность вычислений), по горизонтали — арифметическая интенсивность алгоритма. Оба масштаба логарифмические.

Рисунок в оригинале (roofline-improved.png): пример roofline-графика с двумя алгоритмами разной арифметической интенсивности (Algo 1 и Algo 2) и их теоретическими пиковыми скоростями при двух разных пропускных способностях (BW1 и BW2). Линия сначала идёт наклонно вверх, затем упирается в горизонтальный «потолок» пиковых FLOPs/s. В красной области алгоритм ограничен пропускной способностью при обеих пропускных способностях и теряет часть пиковых FLOPs/s железа. Жёлтая область ограничена пропускной способностью только при меньшей (BW1). Зелёная область упирается в вычисления при любой пропускной способности: здесь мы уже используем пиковые FLOPs/s ускорителя, и ни рост пропускной способности, ни рост интенсивности ничего не дают.

Выше видно: с ростом интенсивности (слева направо) производительность алгоритма (в FLOPs/s) сначала растёт линейно, пока не достигнет критической арифметической интенсивности железа — у TPU v5e это 240. Любой алгоритм с меньшей интенсивностью ограничен пропускной способностью (BW) и пиковой пропускной способностью памяти (красная область). Любой алгоритм правее полностью использует наши FLOPs (зелёная область). Здесь Algo 1 упирается в связь и использует лишь долю всех FLOPs/s железа. Algo 2 упирается в вычисления. В общем случае производительность алгоритма можно повысить, либо увеличив его арифметическую интенсивность, либо увеличив доступную пропускную способность памяти (переход от BW1 к BW2).

flowchart LR
A["Интенсивность алгоритма I"] --> C{"I больше пика железа?"}
C -->|"нет"| M["Упор в память: скорость = I × пропускная способность"]
C -->|"да"| F["Упор в вычисления: скорость = пиковые FLOPs/s"]

Roofline одним правилом: достижимые FLOPs/s = min(пик, I × пропускная способность) (схема переводчика)

Умножение матриц

Рассмотрим алгоритм, который скоро станет нашим любимым: умножение матриц (matmul). Даны две матрицы $X$ и $Y$, пишем $X * Y \rightarrow Z$, где $X$ имеет форму $\text{bf16}[B, D]$, $Y$ — $\text{bf16}[D, F]$, а $Z$ — $\text{bf16}[B, F]$. Чтобы выполнить умножение, нужно загрузить $2DF + 2BD$ байт, выполнить $2BDF$ FLOPs и записать обратно $2BF$ байт⁹¹⁰. Значит:

$$\text{интенсивность}(\text{matmul}) = \frac{2BDF}{2BD + 2DF + 2BF} = \frac{BDF}{BD + DF + BF}$$

Формула хорошо упрощается, если считать, что «размер батча» $B$ мал по сравнению с $D$ и $F$. Тогда

$$\frac{BDF}{BD + DF + BF} \approx \frac{BDF}{DF} = B$$ $$\text{интенсивность}(\text{matmul}) > \text{интенсивность}(\text{TPU}) \implies B > \frac{1{,}97 \cdot 10^{14}}{8{,}20 \cdot 10^{11}} = 240$$

Для матмулов трансформера это разумное предположение: локальный (на одну реплику) батч в токенах обычно $B$ меньше 1024 (заметьте, это токены, а не последовательности), а $D$ и $F$ больше 8000. Поэтому в общем случае мы упираемся в вычисления, когда батч на реплику¹¹ больше 240 токенов, — очень простое правило!

Вывод

Чтобы умножение матриц в bfloat16 на большинстве TPU упиралось в вычисления, размер батча должен быть больше 240 (токенов)¹².

У этого правила есть несколько заметных оговорок, которые мы разберём в задачах ниже, особенно связанных с квантованием (например, если активации квантованы, а FLOPs выполняются в полной точности), но запомнить его стоит. Для GPU это число немного выше (ближе к 300), но вывод в целом тот же. Когда мы разбиваем большое умножение матриц на маленькие, важны и размеры тайлов¹³. Детали GPU и TPU нижнего уровня мы обсудим в следующей главе.

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

У RTX 4080 пик в BF16 — 97,5 TFLOPS, пропускная способность памяти — 716,8 ГБ/с. Критическая интенсивность: $97{,}5 \cdot 10^{12} / 716{,}8 \cdot 10^{9} \approx 136$ FLOPs/байт. Значит, матмул в bf16 на вашей карте упирается в вычисления уже при батче больше ~136 токенов. При генерации текста по одному токену с батчем 1 интенсивность ≈ 1 — карта использует меньше 1% своих FLOPs, и скорость генерации определяет только то, как быстро читаются веса. Проверить можно так (запуск — секунды):

import time, torch

D = F = 8192
W = torch.randn(D, F, device="cuda", dtype=torch.bfloat16)
for B in [1, 16, 64, 128, 256, 512, 1024, 4096]:
    X = torch.randn(B, D, device="cuda", dtype=torch.bfloat16)
    for _ in range(3):                       # прогрев
        X @ W
    torch.cuda.synchronize(); t = time.perf_counter()
    for _ in range(20):
        X @ W
    torch.cuda.synchronize(); dt = (time.perf_counter() - t) / 20
    print(f"B={B:5d}  {2 * B * D * F / dt / 1e12:6.1f} TFLOPS")   # растёт ~линейно, потом плато

Roofline для сетевой связи

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

Возьмём несколько искусственный пример: мы хотим перемножить две большие матрицы $X\sim \text{bf16}[B, D]$ и $Y \sim \text{bf16}[D, F]$, поровну разбитые между 2 TPU/GPU (по измерению $D$). Чтобы выполнить это умножение (как мы увидим в главе 3), можно перемножить половины каждой матрицы на каждом TPU (Z0 = X[:, :D // 2] @ Y[:D // 2, :] на TPU 0 и Z1 = X[:, D // 2:] @ Y[D // 2:, :] на TPU 1), а затем скопировать получившиеся «частичные суммы» на другой TPU и сложить их. Допустим, мы можем копировать 4.5e10 байт/с в каждую сторону и выполнять 1.97e14 FLOPs/s на каждом чипе. Чему равны $T_\text{math}$ и $T_\text{comms}$?

$T_\text{math}$, очевидно, вдвое меньше, чем раньше: каждый TPU делает половину работы¹⁴, то есть

$$T_\text{math} = \frac{2BDF}{2 \cdot \text{FLOPs/s ускорителя}} = \frac{BDF}{1{,}97 \cdot 10^{14}}$$

А $T_\text{comms}$? Теперь это время связи между чипами! Это просто общее число переданных байт, делённое на пропускную способность сети, то есть

$$T_\text{comms} = \frac{2BF}{\text{пропускная способность сети}} = \frac{2BF}{4{,}5 \cdot 10^{10}}$$

Поэтому мы упираемся в вычисления (теперь относительно сети между чипами), когда $\text{интенсивность}(\text{matmul на 2 чипах}) > \text{интенсивность}(\text{TPU относительно сети})$, или, что то же самое, когда $\frac{BDF}{2BF} = \frac{D}{2} > \frac{1{,}97 \cdot 10^{14}}{4{,}5 \cdot 10^{10}} = 4377$, то есть $D > 8755$. Заметьте: в отличие от прошлого случая, критический порог теперь зависит от $D$, а не от $B$! Подумайте, почему. Это лишь один пример, но подчеркнём: roofline такого рода критически важен, чтобы понимать, когда операцию можно распараллелить на несколько TPU.

Почему от D (примечание переводчика)

Каждый чип считает $2B \cdot (D/2) \cdot F$ операций, а передаёт только результат — матрицу $B \times F$. Вычислений пропорционально $D$, а передача от $D$ не зависит. Чем «толще» сокращаемое измерение, тем больше работы приходится на каждый переданный байт. На вашей машине карты связаны только PCIe 4.0: теоретически до 32 ГБ/с, а без прямого обмена P2P (у GeForce он обычно запрещён) на практике порядка 10 ГБ/с — эту оценку использует урок «Четыре видеокарты: DDP и FSDP на практике». Порог для RTX 4080: $D/2 > 97{,}5 \cdot 10^{12} / 10^{10} \approx 9750$, то есть $D$ больше ~19 500. У моделей до 70B $D$ не больше 8192, так что такое разбиение одного матмула на две карты через PCIe всегда упрётся в связь.

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

Вопрос 1 [матмул в int8]. Допустим, мы хотим выполнить умножение $X[B, D] \cdot_D Y[D, F] \rightarrow Z[B, F]$¹⁵ в точности int8 (1 байт на параметр) вместо bfloat16 (2 байта на параметр), потому что TPU/GPU умеют умножать матрицы в пониженной точности быстрее.

1. Сколько байт нужно загрузить из памяти? Сколько записать обратно?
2. Сколько всего операций (OPs) выполняется?
3. Какова арифметическая интенсивность?
4. Какова roofline-оценка $T_\text{math}$ и $T_\text{comms}$? Каковы разумные верхняя и нижняя оценки времени всей операции?

Считайте, что пропускная способность HBM — 8.2e11 байт/с, а пиковая скорость int8 — 3.94e14 OPs/s (примерно вдвое больше, чем в bfloat16).

Ответ

1. Параметры хранятся в int8, по 1 байту на параметр, поэтому из HBM загружается $BD + DF$ байт, а записывается $BF$.
2. Столько же, сколько в bfloat16, но в теории операции int8 в секунду выполняются быстрее. Так что это по-прежнему $2BDF$ OPs.
3. Арифметическая интенсивность — $2BDF / (BD + DF + BF)$. Если, как выше, считать $B \ll D$ и $B \ll F$, получаем арифметическую интенсивность $2B$, то есть правило превращается в $B > \text{интенсивность HBM для int8} / 2$. По данным числам интенсивность для int8 — 3.94e14 / 8.2e11 = 480, так что правило: $B > 480 / 2 = 240$. Заметьте: почти ничего не изменилось!
4. $T_\text{math} = 2BDF / 3{,}94 \cdot 10^{14}$ и $T_\text{comms} = (BD + DF + BF) / 8{,}2 \cdot 10^{11}$, так что разумная нижняя оценка — $\max(T_\text{math}, T_\text{comms})$, а верхняя — $T_\text{math} + T_\text{comms}$.

Вопрос 2 [матмул int8 + bf16]. На практике веса и активации часто квантуют по-разному: веса храним в очень низкой точности, а активации (и вычисления) — в более высокой. Допустим, мы хотим квантовать веса в int8, а активации (и вычисления) оставить в bfloat16. При каком размере батча мы упрёмся в вычисления? Считайте, что скорость — 1.97e14 FLOPs/s в bfloat16.

Подсказка: речь именно о bf16[B, D] * int8[D, F] -> bf16[B, F], где $B$ — «размер батча».

Ответ

Снова считая $B$ малым, имеем $2BDF$ FLOPs в bfloat16, но лишь $DF$ байт весов (вместо $2DF$ в bfloat16). Значит, мы упираемся в вычисления при $2B > 240$, то есть $B > 120$. Это намного меньше: если квантовать веса в int8 (что довольно легко), но считать FLOPs в bfloat16, эффективность заметно выигрывает (хотя операции в int8 были бы ещё лучше).

Вопрос 3. В условиях вопроса 2 постройте roofline-график пиковых FLOPs/s в зависимости от $B$ для $F = D = 4096$ и $F = D = 1024$. Используйте точное число загружаемых байт, а не приближение.

Ответ

Нужный график — roofline-plot-q3.png в оригинале: две кривые растут с ростом батча и выходят на одно и то же плато пиковых FLOPs/s.

Обе модели в итоге выходят на пиковые FLOPs/s железа, но при больших D/F это происходит раньше. При D=F=1024 критический размер батча почти вдвое больше. Код, который строит этот график:

import matplotlib.pyplot as plt
import numpy as np

bs = np.arange(1, 512)

def roofline(B, D, F):
  total_flops = 2*B*D*F
  flops_time = total_flops / 1.97e14
  comms_time = (2*B*D + D*F + 2*B*F) / 8.2e11   # активации bf16 (2 байта), веса int8 (1 байт)
  total_time = np.maximum(flops_time, comms_time)
  return total_flops / total_time

roofline_big = roofline(bs, 4096, 4096)
roofline_small = roofline(bs, 1024, 1024)

plt.figure(figsize=(8, 4))
plt.plot(bs, roofline_big, label='F=D=4096')
plt.plot(bs, roofline_small, label='F=D=1024')
plt.legend()
plt.xlabel('batch size')
plt.ylabel('peak bfloat16 FLOPs/s on TPU v5e')
plt.grid()

Вопрос 4. А если мы хотим выполнить $\text{int8}[B, D] \cdot_D \text{int8}[B, D, F] \rightarrow \text{int8}[B, F]$, представляя, что у каждого элемента батча своя матрица? Какова арифметическая интенсивность этой операции?

Ответ

Начнём с общего числа FLOPs и объёма связи.

1. Всего FLOPs: почти столько же, ведь мы выполняем $B$ независимых произведений $[D] \times [D, F]$ — та же общая работа, что и одно умножение $[B, D] \times [D, F]$ (подробнее — в главе 4). Так что это просто $2BDF$.
2. Всего связи: здесь её намного больше — $BD + BDF + BF$.
3. Поэтому арифметическая интенсивность теперь $2BDF / (BD + BDF + BF)$. Поскольку в знаменателе доминирует $BDF$, это примерно $2$. То есть вместо зависимости от размера батча она, по сути, постоянна. Это плохо: значит, мы почти всегда будем упираться в связь, что бы ни делали.

Вопрос 5 [roofline памяти для GPU]. По спецификации NVIDIA для H100 SXM вычислите размер батча, при котором умножение матриц в bfloat16 начнёт упираться в вычисления. Учтите, что числа FLOPs для Tensor Core там вдвое больше истинных: они достижимы только со структурированной разреженностью.

Ответ

В спецификации указано 1.979e15 FLOPs/s в bfloat16 со звёздочкой «с разреженностью». Без разреженности истинное значение вдвое меньше: 9.89e14 FLOPs/s. Пропускная способность памяти — 3,35 ТБ/с, или 3.35e12 байт/с. Значит, $B_\text{crit}$ = 9.89e14 / 3.35e12 = 295 — довольно близко к TPU.

Это всё для главы 1. Глава 2 о том, как настоящие TPU обращаются с FLOPs и связью, — здесь.

Словарик

English	Русский
roofline	roofline, «потолок» производительности
FLOPs / FLOPs/s	операции с плавающей точкой / операций в секунду
HBM (high bandwidth memory), HBM bandwidth	память HBM, пропускная способность HBM
ICI, DCN, PCIe	межчиповые связи (inter-chip interconnect), сеть дата-центра, шина PCIe
compute-bound / communication-bound, memory-bound	упор в вычисления / в связь, в память
arithmetic (operational) intensity	арифметическая интенсивность
peak / critical arithmetic intensity	пиковая / критическая арифметическая интенсивность
matmul	умножение матриц
per-replica batch size	размер батча на реплику
tile	тайл (блок матрицы)
sharded matrix, partial sums	шардированная (разбитая) матрица, частичные суммы
quantization, int8	квантование, 8-битные целые
structured sparsity	структурированная разреженность

Footnotes

1. bf16 — сокращение от bfloat16, 16-битного формата с плавающей точкой, часто используемого в ML. ↩

2. H100 и B200 обычно достигают лишь 80–85% заявленного пика FLOPs, а TPU при обычном использовании — ближе к 95%. ↩

3. Учтите, что эти чипы стоят по-разному, и сравнение не нормировано на стоимость. ↩

4. NVIDIA называет её также «memory bandwidth» — пропускной способностью памяти. ↩

5. В книге «упор в связь», «упор в память» и «упор в пропускную способность» (communication-bound, comms-bound, memory-bound, bandwidth-bound) используются как синонимы. ↩

6. MXU — блок умножения матриц в TPU. Мы уточняем это, потому что в TPU есть и другие вычислительные блоки, например VPU, отвечающий за поэлементные операции, — у него другие пиковые FLOPs/s. ↩

7. Это верно, только если алгоритм загружает веса из HBM и выполняется на MXU. Как мы обсудим в следующей главе, иногда параметры можно хранить в VMEM, у которой гораздо больше пропускная способность. Многие алгоритмы также выполняются на VPU с другими характеристиками производительности. ↩

8. Число 240 выше здесь не совсем подходит для сравнения: как вы увидите в следующей главе, скалярное произведение выполняется на VPU, а не на MXU. VPU в TPU v5p выполняет примерно $7 \cdot 10^{12}$ FLOPs в секунду на ядро, так что его критическая интенсивность около 3, и мы всё равно в некоторой степени упираемся в связь. В любом случае интенсивность низкая и постоянная, а значит, на большинстве железа упереться в вычисления трудно. ↩

9. Строго говоря, выполняется $BF \times (2D - 1)$ FLOPs, но это достаточно близко: $BDF$ умножений и $BF \cdot (D-1)$ сложений. Подробнее — в главе 4. ↩

10. Хотя результат матмула формально в float32, перед копированием в HBM его обычно приводят к bfloat16. ↩

11. «На реплику» — потому что, если мы как-то шардируем модель, чтобы увеличить число чипов в матмуле, мы одинаково масштабируем и доступные вычисления, и пропускную способность памяти. Поэтому критический размер батча верен для каждой независимой копии весов модели. ↩

12. Заметьте: это не размер батча в обычном смысле — в последовательностях. Оказывается, большинство roofline зависят только от числа токенов, принадлежат ли они одной последовательности или разным. Например, если у вас батч из 512 последовательностей по 4096 токенов на 2048 GPU, то общий батч — 512 * 4096 = 2M токенов, а локальный батч — 1 тысяча токенов. ↩

13. Большое умножение матриц приходится разбивать на маленькие тайлы, помещающиеся в VMEM/SMEM/TMEM — более быструю память на кристалле. Из-за этого блоки загружаются по нескольку раз, и уже не совсем верно, что мы загружаем лишь $O(N^2)$ байт. Рассмотрим умножение $(m, k) \cdot (k, n)$ с размерами тайлов $bm$, $bk$, $bn$. Пусть $tm = m / bm$ и т. д. Тогда всего FLOPs $2 \cdot tm \cdot tn \cdot tk \cdot bm \cdot bn \cdot bk$, а всего байт $2 \cdot tm \cdot tn \cdot (tk \cdot (bm \cdot bk + bk \cdot bn) + bm \cdot bn)$. Отбросив последнее слагаемое, получаем интенсивность $bm \cdot bn / (bm + bn)$, похожую на формулу выше. ↩

14. Мы пренебрегаем FLOPs на сложение двух частичных сумм (ещё $BF$ сложений) — это практически ничто. ↩

15. Здесь и далее запись $A \cdot_D B$ означает умножение со свёрткой (суммированием) по измерению D. Это вольное обращение с нотацией einsum. ↩