How To Scale Your Model, глава 2: как устроены TPU

Глава о том, как устроены TPU, как они соединяются друг с другом в многочиповые срезы и поды и как всё это влияет на производительность наших моделей. — подзаголовок главы.

О переводе

Полный перевод главы 2 книги How To Scale Your Model (лицензия MIT, © 2025 Google LLC), включая задачи и приложения. Анимации и схемы оригинала описаны словами, ссылки на них даны текстом. Предыдущая глава — «Всё о roofline», следующая — «Шардированные матрицы». Глава о GPU — здесь.

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

TPU у вас нет, но глава нужна ради трёх идей, которые работают и на GPU: (1) ускоритель — это блок умножения матриц, быстрая память рядом с ним и медленная память дальше; (2) каналы связи образуют иерархию скоростей; (3) нагрузка на каждый канал должна быть пропорциональна его скорости. Соответствия: MXU ≈ тензорные ядра, VMEM ≈ общая память (shared memory) и кэш L2, HBM ≈ видеопамять GDDR6X вашей RTX 4080, ICI ≈ NVLink (у вас его нет — только PCIe). Сначала прочитайте урок «Как устроена видеокарта и куда уходит время».

Что такое TPU?

TPU — это, по сути, вычислительное ядро, специализированное на умножении матриц (оно называется TensorCore), подключённое к стопке быстрой памяти (она называется памятью с высокой пропускной способностью, HBM).

Рисунок в оригинале (tpu-chip.png): основные компоненты чипа TPU. TensorCore — серый прямоугольник слева, внутри него блок умножения матриц (MXU), векторный блок (VPU) и векторная память (VMEM); справа — HBM, связанная с TensorCore широким каналом.

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

• MXU (Matrix Multiply Unit, блок умножения матриц) — сердце TensorCore. В большинстве поколений TPU он выполняет одно умножение bf16[8,128] @ bf16[128,128] -> f32[8,128]¹ за 8 тактов с помощью систолического массива (подробности — в приложении Б).
  • Это около 5e13 FLOPs/s в bf16 на один MXU при частоте 1,5 ГГц у TPU v5e. В большинстве TensorCore по 2 или 4 MXU, так что, например, общая скорость TPU v5e в bf16 — 2e14 FLOPs/s.
  • TPU поддерживают и умножения в пониженной точности с большей скоростью (например, чип TPU v5e выполняет 4e14 операций int8 в секунду).
• VPU (Vector Processing Unit, блок векторной обработки) выполняет общие математические операции вроде активаций ReLU или поэлементного сложения и умножения векторов. Здесь же выполняются редукции (суммы). Подробнее — в приложении А.
• VMEM (Vector Memory, векторная память) — сверхоперативная память на кристалле, в TensorCore, рядом с вычислительными блоками. Она намного меньше HBM (например, 128 МиБ у TPU v5e), но у неё гораздо больше пропускная способность к MXU. VMEM работает отчасти как кэш L1/L2 у процессоров, но она намного больше и управляется программистом. Данные из HBM нужно скопировать в VMEM, прежде чем TensorCore сможет с ними что-то вычислять.

TPU очень, очень быстро умножают матрицы. В основном это они и делают — и делают хорошо. TPU v5p, один из самых мощных TPU на сегодня, выполняет 2.5e14 FLOPs в bf16 в секунду на ядро, или 5e14 на чип. Один под из 8960 чипов выполняет 4 экзаFLOPs/s в bf16. Это очень много. Это один из самых мощных суперкомпьютеров в мире. А у Google их много².

На схеме выше есть и другие компоненты, например SMEM и скалярный блок: они отвечают за управление ходом программы и кратко описаны в приложении А, но понимать их не обязательно. А вот HBM важна и довольно проста:

• HBM (High Bandwidth Memory, память с высокой пропускной способностью) — большой кусок быстрой памяти, где хранятся тензоры для TensorCore. Объём HBM обычно — десятки гигабайт (например, у TPU v5e 16 ГиБ HBM).
  • Когда тензоры нужны для вычисления, они потоком идут из HBM через VMEM (см. ниже) в MXU, а результат записывается из VMEM обратно в HBM.
  • Пропускную способность между HBM и TensorCore (через VMEM) называют «пропускной способностью HBM» (обычно около 1–2 ТБ/с); она ограничивает скорость вычислений в задачах, упирающихся в память.

Как правило, все операции TPU конвейеризованы и перекрываются. Чтобы выполнить умножение $X \cdot A \to Y$, TPU сначала копирует куски матриц $A$ и $X$ из HBM в VMEM, затем загружает их в MXU, который перемножает куски 8×128 (для $X$) и 128×128 (для $A$), затем по кускам копирует результат обратно в HBM. Чтобы делать это эффективно, умножение конвейеризовано: копирование в VMEM и из неё перекрывается с работой MXU. Так MXU продолжает работать, а не ждёт передачи данных, и умножения упираются в вычисления, а не в память.

Вот пример того, как можно выполнить поэлементное произведение из HBM.

