interpres 3 часа назад

Что же такое TPU

Простой

14 мин

1.2K

Блог компании RUVDS.comМашинное обучение*Компьютерное железоПроцессорыИскусственный интеллект

Обзор

Перевод

Автор оригинала: Henry Ko

В последнее время я много работал с TPU, и мне было интересно наблюдать такие сильные различия в их философии дизайна по сравнению с GPU.

Главная сильная сторона TPU — это их масштабируемость. Она достигается благодаря и аппаратной (энергоэффективности и модульности), и программной стороне (компилятору XLA).

Общая информация

Если вкратце, то TPU — это ASIC компании Google, делающий упор на два фактора: огромную производительность перемножения матриц + энергоэффективность.

Их история началась в Google в 2006 году, когда компания впервые начала размышлять о том, что же ей стоит реализовывать: GPU, FPGA или специализированные ASIC. В те времена было лишь несколько областей применения, в которых требовалось специализированное оборудование, поэтому было решено, что потребности компании можно удовлетворить при помощи незадействованных вычислительных ресурсов (compute) CPU её крупных датацентров. Но в 2013 году ситуация изменилась: функция голосового поиска Google начала использовать нейросети, и по расчётам для её реализации потребовалось бы гораздо больше compute.

Перенесёмся в настоящее: сегодня TPU лежат в основе большинства ИИ-сервисов Google. Разумеется, сюда включены обучение и инференс Gemini и Veo, а также развёртывание моделей рекомендаций (DLRM).

Давайте начнём разбирать внутренности TPU с самого нижнего уровня.

TPU на уровне отдельного чипа

В своих схемах я в основном буду рассматривать TPUv4, но эта структура более-менее применима и к TPU нового поколения (то есть TPUv6p Trillium; на момент написания статьи технические характеристики TPUv7 Ironwood ещё не были опубликованы).

Вот, как выглядит структура отдельного чипа TPUv4:

chip diagram — TPU Single Chip + TensorCore

В каждом чипе есть два тензорных ядра (TensorCore) TPU, отвечающих за compute. (Примечание: специализирующиеся на инференсе TPU имеют только одно TensorCore). Оба TensorCore имеют общие блоки памяти: CMEM (128 МиБ) и HBM (32 ГиБ).

Внутри каждого TensorCore находятся блоки compute и буферы памяти меньшего размера:

1. Matrix Multiply Unit (MXU, блок перемножения матриц)
- Это ключевой компонент TensorCore, представляющий собой систолический массив 128x128.
- Подробнее о систолических массивах мы поговорим ниже.
2. Vector Unit (VPU, векторный блок)
- Общие поэлементные операции (например, ReLU, поточечное сложение/умножение, редукции)
3. Vector Memory (VMEM, векторная память; 32 МиБ)
- Буфер памяти. Перед тем, как TensorCore сможет выполнять вычисления, данные из HBM нужно скопировать в VMEM.
4. Scalar Unit + Scalar Memory (SMEM, скалярный блок + скалярная память; 10 МиБ)
- Сообщает VPU и MXU, что им делать.
- Занимается потоком управления, скалярными операциями и генерацией адресов памяти.

Если вы имеете опыт работы с GPU NVIDIA, то вас могут сбить с толку следующие наблюдения:

Блоки памяти на чипе (CMEM, VMEM, SMEM) TPU гораздо больше, чем кэши L1, L2 в GPU.
HBM в TPU гораздо меньше, чем HBM в GPU.
Похоже, за вычисления отвечает гораздо меньшее количество «ядер».

Это полностью противоположно ситуации с GPU, в которых есть маленькие кэши L1, L2 (для H100, соответственно, 256 КБ и 50 МБ), HBM большего размера (80 ГБ в случае H100) и десятки тысяч ядер.

Прежде чем двигаться дальше, вспомним, что TPU способны, как и GPU, обеспечивать чрезвычайно высокую производительность. TPU v5p может достигать 500 терафлопс/с на чип, а при полном поде из 8960 чипов мы можем достигать приблизительно 4,45 экзафлопс/с. Заявляется, что новые TPUv7 Ironwood будут достигать до 42,5 экзафлопс/с на каждый под (9216 чипов).

