How To Scale Your Model, глава 3: шардированные матрицы и как их умножать

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

О переводе

Перевод главы 3 книги How To Scale Your Model (лицензия MIT, © 2025 Google LLC): весь основной текст и задачи 1–4 с ответами; условия задач 5–10 перечислены в конце, их ответы в оригинале. Анимации описаны словами. Предыдущая глава — «Как устроены TPU», следующая — «Вся математика трансформера».

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

Глава даёт язык, на котором дальше разбирается любой параллелизм: «шардировать по оси», «AllGather», «ReduceScatter», «AllReduce». В уроке «Распределённое обучение» те же операции появляются как all_reduce и reduce_scatter в torch.distributed, а в уроке «Четыре видеокарты: DDP и FSDP на практике» вы их запускаете. AllToAll понадобится для смеси экспертов (урок «Смесь экспертов»). Код в главе — на JAX; соответствия PyTorch даны в конце.

Нотация разбиения и коллективные операции

Когда мы обучаем LLM на десяти тысячах TPU или GPU, мы абстрактно выполняем то же вычисление, что и при обучении на одном. Разница в том, что наши массивы не помещаются в HBM одного TPU/GPU, поэтому их приходится делить¹. Это называют «шардированием» (sharding) или «разбиением» (partitioning) массивов. Искусство масштабирования — понять, как шардировать модель, чтобы вычисления оставались эффективными.

Рисунок в оригинале (sharding-example.png): массив формы A[I, J] шардирован между 4 устройствами. Оба измерения поровну разделены между 2 устройствами с шардированием A[I_X, J_Y]. Каждый TPU хранит 1/4 всей памяти.

Заметьте: у шардированного массива по-прежнему та же глобальная, или логическая форма, что и у нешардированного, например (4, 128), но есть и локальная форма на устройстве, например (2, 64), — она показывает реальный размер в байтах, который хранит каждый TPU (на рисунке выше каждый TPU хранит ¼ всего массива). Обобщим это на произвольные массивы.

Единая нотация шардирования

Чтобы описать, как тензор разделён на блоки между устройствами, мы используем вариант нотации с именованными осями: предполагаем, что есть двумерная или трёхмерная сетка устройств, называемая сеткой устройств (device mesh), где каждой оси дано имя оси сетки, например X, Y и Z. Тогда расположение данных матрицы по сетке устройств можно задать, описав, как каждое именованное измерение массива разделено по физическим осям сетки. Такое назначение мы называем шардированием (sharding).

Пример (рисунок выше). Для рисунка выше:

• Сетка: Mesh(devices=((0, 1), (2, 3)), axis_names=('X', 'Y')) — 4 TPU в сетке 2×2 с именами осей $X$ и $Y$.
• Шардирование: $A[I_X, J_Y]$ — первую ось, $I$, шардировать по оси сетки $X$, а вторую, $J$, — по оси сетки $Y$. Значит, каждый шард хранит $1 / (\lvert X\rvert \cdot \lvert Y\rvert)$ массива.

Вместе это говорит нам, что локальная форма массива (размер шарда, который хранит отдельное устройство) — $(\lvert I\rvert / 2, \lvert J\rvert / 2)$, где $\lvert I\rvert$ — размер первого измерения A, а $\lvert J\rvert$ — второго.

Быстрый вопрос [двумерное шардирование по одной оси массива]. Рассмотрим массив fp32[1024, 4096] с шардированием $A[I_{XY}, J]$ и сеткой {'X': 8, 'Y': 2}. Сколько данных хранит каждое устройство? Сколько времени займёт загрузка этого массива из HBM на H100 (пропускная способность памяти 3.4e12 на чип)?

Ответ

$A[I_{XY}, J]$ шардирует первое измерение (I) сразу по аппаратным осям X и Y. Здесь локальная форма — $(\lvert I\rvert /(\lvert X\rvert \cdot \lvert Y\rvert), \lvert J\rvert)$. Глобальная форма — fp32[1024, 4096], значит, локальная — fp32[64, 4096].

У каждого GPU 4 * 64 * 4096 = 1MiB байт, так что загрузка займёт около 1e6 / 3.4e12 = 294ns, хотя, вероятно, заметно дольше из-за разных накладных расходов — массив очень мал.

Как это выглядит. Представим себе эти шардирования на примере двумерного массива, разделённого между 4 устройствами (рисунки sharding-colored1–6 в оригинале, где каждое устройство закрашено своим цветом):

• Полностью реплицированную (размноженную) форму матрицы пишем просто $A[I, J]$, без назначения шардирования. Это значит, что каждое устройство хранит полную копию всей матрицы.
• То, что одно из измерений разделено по оси сетки, обозначаем нижним индексом — осью сетки. Например, $A[I_X, J]$ значит, что логическая ось I разделена по измерению сетки X, а измерение J не разделено, и блоки остаются частично реплицированными по оси сетки Y.
• $A[I_X, J_Y]$ значит, что логическая ось I разделена по оси сетки X, а измерение J — по оси сетки Y.
• $A[I_{XY}, J]$ значит, что оси сетки X и Y мы рассматриваем как одно большое «сплющенное» измерение и делим именованную ось I между всеми устройствами. Порядок нескольких осей сетки в индексе важен: он задаёт порядок обхода сетки при разбиении.

Наконец, нельзя шардировать несколько именованных осей по одному и тому же измерению сетки. Например, $A[I_X, J_X]$ — бессмысленное, запрещённое шардирование. Как только измерение сетки использовано для шардирования одного измерения массива, оно в каком-то смысле «израсходовано».

Быстрый вопрос. Пусть A — массив формы int8[128, 2048] с шардированием $A[I_{XY}, J]$ на сетке Mesh({'X': 2, 'Y': 8, 'Z': 2}) (всего 32 устройства). Сколько памяти A занимает на устройстве? Сколько всего на всех устройствах?

Ответ

Массив A шардирован по X и Y и реплицирован по Z, так что на устройстве он имеет форму int8[128 / (2 * 8), 2048] = int8[8, 2048] и размер 8 * 2048 = 16,384 байт. Поскольку он реплицирован по Z, а внутри каждой Z-плоскости полностью шардирован по X и Y, получается 2 полные копии исходного массива (по одной на Z-плоскость). Значит, всего на всех устройствах: размер исходного массива × число реплик по Z = 128 × 2048 × 2 = 512 КиБ. Проверка: 32 устройства × 16 384 байт = 512 КиБ.