Анимация в оригинале (pointwise-product.gif): поэлементное произведение на TPU с загрузкой байт из HBM. Байты выходят из памяти кусками, а частичные результаты по конвейеру отправляются обратно, не дожидаясь, пока весь массив будет готов.

Умножение матриц выглядело бы почти так же, только данные загружались бы в MXU вместо VPU, а загрузки и сохранения шли бы в другом порядке, потому что один и тот же кусок весов используется для нескольких кусков активаций. Видно, как куски данных идут в VMEM, затем в VREG (векторные регистры), затем в векторный блок, затем обратно в VMEM и HBM. Как мы сейчас увидим, если загрузка из HBM в VMEM медленнее, чем FLOPs в векторном блоке (или MXU), мы «упираемся в пропускную способность»: VPU или MXU простаивают без работы.

Главный вывод

TPU очень просты. Они загружают веса из HBM в VMEM, затем из VMEM в систолический массив, который выполняет около 200 триллионов умножений-сложений в секунду. Пропускные способности HBM ↔ VMEM и VMEM ↔ систолический массив задают фундаментальные пределы того, какие вычисления TPU могут выполнять эффективно.

VMEM и арифметическая интенсивность. VMEM намного меньше HBM, но её пропускная способность к MXU гораздо выше. Как мы видели в главе 1, значит, если все входы и выходы алгоритма помещаются в VMEM, он гораздо реже упирается в связь. Это особенно полезно, когда у вычисления низкая арифметическая интенсивность: пропускная способность VMEM примерно в 22 раза выше, чем у HBM, поэтому операции MXU, читающей из VMEM и пишущей в неё, для выхода на пиковые FLOPs достаточно арифметической интенсивности всего 10–20. Значит, если удастся поместить веса в VMEM, а не в HBM, умножения матриц будут упираться в FLOPs при гораздо меньших размерах батча. И алгоритмы с принципиально низкой арифметической интенсивностью всё равно могут быть эффективными. Просто VMEM настолько мала, что часто это трудно³.

Рисунок в оригинале (tpu-bandwidth.png): схема пропускных способностей внутри чипа: HBM ↔ VMEM — порядка 1–3 ТБ/с, VMEM ↔ MXU — примерно в 22 раза быстрее.

Чип TPU обычно (но не всегда) состоит из двух ядер TPU с общей памятью, и его можно считать одним большим ускорителем с вдвое большими FLOPs (такую конфигурацию называют «мегаядром», megacore). Так устроены TPU v4, v5 и v6 (в TPU v7 мегаядра нет, вместо него — быстрый канал между двумя ядрами). У более старых чипов память раздельная, и их считают двумя отдельными ускорителями (TPU v3 и старше). Чипы, оптимизированные для инференса, вроде TPU v5e, имеют лишь одно ядро на чип.

Чипы собраны по 4 на «плате» (tray), которая подключена к хосту с процессором через PCIe. Большинству читателей знаком именно этот формат: 4 чипа (8 ядер, хотя обычно их считают 4 логическими мегаядрами), доступные через Colab или одну TPU-VM. У инференс-чипов вроде TPU v5e на хост приходится 2 платы, а не 1, но и ядро на чип только одно, так что получается 8 чипов = 8 ядер⁴.

Пропускная способность PCIe ограничена. Как и канал HBM ↔ VMEM, соединение процессор ↔ HBM по PCIe имеет свою пропускную способность, которая ограничивает скорость загрузки из памяти хоста в HBM и обратно. Например, у TPU v4 пропускная способность PCIe — 16 ГБ/с в каждую сторону, то есть почти в 100 раз медленнее HBM. Выгружать данные в оперативную память хоста (процессора) и загружать их оттуда можно, но не очень быстро.

Сеть TPU

Внутри пода чипы связаны друг с другом сетью ICI. В старых поколениях (TPU v2 и TPU v3), у инференс-чипов (например, TPU v5e) и у Trillium (TPU v6e) ICI («inter-chip interconnects», межчиповые связи) соединяет 4 ближайших соседей (а связи по краям замыкают двумерный тор). TPU v4 и TPU v5p связаны с 6 ближайшими соседями (трёхмерный тор). Заметьте: эти связи не проходят через хосты, это прямые каналы между чипами.

Рисунок в оригинале (ici-wraparound.png): решётка чипов, где каждый связан с соседями, а крайние чипы каждой строки и столбца связаны между собой «замыкающими» (wraparound) каналами.

Тороидальная структура уменьшает максимальное расстояние между любыми двумя узлами с $N$ до $N / 2$, и связь становится гораздо быстрее. У TPU есть ещё конфигурация «скрученного тора» (twisted torus), которая сворачивает тор наподобие ленты Мёбиуса и ещё сильнее сокращает среднее расстояние между узлами.

Поды TPU (связанные ICI) бывают очень большими: максимальный размер пода (его называют суперподом) — 16x16x16 у TPU v4 и 16x20x28 у TPU v5p. Такие большие поды собраны из переконфигурируемых кубов 4x4x4 чипов, соединённых оптическими замыкающими каналами⁵, которые можно переконфигурировать, чтобы соединять очень большие топологии.