Чтобы понять, как у TPU это получается, надо разобраться в философии их дизайна.

Философия дизайна TPU

TPU обеспечивают потрясающую пропускную способность и энергоэффективность благодаря двум главным столпам и ключевому допущению: систолическим массивам + конвейерам, компиляции Ahead-of-Time (AoT) и допущению, что большинство операций можно выразить способом, хорошо отображаемым в систолические массивы. К счастью, в нашу современную эпоху глубокого обучения перемножение матриц составляет подавляющую часть вычислений.

Архитектурное решение TPU №1: систолические массивы + конвейеры

Вопрос: что такое систолический массив?

Систолический массив — это аппаратная архитектура, состоящая из сетки взаимосвязанных обрабатывающих элементов (processing element, PE). Каждый PE выполняет небольшие вычисления (например, умножение и суммирование), а затем передаёт результаты соседним PE.

systolic array diagram — Схема систолического массива

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

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

Более того, существуют очевидные возможности связывания в конвейере вычислений с перемещением данных. Вот схема конвейерной поточечной операции в TPU.

Примечание: слабая сторона систолических массивов — разреженность

Из схемы видно, что систолические массивы любят плотные матрицы (то есть когда каждый PE активен почти в каждом цикле). Однако их недостаток заключается в том, что в случае работы с разреженными матрицами того же размера нет никаких улучшений производительности: необходимо выполнять то же самое количество тактов, пусть даже PE работают с элементами, имеющими нулевые значения.

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

Архитектурное решение TPU №2: компиляция Ahead of Time (AoT) + меньшая зависимость от кэшей

В этом разделе мы поговорим о том, как TPU достигают высокой энергоэффективности благодаря избеганию кэшей при помощи аппаратной и программной связки TPU + компилятора XLA.

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

Однако доступ к кэшу (и доступ к памяти в целом) требует существенной энергии. Ниже представлена приблизительная оценка энергозатратности операций в чипе (45 нм, 0,9 В; [18]). Из этого можно сделать вывод, что доступ к памяти и управление тратят основную часть энергии, а арифметика занимает существенно меньше.

Но что, если наша область применения специфична, а её паттерны вычислений/доступа к памяти хорошо прогнозируются?

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

Именно к этому стремится философия TPU и именно поэтому TPU проектируются параллельно с компилятором XLA. Компилятор XLA предварительно анализирует графы вычислений и генерирует оптимизированные программы.

Вопрос: но JAX тоже хорошо работает с TPU, хотя там используется @jit?

JAX+XLA в TPU — это гибридное пространство JIT и AOT, отсюда и берётся недопонимание. Когда мы впервые вызываем jit-функцию в JAX, то JAX трассирует её, чтобы создать статический граф вычислений. Этот граф передаётся компилятору XLA, где превращается в полностью статичный двоичный файл для TPU. Чтобы адаптировать процесс под TPU, специфичные для TPU оптимизации (например, минимизация количества операций доступа к памяти) происходит на последнем этапе преобразования.

Но здесь нужно быть аккуратными: jit-функции необходимо компилировать и кэшировать повторно, если они запускаются с другими формами входных данных. Именно поэтому JAX плохо проявляет себя при любом динамическом заполнении или слоях с циклами for разной длины, зависящими от входных данных.

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

Но почему же Google всё равно стремится к реализации этой философии дизайна?

TPU и энергоэффективность (на примере TPUv4)

Предыдущая диаграмма энергопотребления неточно отражает показатели для TPU, поэтому покажем разбивку энергопотребления для TPUv4. Стоит отметить, что TPUv4 изготовлены по техпроцессу 7 нм, а 45 нм указаны здесь для сравнения ([3], [16]).

В столбчатой диаграмме показаны значения, но стоит отметить, что в современных чипах применяется HBM3, которая потребляет гораздо меньше энергии, чем показанные выше DRAM DDR3/4. Тем не менее, из диаграммы видно, что операции с памятью требуют на пару порядков величин больше энергии.

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

TPU на уровне набора чипов