Как это описать в коде?

До сих пор мы обходили стороной код, но сейчас хороший момент заглянуть вперёд. JAX использует синтаксис именованного шардирования, очень близкий к абстрактному синтаксису выше. Подробнее мы поговорим об этом в главе 10, а пока — краткий анонс. Поиграть с этим можно в Google Colab и, профилируя результат, посмотреть, как JAX обрабатывает разные шардирования. Этот фрагмент делает 3 вещи:

1. Создаёт jax.Mesh, который раскладывает наши 8 TPU в сетку 4×2 с именами осей ‘X’ и ‘Y’.
2. Создаёт матрицы A и B, где A шардирована по обоим измерениям, а B — по выходному.
3. Компилирует и выполняет простое умножение матриц, которое возвращает шардированный массив.

import jax
import jax.numpy as jnp

# Создаём сетку! Работаем на срезе TPU v2-8 4x2 с именами осей 'X' и 'Y'.
# Тип оси Auto говорит JAX позволить компилятору XLA самому выводить промежуточные шардирования.
assert len(jax.devices()) == 8
Auto = jax.sharding.AxisType.Auto
mesh = jax.make_mesh(axis_sizes=(4, 2), axis_names=('X', 'Y'), axis_types=(Auto, Auto))

# Небольшая вспомогательная функция для задания шардирования. PartitionSpec — это наше
# шардирование (отображение осей массива в имена осей сетки).
def P(*args):
  return jax.NamedSharding(mesh, jax.sharding.PartitionSpec(*args))

# Шардируем и A, и B по несворачиваемому измерению, а A — ещё и по сворачиваемому.
A = jnp.zeros((8, 2048), dtype=jnp.bfloat16, device=P('X', 'Y'))
B = jnp.zeros((2048, 8192), dtype=jnp.bfloat16, device=P(None, 'Y'))

# На этих шардированных массивах можно выполнить умножение! out_shardings говорит, как
# шардировать результат. Остальное шардирование JAX/XLA сделают за нас.
y = jax.jit(lambda A, B: jnp.einsum('BD,DF->BF', A, B), out_shardings=P('X', 'Y'))(A, B)

Самое приятное в JAX — эти массивы ведут себя так, будто они не шардированы! B.shape покажет глобальную, или логическую, форму (2048, 8192). Чтобы увидеть локальное шардирование, нужно посмотреть на B.addressable_shards. С этими массивами можно выполнять операции, и JAX попытается понять, как их транслировать (broadcast) или менять форму, чтобы выполнить операции. Например, в примере выше локальная форма A — [2, 1024], а B — [2048, 4096]. JAX/XLA автоматически добавят между этими массивами связь, необходимую для итогового умножения.

🖥 Запуск без TPU (примечание переводчика)

Код можно запустить на обычном процессоре: перед импортом JAX задайте восемь «фальшивых» устройств — import os; os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=8" (или, в новых версиях JAX, jax.config.update('jax_num_cpu_devices', 8), как советует глава 10). Скорость будет ни о чём не говорить, но шардирования и addressable_shards вы увидите те же.

Вычисления с шардированными массивами

Если у вас массив данных, распределённый между многими устройствами, и вы хотите выполнять над ним математические операции, каковы накладные расходы на шардирование и данных, и вычислений?

Очевидно, это зависит от вычисления.

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

Остаток раздела — о том, как умножать шардированные матрицы. В первом приближении это значит перемещать куски матриц так, чтобы каждый кусок можно было полностью перемножить или просуммировать. Каждое шардирование потребует своей связи. Например, $A[I_X, J] \cdot B[J, K_Y] \to C[I_X, K_Y]$ можно перемножить вообще без связи, потому что сворачиваемое измерение (contracting dimension; J — то, по которому мы на самом деле суммируем) не шардировано. Но если бы мы хотели получить нешардированный результат ($A[I_X, J] \cdot B[J, K_Y] \to C[I, K]$), пришлось бы скопировать на каждое устройство либо $A$ и $B$, либо $C$ (с помощью AllGather). У этих двух вариантов разная стоимость связи, так что её нужно посчитать и выбрать наименьшую.

Это удобно понимать через «блочное умножение матриц»

Здесь полезно вспомнить понятие «блочной матрицы» — матрицы, составленной из матриц:

$$\begin{pmatrix} a_{00} & a_{01} & a_{02} & a_{03} \\ a_{10} & a_{11} & a_{12} & a_{13} \\ a_{20} & a_{21} & a_{22} & a_{23} \\ a_{30} & a_{31} & a_{32} & a_{33} \end{pmatrix} = \begin{pmatrix} \mathbf{A_{00}} & \mathbf{A_{01}} \\ \mathbf{A_{10}} & \mathbf{A_{11}} \end{pmatrix}$$

где каждый $\mathbf{A_{ij}}$ — блок 2×2. У умножения матриц есть приятное свойство: если сомножители записаны через блоки, произведение можно записать через блочные умножения по обычному правилу:

$$\begin{pmatrix} A_{00} & A_{01} \\ A_{10} & A_{11} \end{pmatrix} \cdot \begin{pmatrix} B_{00} & B_{01} \\ B_{10} & B_{11} \end{pmatrix} = \begin{pmatrix} A_{00}B_{00} + A_{01}B_{10} & A_{00}B_{01} + A_{01}B_{11} \\ A_{10}B_{00} + A_{11}B_{10} & A_{10}B_{01} + A_{11}B_{11} \end{pmatrix}$$

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

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

1. Случай 1: ни один вход не шардирован по сворачиваемому измерению. Можно перемножить локальные шарды без всякой связи.
2. Случай 2: у одного входа сворачиваемое измерение шардировано. Обычно делают «AllGather» шардированного входа по сворачиваемому измерению.
3. Случай 3: оба входа шардированы по сворачиваемому измерению. Можно перемножить локальные шарды, а затем сделать «AllReduce» результата.
4. Случай 4: у обоих входов несворачиваемое измерение шардировано по одной и той же оси. Двигаться дальше нельзя, пока не сделан AllGather одного из входов.

Это можно считать просто правилами, которым нужно следовать, но полезно понимать, почему они верны и сколько стоят. Разберём каждое подробно.