Рисунок в оригинале (tpu-rack.png): стойка с кубом 4x4x4 чипов и оптическими переключателями.

Можно запросить и топологии поменьше (например, 2x2x1, 2x2x2), но без замыкающих каналов. Это важная оговорка: обычно из-за этого время большинства операций связи удваивается. Любое кратное полного куба (например, 4x4x4 или 4x4x8) будет иметь замыкающие каналы, которые обеспечивают оптические переключатели⁶.

Рисунки в оригинале (subslices.png, more-subslices.png): примеры срезов разной формы внутри пода.

Поды TPU v5e и Trillium — это один двумерный тор 16x16 с замыкающими каналами по любой оси размера 16 (то есть у 8x16 замыкающий канал есть по длинной оси). TPU v5e и v6e (Trillium) не могут расширяться дальше тора 16×16, но поды могут общаться друг с другом по обычной сети дата-центра (DCN), которая связывает хосты TPU. И снова: можно запросить топологии поменьше, без замыканий по измерениям меньше 16.

Связь только с ближайшими соседями — ключевое отличие TPU от GPU. GPU соединены иерархией коммутаторов, которая приближает соединение «каждый с каждым» между всеми GPU, а не локальными связями, как у TPU. Обычно GPU внутри узла (8 GPU у H100 или до 72 у B200 NVL72) соединены напрямую, а в больших топологиях между GPU приходится $O(\log N)$ переходов. С одной стороны, это значит, что GPU могут отправлять любые данные за небольшое число переходов. С другой стороны, TPU намного дешевле (коммутаторы NVLink дороги), их проще соединять, и они масштабируются до гораздо больших топологий, потому что число каналов и пропускная способность на устройство постоянны. Подробнее — в главе о GPU.

ICI очень быстра по сравнению с DCN, но всё же медленнее, чем HBM. Например, у TPU v5p:

• 2.8e12 байт/с (2,8 ТБ/с) пропускной способности HBM на чип;
• 9e10 байт/с (90 ГБ/с) пропускной способности ICI на ось, по 3 оси на чип⁷;
• 6.25e9 байт/с (6,25 ГБ/с) исходящей пропускной способности DCN на TPU (через 1–2 сетевые карты на каждом хосте)⁸.

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

Обучение на нескольких срезах (multi-slice). Набор TPU, связанных ICI, называют срезом (slice). Разные срезы можно связать между собой через DCN, например чтобы соединить срезы на разных подах. DCN гораздо медленнее ICI, поэтому нужно стараться, чтобы вычисления как можно меньше ждали данных по DCN. DCN соединяет хосты, поэтому, чтобы передать буферы с TPU на TPU по DCN, их нужно сначала передать по PCIe на хост, затем отправить по сети, затем принять на целевом хосте и по PCIe загрузить в HBM.

flowchart LR
MXU["MXU: ~2e14 FLOPs/s"] <-->|"VMEM: ~22× HBM"| V["VMEM 128 МиБ"]
V <-->|"HBM: ~1–3 ТБ/с"| H["HBM 16–96 ГБ"]
H <-->|"ICI: ~45–90 ГБ/с на канал"| N["Соседний TPU"]
H <-->|"PCIe: ~16 ГБ/с"| C["Хост (CPU)"]
C <-->|"DCN: ~3–12 ГБ/с на TPU"| C2["Другой хост"]

Иерархия каналов связи TPU — от самого быстрого к самому медленному (схема переводчика по числам главы)

Главные выводы

• TPU просты, и в большинстве случаев их можно представлять как блок умножения матриц, связанный с памятью (сверхбыстро), с другими чипами по ICI (довольно быстро) и с остальным дата-центром по DCN (более-менее быстро).
• Связь ограничена разными сетевыми пропускными способностями, по убыванию скорости:
  • пропускная способность HBM — между TensorCore и его HBM;
  • пропускная способность ICI — между чипом TPU и его 4 или 6 ближайшими соседями;
  • пропускная способность PCIe — между хостом с процессором и его платой (платами) чипов;
  • пропускная способность DCN — между несколькими хостами, обычно не связанными ICI.
• Внутри среза TPU связаны по ICI только с ближайшими соседями. Значит, связь по ICI между далёкими чипами среза должна сначала пройти через промежуточные чипы.
• Матрицы весов нужно дополнять хотя бы до размера 128 (256 у TPU v6e) по обоим измерениям, чтобы заполнить MXU (на самом деле меньшие оси и так дополняются до 128).
• Умножение матриц в пониженной точности обычно быстрее. В поколениях, которые это поддерживают, TPU выполняют операции int8 или int4 примерно в 2 или 4 раза быстрее, чем FLOPs в bfloat16. Операции VPU по-прежнему выполняются в fp32.
• Чтобы не делать вычислительный блок TPU узким местом, нужно следить, чтобы объём связи по каждому каналу был пропорционален его скорости.

