В последние годы стало очевидно, что классические центральные процессоры (CPU) и видеокарты (GPU) уже не всегда поспевают за непрерывным ростом и усложнением нейронных сетей. Вместо бесконечного наращивания «универсального» железа, компании начали разрабатывать и внедрять в своих дата-центрах Domain-Specific Architecture (DSA) — аппаратные ускорители, заточенные под конкретные задачи.
Google TPU (Tensor Processing Unit) — одно из первых крупных решений такого рода. Начиная с 2015 года (поколение TPUv1), Google успела вывести на рынок несколько поколений TPU для внутренних нужд: TPUv1 и TPUv2/v3, а в 2020 году — новое решение TPUv4i. Если первые версии TPU были ориентированы исключительно на ускорение инференса (выполнение уже обученных моделей), то TPUv2 и TPUv3 смогли взять на себя ещё и тренировку крупных нейросетей. Но в дальнейшем выяснилось, что для оптимальной работы дата-центров в масштабах Google рациональнее разделить решения для тренировки и инференса. TPUv4i — это результат учёта многих уроков и ограничений, проявившихся в предыдущих чипах.
В этом материале разберём, какие «десять уроков» сформировали подход Google к созданию TPUv4i, что это за архитектура и какие проблемы дата-центров она решает.
Краткая история развития TPU
В TPUv1 (2015 год) главной задачей было выполнение задач, связанных с инференсом статически обученных моделей. На тот момент считалось, что для этого достаточно использовать «жёстко запрограммированный» вычислительный конвейер для матриц, основанный на систолических массивах 8-битных умножений, а веса нейронной сети хранить во внешней памяти DDR3. Управление инструкциями осуществлялось с помощью CPU через PCIe. Такой подход позволил значительно превзойти классические CPU и GPU по энергоэффективности, обеспечивая высокую производительность на ватт.
Однако вскоре выяснилось, что научное сообщество и разработчики внутри Google и за её пределами каждый год усложняют модели, повышая требования к точности вычислений и пропускной способности памяти. В результате появилась TPUv2 (2017), поддерживающая как инференс, так и тренировку. В этой версии использовались более широкая шина памяти (HBM вместо DDR3), собственная векторная логика, управляемая VLIW, и специальная топология межчиповых связей (Inter-Chip Interconnect), что позволяло собирать огромные «суперкомпьютеры» из TPU для обучения больших моделей.
TPUv3 (2018) стал логическим продолжением TPUv2: была увеличена тактовая частота, масштабировано количество умножителей (MXU), а энергопотребление достигло 450 Вт на чип.
Десять уроков, которые предопределили TPUv4i
Урок 1. Логика совершенствуется быстрее, чем память и проводники
В отличие от предыдущих «золотых» эпох масштабирования, современные техпроцессы дают неравномерный прирост для логики, SRAM и пропускной способности памяти. Логические элементы становятся дешевле и энергоэффективнее, особенно при переходе на 7 нм, в то время как проводники и SRAM развиваются гораздо медленнее. Из этого следует, что эффективнее увеличивать количество вычислительных блоков (матриц умножения), чем надеяться на лёгкое масштабирование объёмов быстрой памяти и пропускной способности шин.
Урок 2. Компиляторы оказывают решающее влияние на раскрытие аппаратного потенциала.
Успехи TPU стали возможны благодаря системному стеку и компилятору XLA (Accelerated Linear Algebra). Для архитектуры типа VLIW очень важно оптимизировать производительность через грамотные компиляции. Новая архитектура без готового компилятора часто не раскрывает свой потенциал в первые пару лет. Google сделала ставку на единый стек XLA, который поддерживает несколько высокоуровневых фреймворков (TensorFlow, PyTorch, JAX).
Урок 3. Приоритет совокупной стоимости владения (TCO) важнее фокуса на цене чипа.
Для дата-центра итоговые затраты складываются не только из стоимости ASIC, но и из расходов на электроэнергию, системы охлаждения, распределения питания и так далее. В результате большую роль играет потребляемая мощность (TDP). Иногда чип, который стоит немного дешевле (CapEx), потребляет столько энергии (OpEx), что за три-пять лет эксплуатации становится дороже. Тренд Google — смотреть прежде всего на производительность на ватт и на общую «стоимость владения», а не только на стоимость приобретения чипа.
Урок 4. «Обратная совместимость» с точки зрения машинного обучения
Ключевым оказался принцип, что новый чип для инференса должен численно совпадать с тем, как модель вела себя на обучении. Если TPUv3 обучал в формате bfloat16/float32, то и TPUv4i (при инференсе) должна возвращать совместимые результаты, включая крайние случаи округлений. Только так большие команды могут «безболезненно» перенести новую модель в продакшн. Поэтому в TPUv4i не стали убирать поддержку bfloat16 и fp32 в пользу «чисто int8»: да, это где-то снимает необходимость постквантования и риск потери качества.
Урок 5. Для масштабируемого инференса необходимо воздушное охлаждение.
Глобальные сервисы нужно разворачивать по всему миру, и зачастую в «плотных» стойках, где сложно организовать жидкостное охлаждение (как в TPUv3 для обучения). Поэтому TPUv4i рассчитан на TDP 175 Вт и классические системы охлаждения.
Урок 6. Точность с плавающей запятой критически важна для ряда инференс-задач. Использование квантования с помощью int8 позволяет экономить память и иногда даже ускоряет вычисления. Однако для некоторых приложений, таких как сложная сегментация изображений или работа с моделями, обладающими тонкой лингвистической детализацией, потеря качества после квантования может достигать 1%, что уже является критичным. Кроме того, оптимизация модели с целью уменьшения ошибки после квантования может занять дополнительные месяцы работы. Поэтому, чтобы обеспечить поддержку широкого спектра приложений, чип должен иметь возможность работать с 32-битными или 16-битными (bfloat16) числами с плавающей запятой.
Урок 7. Многозадачность и быстрая смена контекста усиливают гибкость и эффективность чипа.
В реальных условиях на одном и том же ускорителе одновременно работает множество различных нейросетей: от распознавания речи и изображений до рекомендаций и сервисных подсказок. Чтобы эффективно справляться с такими задачами, требуется быстро переключаться между моделями, которые могут весить десятки или даже сотни мегабайт. Использование только CPU и PCIe для загрузки весов было бы слишком медленным. Однако TPUv4i оснащён большой локальной памятью и HBM, что позволяет хранить и быстро переключаться между несколькими моделями.
Урок 8. Ежегодный рост размеров моделей ~1.5× требует запаса по производительности.
Объёмы моделей и сложность инференса растут очень быстро — не меньше, чем по «закону Мура». Если сделать чип, который идеально подойдёт для сегодняшних моделей, он устареет уже через год. В Google предпочитают иметь запас мощности и памяти, чтобы несколько лет запускать обновлённые DNN и не выходить за пределы SLO.
Урок 9. Динамичная эволюция архитектур нейросетей повышает ценность универсальной программируемой логики.
Помимо стремительного увеличения размеров, меняются и типы нейронных сетей. Если раньше преобладали MLP, CNN и LSTM, то уже сейчас популярны различные варианты transformer-архитектуры. Чтобы аппаратная платформа оставалась актуальной, она должна быть универсальной и гибкой, как, например, XLA.
Урок 10. Лимиты по задержке (SLO) определяют стратегию обработки данных сильнее, чем фиксированный batch.
Для реальных сервисов, таких как рекомендации и поиск, важно, чтобы время ответа составляло не менее 99-го процентиля. Можно увеличить размер пакета данных (batch) для повышения эффективности, но если задержки начнут выходить за пределы Service Level Agreement (SLO), модель станет бесполезной. На практике размер пакета данных обычно довольно большой, а «строгий» лимит чаще всего устанавливается по времени.
Особенности архитектуры TPUv4i
Опираясь на накопленный опыт, Google приняла решение отойти от концепции универсального чипа, способного выполнять две задачи одновременно. Вместо этого компания предлагает более специализированный подход, предусматривающий возможность выбора между двумя вариантами: TPUv4 — двухъядерная версия для обучения и TPUv4i — чип с одним блоком вычислений для выполнения инференса. Оба варианта будут основаны на новой технологической базе с использованием 7-нанометровых транзисторов.
-
Логика vs. память: у TPUv4i четыре систолических массива (MXU) на ядро, поддерживающих bfloat16, fp32 и int8. Благодаря тому, что логика «подешевела» при переходе на 7 нм, стало оправдано «умножить» матричные блоки.
-
Расширенная иерархия памяти: помимо традиционного HBM, в TPUv4i добавили крупный Common Memory (CMEM) объёмом 128 MB. CMEM работает быстрее, чем HBM, и это принципиально помогает обрабатывать промежуточные тензоры и быстро переключаться между моделями. В результате ряд задач даёт +1.3–2.2x производительности при использовании CMEM.
-
Воздушное охлаждение, TDP = 175 Вт: сниженная тактовая частота (около 1 ГГц) и умеренный уровень энергопотребления позволяют разворачивать такие чипы во множестве дата-центров по всему миру.
-
Гибкая система DMA: четырёхмерная адресация (4D DMA) ускоряет манипуляции с тензорами и упрощает компилятору размещение данных в иерархии памяти.
-
Упор на производительность на ватт (perf/TDP): по сравнению с TPUv3, хоть в абсолютных FLOPS прирост не гигантский, но в пересчёте на потребляемую мощность TPUv4i показал заметное превосходство (в среднем +2.3x в крупных приложениях Google).
-
Совместимость на уровне компилятора XLA, а не инструкций: поскольку TPU использует VLIW, Google не преследует «побитовую» бинарную совместимость. Главное — чтобы новый чип воспринимался XLA как достаточно похожая целевая платформа, и старые оптимизации компилятора приносили эффект.
Заключение
TPUv4i получился логическим итогом череды решений, связанных с практическими ограничениями эксплуатации больших нейронных сетей в дата-центрах Google. Уроки, извлечённые из эксплуатации TPUv1 и TPUv2/v3, позволили иначе взглянуть на задачу инференса, где важны не только голая производительность (FLOPS), но и:
-
минимизация совокупной стоимости владения (TCO),
-
глобальная масштабируемость,
-
поддержка многозадачности,
-
совместимость с тренировочными моделями (bfloat16, fp32),
-
способность быстро реагировать на рост и усложнение DNN.
Несмотря на то, что TPUv4i уступает TPUv3 в максимальной производительности (FLOPS), по метрике “производительность на ватт” он опережает предшественников. Это важнее для масштабных продакшн-сервисов, так как электроэнергия и системы охлаждения — значительные статьи расходов.
История TPU демонстрирует общую закономерность: DSA-подход эффективен, когда спроектирован под реальные нагрузки и эволюцию нейросетей, а не только ради впечатляющих цифр в одном синтетическом бенчмарке. Будущие версии TPU будут следовать этой стратегии, вбирая в себя как уроки общих технологических ограничений, так и наработки в компиляторах и алгоритмах.
Автор: tensorprogrammer