flowchart TD
S["Умножаем A · B, сворачивая по J"] --> Q1{"J шардировано?"}
Q1 -->|"ни у кого"| Q4{"Несворачиваемые оси шардированы по одной оси сетки?"}
Q4 -->|"нет"| C1["Случай 1: связь не нужна"]
Q4 -->|"да"| C4["Случай 4: сначала AllGather одного входа"]
Q1 -->|"у одного"| C2["Случай 2: AllGather этого входа, затем локальный матмул"]
Q1 -->|"у обоих по одной оси"| C3["Случай 3: локальный матмул, затем AllReduce или ReduceScatter"]

Четыре случая умножения шардированных матриц (схема переводчика)

Случай 1: ни у одного сомножителя сворачиваемое измерение не шардировано

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

$$\mathbf{A}[I_X, J] \cdot \mathbf{B}[J, K_Y] \rightarrow \mathbf{C}[I_X, K_Y]$$

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

$$\begin{aligned} \mathbf{A}[I, J] \cdot \mathbf{B}[J, K] &\rightarrow \mathbf{C}[I, K] \\ \mathbf{A}[I_X, J] \cdot \mathbf{B}[J, K] &\rightarrow \mathbf{C}[I_X, K]\\ \mathbf{A}[I, J] \cdot \mathbf{B}[J, K_Y] &\rightarrow \mathbf{C}[I, K_Y]\\ \mathbf{A}[I_X, J] \cdot \mathbf{B}[J, K_Y] &\rightarrow \mathbf{C}[I_X, K_Y] \end{aligned}$$

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

Случай 2: у одного сомножителя сворачиваемое измерение шардировано

Посмотрим, что делать, когда один вход A шардирован по сворачиваемому измерению J, а B полностью реплицирован:

$$\mathbf{A}[I, J_X] \cdot \mathbf{B}[J, K] \rightarrow \mathbf{C}[I, K]$$

Просто перемножить локальные куски A и B нельзя: нужно просуммировать по всему сворачиваемому измерению A, а оно разделено по оси X. Обычно сначала делают «AllGather» шардов A, чтобы у каждого устройства была полная копия, и только потом умножают на B:

$$\textbf{AllGather}_X[I, J_X] \rightarrow \mathbf{A}[I, J]$$ $$\mathbf{A}[I, J] \cdot \mathbf{B}[J, K] \rightarrow \mathbf{C}[I, K]$$

Так само умножение целиком выполняется на каждом устройстве.

Вывод

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

Заметьте: если B не шардирована по X, можно также выполнить локальное частичное умножение, а затем просуммировать (сделать AllReduce) шардированные частичные суммы — так вычисления тоже шардируются, но связь обычно дороже. Иногда это быстрее, хотя на практике B обычно тоже шардирована. Когда это лучше, разбирает вопрос 4 ниже.

Что такое AllGather? AllGather — первый базовый коммуникационный примитив MPI, который мы обсудим. AllGather снимает шардирование по оси и собирает шарды, разбросанные по устройствам, на каждом устройстве вдоль этой оси. В нашей нотации AllGather убирает нижний индекс у набора осей, например

$$\textbf{AllGather}_{XY}(A[I_{XY}, J]) \rightarrow A[I, J]$$

Убирать все индексы данного измерения не обязательно: например, $A[I_{XY}, J] \rightarrow A[I_Y, J]$ — тоже AllGather, только по одной оси. Также AllGather может понадобиться, чтобы снять шардирование несворачиваемого измерения, например в умножении

$$A[I_X, J] \cdot B[J, K] \rightarrow C[I, K]$$

Можно сначала сделать AllGather A, сняв шардирование входа, а можно выполнить шардированное умножение и затем сделать AllGather результата C.

Как на самом деле выполняется AllGather? Чтобы выполнить одномерный AllGather вдоль одной оси TPU (кольца), каждый TPU, по сути, передаёт свой шард по кольцу, пока копия не окажется у каждого устройства².

Анимация в оригинале (all-gather.gif): AllGather на 8 устройствах TPU или GPU. Каждое устройство начинает с 1/8 массива и заканчивает полной копией.

AllGather можно делать в одном направлении или в обоих (в анимации — оба). При одном направлении каждый TPU отправляет куски размером $\text{bytes} / N$ за $N - 1$ переходов по кольцу. При двух — $\lfloor \frac{N}{2} \rfloor$ переходов с кусками размером $2 \cdot \text{bytes} / N$.

Сколько это длится? Возьмём двунаправленный AllGather и посчитаем его время. Пусть $V$ — число байт в массиве, а $X$ — число шардов по сворачиваемому измерению. Тогда, как видно из анимации, на каждом переходе в каждую сторону отправляется $V / \lvert X\rvert$ байт, так что каждый переход занимает

$$T_{hop} = \frac{2 \cdot V}{\lvert X \rvert \cdot W_\text{ici}}$$

где $W_\text{ici}$ — двусторонняя пропускная способность ICI³. Чтобы дойти до каждого TPU, нужно $\lvert X\rvert / 2$ переходов⁴, так что вся операция занимает

$$T_{total} = \frac{2 \cdot V \cdot X}{2 \cdot X \cdot W_\text{ici}} = \frac{V}{W_\text{ici}}$$

Заметьте: это не зависит от $X$! Это поразительно: хотя наши TPU связаны только локально, локальность связей роли не играет. Узкое место — лишь скорость каждого канала.

Вывод

Когда AllGather (или ReduceScatter, или AllReduce) выполняется в режиме, ограниченном пропускной способностью, реальное время связи зависит только от размера массива и доступной пропускной способности, а не от числа устройств, по которым шардирован массив!

О задержке ICI. У каждого перехода по каналу ICI есть собственные накладные расходы, не зависящие от объёма данных. Обычно это около 1 мкс. Значит, когда массив $A$ очень мал и каждый переход занимает меньше 1 мкс, можно попасть в режим, ограниченный задержкой (latency-bound), где время зависит от $X$.

Подробности

Пусть $T_\text{min}$ — минимальное время одного перехода. Тогда

$$T_{hop} = \max \left[ T_{min}, \frac{2 \cdot V}{X \cdot W_\text{ici}} \right]$$$$T_{total} = \max \left[ \frac{T_{min} \cdot X}{2}, \frac{V}{W_\text{ici}} \right]$$

поскольку переходов $X / 2$. Для больших редукций или сборов мы надёжно ограничены пропускной способностью: данных так много, что накладные расходы на переход практически не важны. Но для маленьких массивов (например, при сэмплировании из модели) они уже не пренебрежимы, а пропускная способность ICI не важна: мы ограничены только задержкой. Иначе говоря, для конкретного TPU, например TPU v5e с односторонней пропускной способностью ICI 4.5e10, отправка любого буфера меньше 4.5e10 * 1e-6 = 45kB будет ограничена задержкой.