Характеристики TPU

Вот конкретные числа для наших чипов:

Модель	Размер пода	Размер хоста	HBM на чип	Пропускная способность HBM на чип (байт/с)	FLOPs/s на чип (bf16)	FLOPs/s на чип (int8)
TPU v3	32x32	4x2	32 ГБ	9.0e11	1.4e14	1.4e14
TPU v4p	16x16x16	2x2x1	32 ГБ	1.2e12	2.75e14	2.75e14
TPU v5p	16x20x28	2x2x1	96 ГБ	2.8e12	4.59e14	9.18e14
TPU v5e	16x16	4x2	16 ГБ	8.2e11	1.97e14	3.94e14
TPU v6e	16x16	4x2	32 ГБ	1.6e12	9.20e14	1.84e15
TPU7x	4x4x576	2x2x1	192 ГБ	7.4e12	2.30e15	4.61e15

Размер хоста — топология TPU, подключённых к одному хосту (например, у TPU v5e один хост с процессором подключён к 8 TPU в топологии 4x2). Подробнее о последнем поколении — в документации TPU7x. Вот числа межчиповых связей:

Модель	ICI на канал (в одну сторону, байт/с)	ICI на канал (в обе стороны, байт/с)
TPU v3	1.0e11	2.0e11
TPU v4p	4.5e10	9.0e10
TPU v5p	9.0e10	1.8e11
TPU v5e	4.5e10	9.0e10
TPU v6e	9.0e10	1.8e11
TPU7x	9.0e10	1.8e11

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

Пропускная способность PCIe обычно около 1.6e10 байт/с на TPU (3.2e10 у TPU v6e), а DCN — около 6.25e9 байт/с на TPU (12.5e9 у TPU v6e и TPU7x и 3.125e9 у TPU v5e).

🖥 Те же числа для вашей машины (примечание переводчика)

	RTX 4080 (ваша)	TPU v5e
Память на чип	32 ГБ GDDR6X	16 ГБ HBM
Пропускная способность памяти	7,2e11 байт/с	8,2e11 байт/с
FLOPs/s в bf16	9,75e13	1,97e14
Связь между чипами	PCIe 4.0: до 3,2e10 байт/с, без P2P на практике ~1e10	ICI: 4,5e10 на канал, 4 канала

Ваши четыре карты связаны не тором, а через процессор по PCIe, и это самый медленный канал в системе — в 20–70 раз медленнее видеопамяти (урок «Четыре видеокарты: DDP и FSDP на практике»). Вывод главы «нагрузка на канал должна быть пропорциональна его скорости» для вас значит: обмениваться между картами как можно реже (DDP с редкой синхронизацией градиентов, накопление градиентов), а не разрезать каждый слой между картами (тензорный параллелизм).

Задачи с решениями

Эти числа немного сухие, но они позволяют делать базовые roofline-оценки производительности модели. Решим несколько задач, чтобы показать, чем это полезно. Больше примеров — в главе 3.

Вопрос 1 [нижняя граница задержки LLM]. Допустим, вы хотите сэмплировать из модели на 200 млрд параметров в bf16, разбитой между 32 TPU v4p. Сколько времени займёт загрузка всех параметров из HBM в систолический массив? Подсказка: используйте числа выше.

Ответ

Мы загружаем sizeof(bf16) * 200e9 = 400e9 байт на 32 чипа, то есть 12,5e9 байт на чип, у каждого пропускная способность HBM 1,23e12. Значит, загрузка занимает около 10 мс.

Это очень полезно: это разумная нижняя граница задержки сэмплирования из модели. На каждом шаге сэмплирования нужно загрузить из HBM все параметры, поэтому шаг не может занять меньше 10 мс. На практике при маленьких батчах это почти достижимо.

Вопрос 2 [детали TPU]. Рассмотрим полный под TPU v5e. Сколько в нём хостов с процессорами? Сколько TensorCore? Каковы общие FLOPs/s всего пода? Сколько всего HBM? Сделайте то же для пода TPU v5p.

Ответ

У TPU v5e под — 16x16, а каждый хост — срез 4x2, значит, хостов 16*16 / 8 = 32. У TPU v5e на чип только одно ядро, значит, TensorCore — 256. Общие FLOPs/s — 16*16*2e14 = 5.1e16 в bfloat16. У каждого чипа 16 ГБ HBM, значит, всего 256 * 16 = 4TB памяти.

У полного пода TPU v5p 16x20x28 чипов, и каждый хост — 2x2x1, значит, хостов (16*20*28) / (2*2) = 2,240. У TPU v5p на чип по два TensorCore, значит, ядер 8960 * 2 = 17,920. Общие FLOPs/s — 8960 * 4.59e14 = 4.1e18 в bfloat16. У каждого чипа 96 ГБ HBM, значит, всего 8960 * 96 = 860TB памяти.