Давайте сделаем шаг вверх по лестнице и посмотрим, как TPU работают в системах из набора чипов.

Уровень трея (или «платы»; 4 чипа)

Отдельный трей (Tray) TPU состоит из 4 чипов TPU, или 8 TensorCore (называемых просто ядрами). И у каждого трея есть собственный CPU-хост (примечание: в случае TPU инференса один хост выполняет доступ к двум треям, поскольку у них всего по одному ядру на чип).

Связь хост⇔чип выполняется по PCIe, но связь чип⇔чип реализуется через Inter-Core Interconnect (ICI), имеющую более широкую полосу пропускания.

Однако связи ICI распространяются и дальше, на группы треев. И чтобы разобраться с ними, нам нужно подняться на уровень стоек (Rack).

Уровень стоек (4x4x4 чипов)

Самое удивительное в TPU — это их масштабируемость, которую мы можем начать наблюдать с уровня стоек.

Стойка TPU состоит из 64 TPU, соединённых в 3D-тор 4x4x4. Вот изображение 8 стоек TPU из рекламных материалов Google.

Но прежде, чем мы приступим к обсуждению стоек, нам нужно разобраться в этой запутанной терминологии: в чём различие между стойками (rack), подами (pod) и срезами (slice).

Вопрос: в чём различия между «стойкой TPU», «подом TPU» и «срезом TPU»?

В разных источниках Google использование этих терминов немного разнится, и иногда «поды» и «срезы» обозначают одно и то же. В этой статье мы будем придерживаться определений из научных статей Google о TPU и из документации GCP по TPU ([3], [7], [9]).

Стойка TPU: физическое устройство, содержащее 64 чипа. Также называется «кубом» (cube).
Под TPU: максимальный блок TPU, который можно соединить при помощи ICI и оптоволокна. Часто также называется «Superpod» или «Full pod». Например, под TPU в случае TPUv4 состоял бы из 4096 чипов, или 64 стоек TPU.
Срез TPU: любая конфигурация TPU в промежутке от 4 чипов и до размера Superpod.

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

Для начала давайте поработаем с физическими единицами: стойками и подами. Понимание способов физического соединения систем TPU позволит нам лучше понять философию дизайна TPU.

Вернёмся к стойкам TPU (на примере TPUv4):

Одна стойка TPU состоит из 64 чипов, связанных друг с другом ICI и Optical Circuit Switching (OCS). По сути, мы соединяем несколько треев, чтобы имитировать систему из 64 чипов. Позже мы вернёмся к этой теме объединения маленьких частей для создания суперкомпьютера.

Ниже показана схема стойки TPUv4. Это 3D-тор размером 4x4x4, в котором каждый узел — это чип, синими стрелками показаны ICI, а линии на гранях — это OCS (перерисовано из [7]).

Однако при изучении этой схемы возникает пара вопросов. Почему OCS используется только на гранях? Иными словами, в чём преимущества применения OCS? Есть три важных преимущества, а чуть позже мы рассмотрим и ещё два.

Преимущество OCS №1: цикличность

Повышение скорости коммуникаций между узлами благодаря цикличности.

OCS используется для зацикливания конфигурации TPU. Это снижает наихудшее количество переходов между двумя узлами с N-1 до (N-1)/2 на ось, поскольку каждая ось превращается в кольцо (1D-тор).

Этот эффект становится важнее при увеличении масштабов, потому что снижение задержек коммуникаций между чипами крайне важно для высокой параллелизации.

Примечание: не все TPU имеют топологии 3D-тора

Старые поколения TPU (TPUv2, v3) и TPU инференса (TPUv5e, TPUv6e) имеют топологию 2D-тора, а не 3D-тора. Однако TPUv7 Ironwood, похоже, имеет топологию 3D-тора, хоть и рекламируется, как чип инференса (Примечание: это лишь моё предположение на основании рекламных материалов).

Уровень Full Pod (Superpod; 4096 чипов в случае TPUv4)

Точно так же, как мы соединяли чипы для создания стойки TPU, можно соединить несколько стоек для создания одного большого Superpod.