Замер пропускной способности AllGather на срезе TPU v5e 8x16 (all-gather-bandwidth.png в оригинале; массив шардирован по оси 16, так что есть полное двунаправленное кольцо) показывает: достигается около 95% заявленной пиковой пропускной способности (4.5e10), причём уже при массиве около 10 МБ, что при 16-путевом шардировании даёт примерно 625 КБ на устройство (к слову, это гораздо лучше, чем у GPU).

Что, если делать AllGather по нескольким осям? Тогда у нас несколько измерений ICI, по которым можно вести сбор. Например, AllGather_XY([B, D_XY]) работает по двум аппаратным осям сетки. Это увеличивает доступную пропускную способность в $N_\text{axes}$ раз.

С учётом задержки получаем общее правило:

$$T_{total} = \max \left[ \frac{T_{min} \cdot \sum_{i} \lvert X_i \rvert}{2}, \frac{V}{W_\text{ici} \cdot N_\text{axes}} \right]$$

где $\sum_i \lvert X_i \rvert / 2$ — длина самого длинного пути в сетке TPU.

Быстрый вопрос 2 [время AllGather]. По числам из главы 2 — сколько займёт AllGather_Y([E_Y, F]) → [E, F] на TPU v5e с двумерной сеткой {'X': 8, 'Y': 4}, $E = 2048$, $F = 8192$ в bfloat16? А при $E=256, F=256$?

Ответ

Начнём с базовых величин:

1. У TPU v5e 4,5e10 байт/с односторонней пропускной способности ICI по каждой из 2 осей.
2. В bfloat16 для случая (1) имеем $A[E_Y, F]$, так что каждое устройство хранит массив формы bf16[512, 8192] размером 512 × 8192 × 2 = 8,4 МБ. Весь массив — 2048 × 8192 × 2 = 34 МБ.

Для случая (1) можно применить формулу выше. Сбор идёт по одной оси, так что $T_{\text{comms}} = 34 \cdot 10^{6} / 9 \cdot 10^{10} = 377$ мкс. Проверим, что мы не упираемся в задержку: по оси размера 4 будет не больше 3 переходов, так что граница по задержке — около 3 мкс, нам до неё далеко. Однако у TPU v5e замыкающий канал есть, только когда ось имеет размер 16, так что здесь полностью двунаправленный AllGather на самом деле невозможен. Данным с краёв нужно 3 перехода, чтобы дойти до другого края, так что теоретически получается скорее $T_{\text{comms}} = 3 \cdot 8{,}4 \cdot 10^{6} / 4{,}5 \cdot 10^{10} = 560\,\mu s$. Вот реальный профиль из этого Colab: $680\,\mu s$ — разумно, ведь мы, скорее всего, не получаем 100% теоретической пропускной способности! Для случая (2) каждый шард имеет размер 64 * 256 * 2 = 32kB. 32e3 / 4.5e10 = 0.7us, так что мы упираемся в задержку. Переходов 3, так что это займёт примерно 3 × 1 мкс = 3 мкс. На практике — ближе к 8 мкс.

Замечание

Когда у нас двумерная сетка вроде {'X': 16, 'Y': 4}, не обязательно, чтобы каждая ось соответствовала конкретной аппаратной оси. Например, сетка выше может описывать куб TPU v5p 4×4×4, где оси $X$ соответствуют 2 аппаратные оси. Это пригодится позже, когда мы будем описывать параллелизм данных по нескольким осям.

Случай 3: у обоих сомножителей сворачиваемые измерения шардированы

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

$$\textbf{A}[I, J_X] \cdot \textbf{B}[J_X, K] \rightarrow C[I, K]$$

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

$$\textbf{A}[I, J_X] \cdot_\text{LOCAL} \textbf{B}[J_X, K] \rightarrow C[I, K] \{\ U_X \}$$

Запись { U_X } читается как «не редуцировано (unreduced) по оси сетки X» и означает, что операция в каком-то смысле «не завершена»: она закончится только после итогового суммирования. Запись $\cdot_\text{LOCAL}$ означает, что мы выполняем локальную сумму, но оставляем результат нередуцированным.

Это можно понимать через связь умножения матриц с внешними произведениями:

$$A \cdot B = \sum_{i=1}^{P} \underbrace{A_{:,i} \otimes B_{i,:}}_{\in \mathbb{R}^{n \times m}}$$

где ⊗ — внешнее произведение. Значит, если у TPU номер i на оси X есть i-й столбец A и i-я строка B, можно локальным умножением получить $A_{:,i} \otimes B_{i,:} \in \mathbb{R}^{n\times m}$. В каждой клетке этой матрицы — i-е слагаемое суммы, которая стоит в этой клетке у A • B. Остаётся выполнить суммирование по P, которое мы шардировали по оси сетки X, чтобы получить полное A • B. Так же это работает, если записать A и B блоками (то есть шардами) и затем просуммировать по каждому получившемуся шарду результата.

Выполнить это суммирование можно полным AllReduce по оси X:

$$\begin{aligned} A[I, J_X] \cdot_\text{LOCAL} B[J_X, K] &\rightarrow C[I, K] \{ U_X \} \\ \textbf{AllReduce}_X C[I, K] \{ U_X \} &\rightarrow C[I, K] \end{aligned}$$

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

$$\textbf{AllReduce}_Y A[I_X, J] \{U_Y\} \rightarrow A[I_X, J]$$

То есть он просто убирает суффикс $\{U_Y\}$, в остальном оставляя результат прежним.

Сколько стоит AllReduce? Одна мысленная модель AllReduce: каждое устройство отправляет свой шард соседям и суммирует все полученные шарды. Ясно, что это дороже AllGather: каждый «шард» имеет ту же форму, что и весь массив. Как правило, AllReduce вдвое дороже AllGather. Один способ это увидеть — заметить, что AllReduce можно выразить как композицию двух других примитивов: ReduceScatter и AllGather. Как и AllReduce, ReduceScatter сводит частичные суммы массива, но результат оказывается «разбросан» (scattered), то есть разделён по заданному измерению. AllGather собирает все эти куски и «снимает разбиение / шардирование / размножает» логическую ось вдоль этой физической оси.