Вопрос 3 [операционная интенсивность PCIe]. Представьте, что мы вынуждены хранить большую матрицу весов $A$ типа $\text{bf16}[D, F]$ и батч активаций $x$ типа $\text{bf16}[B, D]$ в оперативной памяти хоста (DRAM) и хотим их перемножить. Всё работает на одном хосте, и мы используем один подключённый к нему чип TPU v6e. Можно считать, что $B \ll D$ и $F = 4D$ (в следующих главах мы увидим, почему это разумные предположения). При каком наименьшем размере батча $B$ мы останемся в упоре в FLOPs при передаче по PCIe? Пропускная способность PCIe — 1,6e10 байт/с.

Ответ

Нужно выполнить $2BDF$ операций с плавающей точкой, а каждый чип выполняет 9.2e14 операций в секунду. На это уходит $2BDF / 9{,}2 \cdot 10^{14}$ секунд. Нужно загрузить из DRAM $2DF + 2BD$ байт и записать туда $2BF$ байт. Узкое место — скорость передачи по PCIe, поэтому на передачу данных к TPU и обратно нужно $2 \cdot (BD + DF + BF) / 1{,}6 \cdot 10^{10}$ секунд. Мы хотим, чтобы вычисления шли дольше загрузки весов (считая, что всю загрузку можно совместить с вычислениями), то есть $2BDF / 9{,}2 \cdot 10^{14} > 2 \cdot (BD + DF + BF) / 1{,}6 \cdot 10^{10}$. Упрощая с учётом $B \ll D$ и $F = 4D$, получаем

$$\frac{8BD^2}{9{,}2 \times 10^{14}} > \frac{8D^2}{1{,}6 \times 10^{10}}$$

или

$$B > \frac{9{,}2 \times 10^{14}}{1{,}6 \times 10^{10}} \simeq 57\,500$$

Вопрос 4 [задержка умножения матриц в общем случае]. Допустим, мы хотим умножить матрицу весов int8[16384, 4096] на матрицу активаций размера int8[B, 4096], где B — неизвестный размер батча. Для начала мы на 1 TPU v5e.

1. Сколько времени займёт это умножение в зависимости от B? Подсказка: посчитайте, сколько займёт загрузка массивов из HBM и сколько — само умножение. Что из этого — узкое место?
2. А если бы мы выполняли эту операцию из VMEM? Сколько бы она заняла в зависимости от B?

Ответ

(1) Нужно выполнить $2 \cdot 4096 \cdot 16384 \cdot B = 1{,}3 \times 10^{8} \cdot B$ операций. Значит, $T_{\text{math}} = (1{,}3 \times 10^{8} \cdot B) / 3{,}94 \times 10^{14}$ секунд. Из HBM в VMEM нужно загрузить $16384 \cdot 4096 + 4096 \cdot B$ байт и записать из VMEM в HBM $16384 \cdot B$ байт. Значит, $T_{\text{comms}} = (6{,}7 \times 10^{7} + 2 \times 10^{4} \cdot B) / 8{,}2 \times 10^{11}$ секунд. При максимально возможном перекрытии связи и вычислений всё умножение займёт примерно

$$\max\{T_{\text{math}}, T_{\text{comms}}\} = \max\left\{\frac{1{,}3 \times 10^{8} \cdot B}{3{,}94 \times 10^{14}}, \frac{6{,}7 \times 10^{7} + 2 \times 10^{4} \cdot B}{8{,}2 \times 10^{11}}\right\}$$

Мы упрёмся в FLOPs, когда $\frac{6{,}7 \times 10^{7} + 2 \times 10^{4} \cdot B}{8{,}2 \times 10^{11}} < \frac{1{,}3 \times 10^{8} \cdot B}{3{,}94 \times 10^{14}}$, то есть при $B > 267$. Это чуть больше числа 240, которое мы вывели в главе 1, потому что здесь учтено полное влияние $D$ и $F$.

(2) Если загружать из VMEM, будем считать пропускную способность VMEM к MXU в 22 раза больше пропускной способности HBM ↔ VMEM. Тогда знаменатель для загрузки данных превращается из 8,2e11 в 1,80e13, и получаем $B > 11$. Заметьте: на практике отдать всю пропускную способность VMEM под загрузку матрицы весов нельзя, так что на деле получится ближе к 20.

Вопрос 5 [пропускная способность ICI]. Допустим, у нас срез TPU v5e 4x4. Мы хотим отправить массив типа bf16[8, 128, 8192] с TPU{0,0} на TPU{3, 3}. Задержка на один переход у TPU v5e — $1\,\mu s$.

1. Через сколько первый байт придёт в пункт назначения?
2. Сколько займёт вся передача?

Ответ