Суперподом (Superpod) также называют максимальную конфигурацию взаимосвязанных (при помощи одних только ICI и OCS) чипов, достижимую для TPU. Выше идёт уровень мультиподов, но он реализуется через более медленные соединения; его мы рассмотрим позже.

Максимальный размер зависит от поколения; для TPUv4 это 4096 чипов (то есть 64 стоек по 4x4x4 чипов). Для самого нового TPUv7 Ironwood это 9216 чипов.

На схеме ниже показан суперпод TPUv4.

Каждый куб (то есть стойка TPU) связан с другими через OCS. Это позволяет нам брать срезы TPU в поде.

Срезы TPU с OCS

Можно приобретать подмножества TPU внутри пода, называемые срезами TPU. Но даже если вам нужно конкретное количество чипов N, нужно выбрать одну из множества топологий.

Допустим, вам нужно суммарно 512 чипов. Можно приобрести куб (8x8x8), сигару (4x4x32) или прямоугольник (4x8x16). Выбор топологии среза — это само по себе гиперпараметр.

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

Например, куб (скажем, 8x8x8) будет предпочтителен для коммуникаций вида «все со всеми», используемых при параллелизме данных или тензоров, потому что он имеет наибольшую полосу пропускания в сечении. Сигарообразная форма (скажем, 4x4x32) лучше подходит для параллелизма конвейеров, потому что она повышает скорость коммуникаций последовательных слоёв (если один слой умещается в подсрез размером 4x4 чипа).

Но, разумеется, оптимальность топологии будет зависеть от модели; правильный подбор её — это отдельная работа. В научной статье о TPUv4 [9] тоже говорится об этом и демонстрируется, как топология может увеличивать пропускную способность (примечание: я не знаю точно, какая конкретно архитектура LLM имеется в виду в первой строке, потому что она нигде не была указана).

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

Благодаря OCS эти срезы не обязаны быть смежными стойками. Это второе преимущество OCS (и, вероятно, самое важное).

Преимущество OCS №2: несмежные (переконфигурируемые) срезы из множественных узлов

Стоит отметить, что это отличается от связывания нескольких узлов для имитации несмежных срезов. Так как OCS — это коммутатор, а не жёсткая связь, между узлами гораздо меньше физических проводов, что обеспечивает повышенную масштабируемость (то есть бóльшие размеры подов TPU).

Это позволяет в больших масштабах использовать гибкие конфигурации узлов. Допустим, у нас есть три задачи, которые мы хотим выполнять в одном поде. Это можно реализовать и при помощи наивного планирования, но соединения OCS позволяют нам абстрагировать позицию узлов и рассматривать весь под, как «мешок с узлами» (перерисовано из [6]).

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

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

Преимущество OCS №3: перекрученные топологии TPU

Ранее мы видели, что можно создавать различные топологии срезов TPU, меняя размерности (x,y,z) чипов. Однако на этот раз мы будем работать с фиксированными размерностями (x,y,z), меняя их топологии при помощи соединений.

Ниже показан пример перехода от обычного тора сигарообразной формы к перекрученному сигарообразному тору.

В случае перекрученного тора мы повышаем скорость коммуникаций между чипами на перекрученной 2D-плоскости. Это особенно полезно для ускорения связи «все со всеми».

Давайте ещё немного углубимся и представим конкретный сценарий, где это было бы полезно.

Ускорение обучения благодаря перекрученному тору

Теоретически, перекручивание тора обеспечит наибольшую выгоду при параллелизации тензоров (tensor parallel, TP), потому что в каждом слое есть множество операций all-gather и reduce-scatter. Для параллелизации данных (data parallel, DP) выигрыш будет менее существенным, потому что на каждом этапе обучения тоже есть операции all-reduce, но они встречаются реже.

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

Сценарий 1: топология 4x4x16 (TP + PP; всего 256 чипов)

Наша ось z будет размерностью параллелизма конвейеров (pipeline parallel, PP), а 2D-размерность TP будет иметь вид 4x4. По сути, мы предполагаем, что каждый слой k находится на z=k и каждый слой разделён на 16 чипов. Если на схеме не указаны соединения, то используются стандартные соединения OCS (то есть связь с ближайшим соседом).