$$\begin{aligned} \textbf{ReduceScatter}_{Y,J} : A[I_X,J] \{U_Y\} &\rightarrow A[I_X, J_Y] \\ \textbf{AllGather}_Y : A[I_X, J_Y] &\rightarrow A[I_X, J] \end{aligned}$$

А что такое ReduceScatter? Как AllGather собирает шардированный массив (убирая нижний индекс), так ReduceScatter суммирует нередуцированный (частично просуммированный) массив, а затем разбрасывает (шардирует) другую логическую ось по той же оси сетки: $X[F]\{U_Y\} \to X[F_Y]$. Анимация в оригинале (reduce-scatter.gif) показывает, как это делается: очень похоже на AllGather, только вместо того чтобы сохранять каждый шард, мы их суммируем. Поэтому задержка примерно та же, не считая времени на само суммирование.

Время связи на каждом переходе — просто байты шарда $V / Y$, делённые на пропускную способность $W_\text{ici}$, как и для AllGather, так что

$$T_{\text{AllGather или ReduceScatter}} = \frac{V}{W_\text{ici}}$$ $$T_{\text{AllReduce}} = 2 \cdot \frac{V}{W_\text{ici}}$$

где $W_\text{ici}$ — двусторонняя пропускная способность, при условии что есть полное кольцо для редукции.

Случай 4: у обоих сомножителей несворачиваемое измерение шардировано по одной оси

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

$$A[I_X, J] \cdot B[J, K_X] \rightarrow C[I_X, K_X]$$

Это недопустимо, потому что данный шард, скажем i, вдоль измерения X получил бы (i, i)-й шард C, то есть диагональный элемент. Информации во всех шардах вместе не хватает, чтобы восстановить что-то кроме диагональных элементов результата, поэтому такое шардирование запрещено.

Выход — сделать AllGather некоторых измерений. Есть два варианта:

$$\begin{aligned} \textbf{AllGather}_X A[I_X, J] &\rightarrow A[I, J] \\ A[I, J] \cdot B[J, K_X] &\rightarrow C[I, K_X] \end{aligned}$$

или

$$\begin{aligned} \textbf{AllGather}_X B[J, K_X] &\rightarrow B[J, K] \\ A[I_X, J] \cdot B[J, K] &\rightarrow C[I_X, K] \end{aligned}$$

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

Подробнее о коммуникационных примитивах TPU

Четыре случая выше ввели несколько «базовых коммуникационных примитивов», которые используются для умножения шардированных матриц:

1. AllGather: убирает нижний индекс из шардирования, собирая шарды.
2. ReduceScatter: убирает суффикс «нередуцировано», суммируя шарды вдоль этой оси, и оставляет массив шардированным по второй оси.
3. AllReduce: убирает суффикс «нередуцировано», оставляя массив нешардированным вдоль этой оси.

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

Последний коммуникационный примитив: AllToAll

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

$$\textbf{AllToAll}_{X, J} A[I_X, J] \rightarrow A[I, J_X]$$

AllToAll обычно нужны, чтобы перестроить шардированные раскладки между разными частями шардированного вычисления с несовместимыми схемами раскладки. Они естественно возникают в шардированных моделях со смесью экспертов. AllToAll можно представлять как перенос нижнего индекса с одной оси на другую. Поскольку AllToAll не нужно размножать все данные каждого шарда по кольцу, он на самом деле дешевле AllGather (в 4 раза)⁵.

Анимация в оригинале (all-to-all.gif): AllToAll по кольцу.

Если обобщить на N-мерный AllToAll, общая стоимость для массива из $V$ байт (суммарно по всем устройствам) на сетке A×B×C:

$$T_\text{AllToAll} = \frac{V \cdot \max(A, B, C, ...)}{4 \cdot N \cdot W_\text{ici}}$$

где, как обычно, $W_\text{ici}$ — двусторонняя пропускная способность ICI, а $N = A \cdot B \cdot C \cdot \ldots$ — общее число устройств. В терминах байт на устройство $V / N$ стоимость равна $(V / N) \cdot \max(A, B, C, ...) / (4 \cdot W_\text{ici})$. Для одномерной сетки это сводится к $V / (4 \cdot W_\text{ici})$ — ¼ стоимости AllGather. В двумерном случае стоимость даже уменьшается с размером наименьшей оси.

Попутно — нестрогий вывод. Начнём с одномерного тора $\mathbb{Z} / N\mathbb{Z}$. Если выбрать источник и приёмник случайно, в среднем между ними N / 4 переходов, что даёт стоимость $(V \cdot N) / (4 \cdot N)$. В N-мерном торе каждая ось, по сути, независима. У каждого узла $1 / N$ байт, и в среднем его данные должны пройти $\max(A, B, C, …) / 4$ переходов. Можно вывести это и через бисекционную пропускную способность: при AllToAll каждая половина сетки отправляет половину своих данных ($V / 4$ байт) другой половине. Самое узкое сечение перпендикулярно самой длинной оси и пересекает $2 \cdot N / \max(A, B, …)$ каналов (две плоскости разреза с учётом замыкания), что даёт одностороннюю пропускную способность $N \cdot W_\text{ici} / \max(A, B, …)$. Делим — и получаем формулу выше.

Подробнее о ReduceScatter

ReduceScatter — более фундаментальная операция, чем кажется: на самом деле это производная AllGather, и наоборот. То есть если в прямом проходе у нас

$$\textbf{AllGather}_X A[I_X] \rightarrow A[I]$$

то производные обратного режима A’ (в общем случае разные на каждом шарде) мы сводим ReduceScatter, получая шардированное A’:

$$\textbf{ReduceScatter}_X A'[I] \{ U_X \} \rightarrow A'[I_X]$$