У TPU v5e двумерная связность. Поскольку у нас срез всего 4x4 (без осей размера 16), замыкающих каналов нет. Значит, у целевого чипа два порта, через которые он может получать данные, и у исходного — два порта для отправки. Передать нужно 2 * 8 * 128 * 8192 = 1.7e7 байт. Можно передавать через оба порта одновременно (половину массива — вправо, половину — вниз), так что получаем 2 * 4.5e10 = 9e10 байт в секунду, и передача всего массива займёт около 1.7e7 / 9e10 = 188us (если мы упираемся в пропускную способность). В срезе 4x4 между чипами $(0, 0)$ и $(3, 3)$ шесть переходов, ведь для осей меньше 16 чипов замыкающих каналов нет. Задержка каждого перехода около $1\,\mu s$, поэтому первый байт придёт примерно через 6us, а вся передача займёт около 188 + 6 = 194us: последний байт тоже должен пройти шесть переходов после того, как покинет источник (в общем случае задержка и время передачи складываются, хотя здесь задержка — лишь небольшая поправка).

Вопрос 6 [всё вместе, сложный]. Представьте большую матрицу A: int8[128 * 1024, 128 * 1024], поровну разбитую по срезу TPU v5e 4x4, но выгруженную в DRAM хоста для каждого чипа. Допустим, вы хотите скопировать весь массив на TPU{0, 0} и умножить его на вектор bf16[8, 128 * 1024]. Сколько это займёт? Подсказка: используйте числа выше.

Ответ

Сначала перечислим операции, которые нужно выполнить. Наш массив — около 16 ГБ. По таблице выше хост TPU v5e имеет топологию 4x2, так что в 4x4 — 2 хоста. Значит, поскольку массив разбит поровну, на каждом хосте фактически лежит половина массива, 8 ГБ. Все эти куски нужно скопировать на TPU{0,0}, и есть два варианта:

1. Скопировать по DCN, а затем загрузить весь неразбитый массив по PCIe в HBM.
2. Загрузить разбитые куски на соответствующие TPU, затем собрать их (gather) по ICI и выполнить умножение на TPU{0,0}.

Ясно, что вариант (2) лучше. DCN медленная по сравнению с ICI, и большой массив гораздо выгоднее загружать по многим каналам PCIe, а не по нескольким (8 на хосте 0). На схеме в оригинале (challenge-problem.png) показана часть системы: TPU связаны с соседями по ICI (даже через границу хостов), все TPU связаны со своим хостом (по PCIe), а хосты связаны по DCN. (На самом деле у каждого чипа свой канал PCIe к хосту, хотя для наглядности показан только один.)

Теперь посчитаем, сколько займёт каждая часть:

1. Загрузка по PCIe: мы загружаем куски общим размером 16 ГБ по 16 каналам PCIe, у каждого пропускная способность 1.6e10 байт/с. Значит, это займёт около 63 мс.
2. Копирование по ICI: теперь на каждом TPU 16 ГБ / 16 = 1 ГБ массива. Пропускная способность ICI — 9e10 байт/с на канал в обе стороны, и, как видно из схемы, у TPU{0,0} в этой топологии задействованы только 2 из 4 каналов ICI TPU v5e. Поскольку TPU{0,0} должен получить всего 15 ГБ по 2 осям при 4.5e10 байт/с на канал, время можно оценить снизу как 15e9 / (4.5e10 * 2) = 167ms. На практике это, вероятно, недостижимо, потому что нагрузка очень неравномерна, но, скорее всего, в пределах множителя 2. Как вы увидите в главе 3, полный AllGather тоже занял бы примерно 16e9 / (4.5e10 * 2), так что это близко к оптимуму.
3. Загрузка HBM → MXU: для итогового умножения нужно загрузить эти 16e9 байт плюс массив bf16[8, 128 * 1024] (ещё 2 МБ, пренебрежимо мало) с пропускной способностью HBM в MXU, что займёт 16e9 / 8.2e11 = 20ms.
4. FLOPs: всего выполняется $2 \cdot 8 \cdot 128 \cdot 1024 \cdot 128 \cdot 1024 = 2{,}7 \times 10^{11}$ FLOPs, и при 1.97e14 FLOPs/s в bf16 получаем 1,4 мс.

Верхняя граница общего времени — сумма всех этих времён, но, поскольку TPU обычно умеет совмещать эти операции, можно считать это задачей конвейера, узкое место которого — самая медленная часть. Если так, ответ — не меньше 167 мс, а при неидеальном перекрытии, вероятно, ближе к 200 мс.

Это всё для главы 2. Глава 3 о разбиении (partitioning) и связи между TPU — здесь.

Приложение

Приложение А: подробнее об устройстве TPU

Здесь мы глубже погрузимся во внутренние операции TPU. Если не оговорено иное, характеристики даны для TPU v5p.

VPU

VPU — векторное арифметическое ядро TPU. VPU состоит из двумерной векторной SIMD-машины (собственно VPU), которая выполняет поэлементные арифметические операции вроде vadd (сложение векторов) или vmax (поэлементный максимум), и набора векторных регистров — VREG, — которые хранят данные для VPU и MXU.

VREG. В каждом ядре TPU v5p 64 32-битных VREG (у TPU v4 — 32), что даёт в сумме около 64 * 8 * 128 * 4 = 256kB регистровой памяти на ядро (или вдвое больше на весь чип, ведь ядер два). TPU v5p может загружать из VMEM 3 регистра за такт и записывать в VMEM 1 регистр за такт.