Мы перекрутим 2D-тор в каждом z=k, что ускорит коммуникации между чипами в каждом слое TP. Перекручивание вдоль размерности PP необязательно, потому что в основном она зависит от коммуникаций «точка-точка».

Примечание: в реальности перекручивание тора приносит выгоду, когда количество чипов больше 4x4. Мы используем пример с 4x4 только для наглядности.

Сценарий 2: топология 16x4x16 (DP + TP + PP; всего 1024 чипа)

Расширим предыдущий сценарий, добавив размерность DP из четырёх чипов. Это значит, что у нас будет четыре сценария 1 вдоль оси x.

Обратите внимание, что перекручивание тора ограничено только каждой размерностью TP в каждой модели DP (то есть 2D-плоскостью 4x4 для каждого z=k, где k=1…16). Зацикливание используется только для размерности DP, из-за чего каждая строка становится горизонтальным кольцом размера 16.

Можно заметить, что существует и альтернативная топология 8x8x16 (то есть размерность DP 2x2), но при смешении размерностей DP и TP всё становится сложнее. В частности, будет непонятно, как нам организовывать зацикленность OCS для оси y, одновременно уместив перекрученные торы для каждой размерности TP.

Уровень мультиподов (или «мультисрезов»; более 4096 чипов в случае TPUv4)

Последний уровень иерархии TPU — это уровень множественных подов (мультиподов). На нём мы работаем со множеством подов как с одной большой машиной. Однако коммуникации между подами происходят через Data-Center Network (DCN), обладающую меньшей полосой пропускания по сравнению с ICI.

Так обучалась PaLM. Для обучения 6144 TPUv4s (2 пода) понадобилось 56 дней. Ниже показано распределение задач TPU по 6 подам: зелёные относятся к PaLM, у красных отсутствуют задачи, остальное — это другие задачи. Каждый квадрат — это куб из TPU 4x4x4.

Создание подобной системы — само по себе непростая задача, но она покажется ещё более впечатляющей, если вспомнить об удобстве для разработчиков. По сути, проектировщикам важнее всего было ответить на вопрос «Как максимально абстрагировать аппаратную/системную часть масштабирования модели?».

Google ответила на него следующим образом: компилятор XLA отвечает за крупномасштабную координацию коммуникаций между чипами. При задании исследователями нужных флагов (например, размерностей параллелизма для DP, FSDP, TP и так далее), компилятор XLA вставляет подходящие иерархические блоки для имеющейся топологии TPU (Xu et al, 2021: GSPMD [2]). Цель этого заключается в том, чтобы крупномасштабное обучение было возможно при внесении в код как можно меньшего количества изменений.

Например, вот разбивка операции all-reduce среди множественных срезов из блога Google [1].

Она демонстрирует, что компилятор XLA берёт на себя заботу о коллективах коммуникаций как между срезами, так и внутри них.

Вот конкретный пример: такой может быть топология TPU для обучения модели. Коммуникации активации внутри среза происходят через ICI, а коммуникации градиентов между срезами происходят через DCN (то есть по размерности DP DCN) [1].

Соотносим диаграммы с реальностью

Мне кажется, осознать масштабы проще, когда смотришь на реальные фотографии оборудования.

Если вы изучали рекламные материалы Google, посвящённые TPU, то могли встречать такое изображение.

Это 8 подов TPU, где каждый блок — это 3D-тор 4x4x4. Каждая строка в поде состоит из 2 треев, то есть в каждой строке по 8 чипов TPU.

Вот отдельный трей TPUv4:

Стоит отметить, что изображение упрощено и на нём показан только один порт PCIe, однако в реальном трее 4 порта PCIe (слева) — по одному на каждый TPU.

А вот отдельный чип:

Центральная часть — это ASIC, а 4 блока рядом — это стеки HBM. На фото показан TPU v4, так что внутри находятся 2 ядра TensorCore, отсюда и 4 стека HBM.

Мне не удалось найти схему чипа TPUv4, поэтому покажу ниже схему TPUv4i; он похож на TPUv4, но имеет лишь одно TensorCore, потому что это чип инференса [3].