Аналогично $\text{ReduceScatter}_X(A[I] \{U_X\}) \to A[I_X]$ в прямом проходе означает $\text{AllGather}_{X}(A'[I_X]) \to A'[I]$ в обратном.

Почему AllGather и ReduceScatter — производные друг друга

Дело в том, что трансляция (broadcast) и редукция (reduce) — транспонированные друг к другу линейные операторы, а AllGather и ReduceScatter — внешние произведения (их также называют произведениями Кронекера) трансляции и редукции соответственно. Конкретно: если у нас вектор $x \in \mathbb{R}^n$, любое число устройств $p \in \mathbb{N}$ и $u = (1, \ldots, 1) \in \mathbb{R}^p$, то трансляцию и редукцию можно определить так (это должно совпадать с интуитивным пониманием):

$$\begin{aligned} \text{broadcast} &: \mathbb{R}^n \rightarrow \mathbb{R}^{p n}, &\quad \text{broadcast} &= u \otimes \mathbf{I}_n \\ \text{reduce} &: \mathbb{R}^{p n} \rightarrow \mathbb{R}^n, &\quad \text{reduce} &= u^T \otimes \mathbf{I}_n \end{aligned}$$

Пример при $n = 1$, $p = 2$. Если $x = (7)$, то $\text{broadcast}(x) = \begin{pmatrix} 1 \\ 1 \end{pmatrix} x = \begin{pmatrix} 7 \\ 7 \end{pmatrix}$ — как и ожидалось, вектор из $\mathbb{R}^n$ размножен в $\mathbb{R}^{pn}$. При $y = (8, 9)$: $\text{reduce}(y) = \begin{pmatrix} 1 & 1 \end{pmatrix} \begin{pmatrix} 8 \\ 9 \end{pmatrix} = (17)$ — вектор из $\mathbb{R}^{p n}$ сведён к вектору из $\mathbb{R}^{n}$. Поскольку $(A \otimes B)^T = A^T \otimes B^T$ для любых матриц $A$ и $B$, получаем $\text{reduce} = \text{broadcast}^T$. AllGather и ReduceScatter выражаются как внешние произведения:

$$\text{AllGather} = \text{broadcast} \otimes \mathbf{I}_p : \mathbb{R}^{p n} \rightarrow \mathbb{R}^{p^2 n}, \qquad \text{ReduceScatter} = \text{reduce} \otimes \mathbf{I}_p : \mathbb{R}^{p^2 n} \rightarrow \mathbb{R}^{p n}$$

Здесь $\mathbb{R}^{p^2 n}$ понимается как $\mathbb{R}^{p \times p n}$ — по одному вектору $\mathbb{R}^{p n}$ на каждое из $p$ устройств. Советуем поиграть с маленькими примерами, скажем $n = 2$, $p = 3$, чтобы увидеть, как эти операторы выглядят в виде матриц. По тому же свойству транспонирования $\text{AllGather}^T = \text{ReduceScatter}$ и $\text{ReduceScatter}^T = \text{AllGather}$. Транспонирование и возникает при обратном распространении: если $y = Ax$ для линейного оператора $A$ (например, AllGather или ReduceScatter), то при обратном распространении у нас есть производная потерь по $y$, $\frac{\partial L}{\partial y}$, и мы получаем $\frac{\partial L}{\partial x} = A^T \frac{\partial L}{\partial y}$. Отсюда производная AllGather — ReduceScatter, и наоборот.

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

$$A[I, J_X] \cdot B[J_X, K] \rightarrow C[I, K_X]$$

Тогда вместо AllReduce можно выполнить ReduceScatter, а AllGather — при желании позже:

$$\begin{aligned} A[I, J_X] \cdot_{LOCAL} B[J_X, K] &\rightarrow C[I, K] \{ U_X \} \\ \textbf{ReduceScatter}_{X,K} C[I, K] \{ U_X \} &\rightarrow C[I, K_X] \end{aligned}$$

Заметьте: ReduceScatter вводит шардированное измерение, и здесь у него есть естественная свобода шардировать либо по I, либо по K. Используя ReduceScatter, обычно нужно выбрать, какому именованному измерению добавить шардирование (хотя выбор обычно диктуется контекстом модели). Поэтому мы пишем ReduceScatter_X,K, указывая ось, по которой шардировать.

Как совместить связь при умножении матриц с вычислениями

Как мы обсуждали в главе 1, обычно мы считаем, что связь всегда можно совместить с какими-то полезными вычислениями, если она достаточно быстра. Коллективные операции этого раздела, как правило, можно совместить с вычислениями самого умножения матриц, но это нетривиально. Алгоритм для этого называется коллективным умножением матриц (collective matmul); впервые он описан в Wang et al..

Анимация в оригинале (ag_matmul.gif): как одно шардированное умножение матрицы на вектор совмещается с последующим AllReduce (случай 3 выше). Полное умножение матриц состоит из многих умножений матрицы на вектор.

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

Что мы узнали?

• Шардирование массива задаётся сеткой (Mesh), которая называет физические аппаратные оси сетки TPU, и шардированием (Sharding), которое назначает имена осей сетки логическим осям массива.
  • Например, A[I_XY, J] описывает абстрактный массив A, первое измерение которого шардировано по двум осям сетки X и Y. Вместе с Mesh(mesh_shape=(4, 8), axis_names=('X', 'Y')) или сокращённо Mesh({'X': 4, 'Y': 8}) это говорит, что массив разделён на 32 части по первому измерению.
• Арифметика с шардированными массивами работает в точности как с нешардированными, если не выполнять свёртку по шардированной оси. В этом случае нужно добавить связь. Рассматриваем четыре случая:
  1. Ни один массив не шардирован по сворачиваемому измерению: связь не нужна.
  2. Один массив шардирован по сворачиваемому измерению (или сворачиваемые измерения шардированы по разным осям): перед операцией делаем AllGather одного из входов.
  3. Оба массива одинаково шардированы по сворачиваемому измерению: перемножаем шарды локально, затем делаем AllReduce или ReduceScatter.
  4. Оба массива шардированы по одной оси сетки по несворачиваемому измерению: сначала делаем AllGather одного из входов.
• TPU используют примерно 4 базовых коммуникационных примитива:
  1. AllGather: $[A_X, B] \to [A, B]$
  2. ReduceScatter: $[A, B] \{U_X\} \to [A_X, B]$
  3. AllToAll: $[A, B_X] \to [A_X, B]$
  4. AllReduce: $[A_X, B]\{U_Y\} \to [A_X, B]$ (формально не примитив, так как сочетает ReduceScatter + AllGather)
• Стоимость и задержка каждой из этих операций не зависят от размера оси (пока они ограничены пропускной способностью), а только от размера входных массивов и пропускной способности канала. Для одностороннего AllGather/ReduceScatter:

$$T_{\text{AllGather или ReduceScatter}} = \frac{\text{объём данных}}{\text{пропускная способность}} \cdot \frac{\text{ось} - 1}{\text{ось}} \longrightarrow \frac{\text{объём данных}}{\text{пропускная способность (двусторонняя)}}$$

• AllReduce состоит из ReduceScatter и следующего за ним AllGather и поэтому вдвое дороже. AllToAll передаёт шарды лишь на часть кольца и поэтому стоит ¼ от AllGather. Сводка:

Операция	Описание	Запись	Время
AllGather	Собирает все шарды шардированного массива вдоль оси, убирая нижний индекс	$[A_X, B] \to [A, B]$	байты / (двусторонняя пропускная способность ICI × число осей)
ReduceScatter	Суммирует частично просуммированный массив вдоль оси и шардирует его по другой оси (добавляя нижний индекс)	$[A, B] \{U_X\} \to [A_X, B]$	как AllGather
AllReduce	Суммирует частично просуммированный массив вдоль оси, убирает { UX }; сочетает AllGather и ReduceScatter	$[A_X, B]\{U_Y\} \to [A_X, B]$	2 × AllGather
AllToAll	Собирает (размножает) ось и шардирует другое измерение по той же оси	$[A, B_X] \to [A_X, B]$	AllGather / 4 для двунаправленного кольца

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

Книга	torch.distributed	Где в курсе
AllGather	dist.all_gather_into_tensor(out, x)	FSDP собирает веса слоя перед прямым проходом
ReduceScatter	dist.reduce_scatter_tensor(out, x)	FSDP раздаёт градиенты по шардам
AllReduce	dist.all_reduce(x)	DDP усредняет градиенты
AllToAll	dist.all_to_all_single(out, x)	экспертный параллелизм в MoE

Пример из главы на PyTorch: DDP — это случай «батч шардирован, веса реплицированы», градиенты весов приходят нередуцированными { U_batch } и сводятся AllReduce. FSDP (ZeRO-3) — то же, но веса шардированы: перед вычислением AllGather, после — ReduceScatter градиентов. Подробнее — в главе 5 и в уроке «Распределённое обучение».

Задачи

Вот несколько поучительных задач по материалу главы. Ответы есть пока не ко всем — авторы обещают дописывать.

Вопрос 1 [реплицированное шардирование]. Массив шардирован как $A[I_X, J, K, \ldots]$ (то есть только по $X$) на сетке Mesh({'X': 4, 'Y': 8, 'Z': 2}). Каково отношение общего числа байт, которые $A$ занимает на всех чипах, к размеру одной копии массива?

Ответ

Массив шардирован только по X размера 4, так что фактически каждый шард имеет размер $[I / 4, J, K, \ldots] = \text{sizeof}(A) / 4$. Поскольку массив реплицирован по Y и Z, общий размер — $Y \cdot Z \cdot \text{sizeof}(A)$, так что отношение общего размера к размеру одной копии — $Y \cdot Z \cdot \text{sizeof}(A) / \text{sizeof}(A) = 16$.

Вопрос 2 [задержка AllGather]. Сколько должен занимать $\text{AllGather}_X([B_X, D_Y])$ на срезе TPU v4p 4x4x4 с сеткой Mesh({'X': 4, 'Y': 4, 'Z': 4}) при $B=1024$ и $D=4096$ в bfloat16? А $\text{AllGather}_{XY}([B_X, D_Y])$? А $\text{AllReduce}_Z([B_X, D_Y] \{U_Z \})$?

Ответ

У нас полный куб 4x4x4, так что замыкающие каналы есть по всем осям, и доступна двусторонняя пропускная способность 9e10.

1. Мы собираем только по одной оси, а другая шардирована, так что фактически собираем $2BD / Y$ байт по 1 оси. Если представить один шард вдоль оси Y, AllGather по X выглядит как AllGather нешардированного массива с 1 / Y байт. Двусторонняя пропускная способность ICI у TPU v4p — 9e10 байт/с, так что это займёт $2BD / (9 \cdot 10^{10} \cdot Y) = 2 \cdot 1024 \cdot 4096 / (9 \cdot 10^{10} \cdot 4) = 23\,\mu s$.
2. Пропускная способность вдвое больше, но мы собираем весь массив, так что T = 2BD / (2 * W) = 2*1024*4096 / (2 * 9e10) = 46us. Это далеко от границы по задержке в 4 мкс (1 мкс на переход), так что всё в порядке.
3. AllReduce вдвое дороже AllGather. Каждый шард имеет размер $2BD / (X \cdot Y)$, так что стоимость около $4BD / (X \cdot Y \cdot W)$, то есть примерно 4 * 1024 * 4096 / (16 * 9e10) = 11.6us.

Забавный факт: части (1) и (2) на самом деле не оптимальны: массив реплицирован ещё и по неиспользуемой оси Z, и её простаивающие каналы можно задействовать. Сначала бесплатно перешардируем $[B_X, D_Y] \to [B_{XZ}, D_Y]$ (каждое устройство просто отбрасывает часть своего шарда), а затем выполним $\text{AllGather}_{XZ}$ (или $\text{AllGather}_{XYZ}$), придя к тому же итоговому состоянию, но собирая по большему числу осей. Это сокращает часть (1) до 11,5 мкс, а часть (2) — до 31 мкс. На практике этого добиваются, просто сразу шардируя по большему числу осей, — это одна из причин шардировать массивы как можно мельче.

Вопрос 3 [AllGather, ограниченный задержкой]. Допустим, мы выполняем $\text{AllGather}_X([B_X])$, но $B$ очень мало (скажем, 128). Сколько это займёт на срезе TPU v4p 4x4x4 с сеткой Mesh({'X': 4, 'Y': 4, 'Z': 4}) в bfloat16? Подсказка: вероятно, вы упираетесь в задержку.

Ответ

Массив в bfloat16 занимает всего 256 байт, по 64 на устройство. Ось размера 4 у TPU v4p имеет замыкающий канал, так что массив можно отправлять в обе стороны. При односторонней пропускной способности 4.5e10 каждый переход займёт примерно 64 / 4.5e10 ~ 0, так что мы определённо упираемся в задержку. Считаем переходы: полный сбор можно сделать всего за 2 перехода, так что около 2 мкс — хорошая оценка.

Вопрос 4 [стратегии умножения матриц]. Чтобы выполнить $X[B, D] \cdot_D Y[D_X, F] \to Z[B, F]$, в этой главе мы советуем выполнить $\text{AllGather}_X(Y[D_X, F])$ и перемножить полностью реплицированные матрицы (случай 2, стратегия 1). Вместо этого можно перемножить локальные шарды: $X[B, D_X] \cdot_D Y[D_X, F] \to Z[B, F] \{U_X\}$ (случай 3, стратегия 2), — а затем выполнить $\text{AllReduce}_X(Z[B, F] \{ U_X\})$. Сколько FLOPs и связи требует каждая? Какая лучше и почему?

Ответ

Начнём с базовой (стратегия 1). Как мы показали, стоимость AllGather — $2DF / W_\text{ici}$. Когда массивы полностью реплицированы, общее время вычислений — $2BDF / C$ (где $C$ — FLOPs/s ускорителя, поскольку каждый TPU выполняет одни и те же FLOPs). Итак,

$$T_\text{всего (стратегия 1)} = \max\left(\frac{2BDF}{C}, \frac{2DF}{W_\text{ici}}\right)$$

Для сравнения, новая стратегия (стратегия 2) делает AllReduce над $2BF$ байт со стоимостью $4BF / W_\text{ici}$, но выполняет в $X$ раз меньше FLOPs (вычисления шардированы). То есть мы делаем $2\cdot B\cdot D\cdot F / X$ FLOPs, а итоговый AllReduce передаёт $2 \cdot 2 \cdot B \cdot F$ байт в bfloat16. Значит, общее время стратегии 2 (без AllGather, только AllReduce в конце) примерно

$$T_\text{всего} = \max\left(\frac{2BDF}{X \cdot C}, \frac{4BF}{W_\text{ici}}\right)$$

Вопрос: что из этого больше? Стратегия (2) упирается в вычисления, когда $D / (X \cdot C) > 2 / W_\text{ici}$, то есть когда $D / 2X > C / W_\text{ici} \approx 2550 \rightarrow X < D / (2 \cdot 2550)$. Разумно ожидать $D \approx 8$ тысяч, то есть примерно $X < 2$, что маловероятно, — значит, со стратегией 2 мы почти всегда упираемся в связь. С базовой (стратегия 1) мы упираемся в связь, когда $B < C / W_\text{ici} = 2550$, что часто, но не всегда верно.

Так что при $B < 2550$ мы упираемся в связь в обоих случаях, и

$$T_\text{связь, стратегия 2} < T_\text{связь, стратегия 1} \Leftrightarrow \frac{4BF}{W_\text{ici}} < \frac{2DF}{W_\text{ici}}$$

что верно при $D > 2B$, где $2B < 5100$. Это часто так, поэтому стратегия 2 иногда лучше при маленьком батче. При большом батче ($B > 2550$)

$$T_\text{связь, стратегия 2} < T_\text{вычисления, стратегия 1} \Leftrightarrow \frac{4BF}{W_\text{ici}} < \frac{2BDF}{C}$$

что верно при $2 / W_\text{ici} < D / C$, то есть при $D > 2 \cdot 2550 = 5100$, что для больших моделей обычно выполняется. Так что эта альтернативная стратегия обычно лучше для больших моделей, если только $D$ не мал.

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

5. [Минимальная задержка] Как шардировать входы умножения $A[I, J] \cdot_J B[J, K] \to C[I, K]$ на TPU v4p 4x4x4, чтобы получить наименьшую задержку при полностью реплицированном результате? Каковы FLOPs и время связи?
6. Какая связь нужна для $A[I_X, J_Y] \cdot_J B[J_Y, K] \to C[I_X, K]$ на TPU v5e 4x4 и сколько времени уходит на связь и на вычисления?
7. Блок трансформера $In[B, D] \cdot W_\text{in}[D, F] \cdot W_\text{out}[F, D]$ при $D=8192$, $F=32768$, $B=128$ в bfloat16 на срезе TPU v5e 2x2, где у каждого TPU всего 300 МБ свободной памяти: как шардировать вход, $W_\text{in}$, $W_\text{out}$ и выход, чтобы уложиться в память и минимизировать время?
8. [Задание] По фрагменту кода выше создайте шардированный массив и замерьте четыре примитива (jax.lax.all_gather, jax.lax.psum, jax.lax.psum_scatter, jax.lax.all_to_all).
9. Почему при одном шардированном по сворачиваемому измерению входе лучше AllGather, чем «шардированный матмул + AllReduce»?
10. [AllToAll] Откуда в стоимости AllToAll множитель 1/4 и как он изменился бы при односторонних каналах ICI?

Словарик

English	Русский
sharding, partitioning	шардирование, разбиение
global (logical) shape / device local shape	глобальная (логическая) форма / локальная форма на устройстве
device mesh, mesh axis	сетка устройств, ось сетки
fully / partially replicated	полностью / частично реплицированный (размноженный)
contracting dimension	сворачиваемое измерение (по которому суммируем)
block matrix	блочная матрица
collective operation, primitive	коллективная операция, примитив
AllGather, ReduceScatter, AllReduce, AllToAll	сбор, редукция с раздачей, общая редукция, «все всем»
unreduced { UX }	нередуцированный (частичные суммы) по оси X
throughput-bound / latency-bound	ограниченный пропускной способностью / задержкой
outer (Kronecker) product	внешнее произведение (Кронекера)
collective matmul	коллективное умножение матриц
resharding, sharded transposition	перешардирование, шардированное транспонирование

Footnotes

1. Стоит отметить, что параллелить можно и ради скорости. Даже если модель поместилась бы на меньшем числе чипов, масштабирование на большее просто даёт больше FLOPs/s. Например, при инференсе модель иногда помещается на меньших топологиях, но мы масштабируемся на большие, чтобы снизить задержку. Так же и при обучении часто масштабируются на больше чипов, чтобы сократить время шага. ↩

2. AllGather на GPU может работать так же: из GPU в узле строится кольцо, и куски передаются по нему в этом (произвольном) порядке. ↩

3. Двойка в числителе — потому что мы используем двустороннюю пропускную способность: отправляем $V / X$ в каждую сторону, всего $2V / X$. ↩

4. Строго говоря, $\lfloor X / 2 \rfloor$. ↩

5. Для двунаправленных колец чётного размера каждое устройство отправляет $(N/2 + (N/2-1) + … + 1)$ кусков вправо и $((N/2-1) + … + 1)$ кусков влево, итого $0{,}5 \cdot (N / 2) \cdot (N/2 + 1) + 0{,}5 \cdot (N / 2) \cdot (N/2 - 1) = N^2/4$. Размер каждого куска («шарда шарда») — $\text{bytes} / N^2$, так что стоимость на устройство — $(\text{bytes} / N^2) \cdot N^2 / 4 = \text{bytes} / 4$. Этот результат масштабируется на все устройства, поскольку общая пропускная способность растёт с числом устройств. ↩