VPU. VPU — двумерный векторный арифметический блок формы (8, 128), где измерение 128 называют осью дорожек (lane), а измерение 8 — осью поддорожек (sublane). Каждая пара (дорожка, поддорожка) у v5 содержит 4 стандартных арифметических блока (ALU) с плавающей точкой, независимых друг от друга. VPU выполняет большинство арифметических инструкций за один такт в каждом ALU (например, vadd — сложение векторов) с задержкой 2 такта, так что, например, у v5 можно за такт сложить 4 пары значений f32 из VREG. Типичная инструкция VPU выглядит как {v2 = vadd.8x128.f32 v0, v1}, где v0 и v1 — входные VREG, а v2 — выходной.

Все дорожки и поддорожки на каждом такте выполняют одну и ту же программу в чистом режиме SIMD, но каждый ALU может выполнять свою операцию. Так что, например, можно за один такт выполнить 1 vadd и 1 vsub, каждый над двумя полными VREG с записью результата в третий.

Быстрый вопрос [пропускная способность VPU]. По данным выше вычислите, сколько векторных FLOPs/s выполняет TPU v5p. Тактовая частота TPU v5p — около 1,75 ГГц.

Ответ

На каждом такте каждое ядро может выполнить 4 векторные инструкции на 8 * 128 ALU. Это даёт 8 * 128 * 4 FLOPs за такт на ядро, или 8 * 128 * 4 * 1.75e9 = 7e12 FLOPs/s. Заметьте, насколько это меньше FLOPs/s у MXU — около 2e14 на ядро (примерно в 30 раз).

Редукции. Как правило, связь или редукция по измерению поддорожек проще, чем по измерению дорожек. Например, VPU поддерживает операцию перемешивания внутри дорожки, которая за такт или около того сдвигает данные по оси размера 8. С её помощью можно эффективно выполнять редукции по измерению поддорожек (просто сдвиги на 4, 2 и 1 и три пары поэлементных сложений).

Редукции между дорожками гораздо труднее: для них нужен отдельный аппаратный блок — XLU, «блок между дорожками» (cross lane unit), — медленный и довольно дорогой.

Сравнение с GPU. Тем, кто знаком с GPU NVIDIA: каждый ALU в VPU аналогичен ядру CUDA, а одна дорожка VPU — «планировщику варпа», то есть набору обычно из 32 ядер CUDA, выполняющих арифметику в режиме SIMD. Редукции внутри дорожки довольно просты, но, если нужно пересечь дорожки, приходится идти как минимум через VMEM/XLU/SMEM, что гораздо медленнее. Подробнее — в главе о GPU.

Скалярное ядро

Скалярное ядро — управляющий блок TPU. Оно выбирает и раздаёт все инструкции, выполняет передачи из HBM в VMEM и может быть запрограммировано на скалярную работу с метаданными. Скалярное ядро однопоточное, поэтому одно из следствий — каждое ядро TPU может создавать лишь один запрос DMA за такт.

Для масштаба: одно скалярное ядро управляет VPU (4096 ALU), 4 MXU, 2 XLU и несколькими движками DMA. Сильный перекос «мало управления на единицу вычислений» — источник аппаратной эффективности, но он же ограничивает возможность хоть сколько-нибудь интересной векторизации, зависящей от данных.

Приложение Б: как работает систолический массив

В сердце MXU у TPU — систолический массив 128x128 (256x256 у TPU v6e). При полной загрузке систолический массив может выполнять одно умножение bf16[8,128] @ bf16[128,128] -> f32[8,128]¹⁰ за 8 тактов.

• По сути, систолический массив — двумерная сетка 128x128 (=16,384) ALU, каждый из которых умеет выполнять умножение и сложение.
• Веса (W, вход 128x128) передаются сверху вниз (их называют RHS, правым операндом), а входы (X, вход 8x128) — слева (их называют LHS, левым операндом).

Анимация в оригинале (systolic-array.gif): упрощённое умножение набора весов (синие) на набор активаций (зелёные). Сначала по диагонали частично загружаются веса (RHS), затем, тоже по диагонали, подаются активации. На каждом кадре перемножаются все совпавшие зелёные и синие клетки, результат складывается с остатком, пришедшим сверху, и передаётся на одну клетку вниз.

Более общий вариант этой анимации (systolic-array2.gif) показывает, как результат потоком выходит из вычисления, а схема systolic-array-pipelining.png — как это конвейеризуется для нескольких массивов RHS и LHS.

Пока загружаются веса (RHS) и активации (LHS), возникает начальный «пузырь» конвейера. После него новые входы и веса можно загружать без дополнительного пузыря.

Есть и «плохая» анимация умножения bf16[2, 3] × bf16[3, 3] (systolic-array-bad.gif) — её можно представить как умножение матрицы весов 2×3 на входные активации с батчем 1 и размером 3. Она повёрнута относительно прошлых слайдов, и входы выходят вправо, а не вниз, но структуру в общих чертах видно.

Так можно эффективно конвейеризовать умножение больших матриц без слишком большого пузыря. Тем не менее важно, чтобы формы матриц были больше стороны MXU — обычно 128×128. В некоторых TPU (начиная с TPU v3) несколько MXU — 2 у TPU v3 или 4 у TPU v4/5, — поэтому нужно следить, чтобы размеры тайлов были больше 128 × число MXU. Вот хорошая анимация на эту тему.

У Trillium (TPU v6e) систолический массив 256x256, то есть он выполняет в 4 раза больше FLOPs за такт. Это также значит, что для полной загрузки MXU размеры тензоров должны быть вдвое больше.

В этом посте есть ещё одна отличная анимация умножения на систолическом массиве для фиксированной матрицы весов.

# Примечание переводчика: систолический массив в миниатюре — каждая «клетка» (i, j) хранит вес W[i, j],
# получает вход слева и частичную сумму сверху, прибавляет x[i] * W[i, j] и передаёт сумму вниз.
import numpy as np

def systolic_matvec(x, W):
    D, F = W.shape
    partial = np.zeros(F)              # частичные суммы, «текущие» сверху вниз по строкам массива
    for i in range(D):                 # строка i массива получает x[i] (в железе — со сдвигом по тактам)
        partial = partial + x[i] * W[i, :]
    return partial                     # внизу массива выходит x @ W

x, W = np.arange(3.0), np.arange(9.0).reshape(3, 3)
print(systolic_matvec(x, W), x @ W)    # [15. 18. 21.] [15. 18. 21.]

Словарик

English	Русский
TensorCore	TensorCore (вычислительное ядро TPU)
MXU (matrix multiply unit)	блок умножения матриц
VPU (vector processing unit)	блок векторной обработки
VMEM, VREG, SMEM	векторная память, векторные регистры, скалярная память
systolic array	систолический массив
megacore	мегаядро
tray, host	плата, хост
ICI (inter-chip interconnect)	межчиповые связи
DCN (data-center network)	сеть дата-центра
torus, wraparound link	тор, замыкающий канал
pod, superpod, slice	под, суперпод, срез
unidirectional / bidirectional bandwidth	односторонняя / двусторонняя пропускная способность
pipelining, pipeline bubble	конвейеризация, пузырь конвейера
lane, sublane	дорожка, поддорожка
prefetching	предзагрузка

Footnotes

1. У TPU v6e (Trillium) MXU размером 256×256, а у всех предыдущих поколений — 128×128. ↩

2. TPU, и особенно их систолические массивы, — такие мощные ускорители потому, что умножение матриц — один из немногих алгоритмов, которые используют $O(n^3)$ вычислений на $O(n^2)$ байт. Поэтому обычному ALU очень легко упереться в вычисления, а не в пропускную способность памяти. ↩

3. Иногда мы говорим о предзагрузке в VMEM (VMEM prefetching) — заблаговременной загрузке весов в VMEM, чтобы скрыть стоимость загрузки для умножений матриц. Например, в обычном трансформере иногда можно загружать большие веса прямого слоя (feed-forward) в VMEM во время внимания и скрыть стоимость загрузки весов, если мы упираемся в пропускную способность памяти. Для этого веса должны быть достаточно малы или достаточно сильно шардированы, чтобы один слой поместился в VMEM с запасом. ↩

4. В облачных TPU VM каждая плата видна как часть отдельной виртуальной машины, так что снова видны 4 ядра. ↩

5. Оптический переключатель — просто переконфигурируемое соединение с той же пропускной способностью ICI. Он лишь позволяет соединять кубы, сохраняя замыкающий канал. ↩

6. Заметьте, что у 2x2x4 замыкающих каналов не будет: их обеспечивают оптические переключатели, доступные только для полного куба. А вот у TPU v5e 8x16 замыкающий канал по длинной оси будет, потому что там не используется переконфигурируемая оптическая сеть. ↩

7. На странице по ссылке указано 100 ГБ/с, что немного отличается от приведённого здесь. Каналы ICI у TPU имеют немного разную пропускную способность в зависимости от выполняемой операции. В целом числами из этого документа можно пользоваться без опаски. ↩

8. У TPU v6e и TPU7x — 12,5e9 байт/с, у v5e — 3,125e9 байт/с. ↩

9. Под двусторонней пропускной способностью мы понимаем общее число байт, которое можно отправить по одному каналу в обоих направлениях, или, что то же самое, общее число исходящих байт с одного TPU вдоль определённой оси при условии, что оба канала используются эффективно. Это верно, когда у нас работающее кольцо, то есть когда по данной оси есть замыкающий канал. Так бывает у инференс-чипов при полной оси 16 или у чипов для обучения (v*p) при оси, кратной 4. Мы предпочитаем двустороннюю пропускную способность, потому что она часто возникает в вычислениях с двусторонней связью. ↩

10. Если вы не знакомы с такой записью: это умножение матрицы 8×128 с элементами bfloat16 на матрицу 128×128 с элементами bfloat16 с сохранением результата в матрице 8×128 с элементами float32. ↩