TPU

TPU (от англ. Tensor Processing Unit — тензорный процессор) — это специализированная интегральная схема (ASIC, специализированная интегральная схема), разработанная компанией Google для ускорения операций машинного обучения, в первую очередь — для обучения и выполнения логических выводов (инференса) нейронных сетей на основе библиотеки TensorFlow. В отличие от центральных процессоров (CPU) и графических процессоров (GPU), TPU спроектированы для выполнения большого количества параллельных матричных операций, характерных для тензорных вычислений, что обеспечивает высокую производительность и энергоэффективность при решении задач искусственного интеллекта.

История

Разработка TPU началась в 2013 году в связи с необходимостью масштабирования вычислительных мощностей для внутренних сервисов Google, таких как поиск, Google Translate, Google Photos и Google Assistant, которые активно используют нейронные сети. Первое поколение TPU (TPU v1) было представлено в 2016 году на конференции Google I/O. Изначально процессоры предназначались исключительно для выполнения логических выводов (инференса) и не поддерживали обучение моделей. TPU v1 были установлены в серверных стойках Google и использовались для ускорения работы сервисов, таких как RankBrain (система ранжирования поиска) и обработка изображений.

В 2017 году Google анонсировала второе поколение — TPU v2, которое уже могло использоваться как для обучения, так и для инференса. Эти процессоры были интегрированы в облачную платформу Google Cloud Platform (GCP) и стали доступны сторонним разработчикам через сервис Cloud TPU. В 2018 году вышло третье поколение — TPU v3, отличавшееся повышенной производительностью (в 2,7 раза выше по сравнению с TPU v2) и увеличенным объёмом памяти HBM (High Bandwidth Memory).

В 2021 году Google представила TPU v4, которые использовались для обучения моделей-гигантов, таких как PaLM (Pathways Language Model) с 540 миллиардами параметров. TPU v4 стали первыми процессорами, объединёнными в оптическую сеть с топологией «трёхмерный тор», что позволило значительно ускорить межпроцессорное взаимодействие. В 2023 году было анонсировано пятое поколение — TPU v5e, ориентированное на эффективность и снижение стоимости вычислений, а также TPU v5p — флагманская версия для самых ресурсоёмких задач.

Архитектура и устройство

Основные компоненты

TPU представляет собой ASIC, оптимизированный для выполнения операций умножения и накопления (MAC, multiply-accumulate), которые составляют основу свёрточных и полносвязных слоёв нейронных сетей. Ключевые элементы архитектуры:

• Матричный умножитель (Systolic Array) — массив из тысяч вычислительных ядер, расположенных в виде решётки. Данные (входные активации и веса) подаются на входы массива, и за один такт выполняется умножение матриц большого размера (например, 128×128 для TPU v1).
• Буферы весов (Weight Buffer) — локальная память, в которой хранятся параметры модели (веса нейронов). Загрузка весов из внешней памяти осуществляется однократно перед началом вычислений, что снижает задержки.
• Аккумуляторы (Accumulators) — модули для накопления результатов частичных произведений. Для предотвращения потери точности аккумуляторы используют 32-битные числа с плавающей запятой, в то время как сами вычисления могут выполняться в формате bfloat16 (Brain Floating Point 16) или INT8.
• Контроллер (Control Unit) — управляет потоком данных и последовательностью операций, синхронизируя работу матричного умножителя и памяти.

Память и интерконнект

TPU используют высокоскоростную память HBM (High Bandwidth Memory) с пропускной способностью до 1,6 ТБ/с (TPU v3) и более. Для связи между несколькими TPU в кластере применяется специализированная сеть — Interconnect, которая в TPU v4 и v5 реализована на основе оптических каналов с топологией «трёхмерный тор». Это позволяет объединять до 4096 процессоров в единую вычислительную систему.

Программная модель

TPU работают под управлением программного стека TensorFlow (и, частично, JAX и PyTorch). Разработчик описывает модель на высоком уровне, а компилятор XLA (Accelerated Linear Algebra) преобразует граф вычислений в машинные инструкции для TPU. Операции, не поддерживаемые TPU (например, рекуррентные вычисления или операции с динамическими размерами), выполняются на CPU.

Поколения и характеристики

Применение

Облачные вычисления

Основное применение TPU — облачные вычисления через сервис Google Cloud TPU. Пользователи арендуют кластеры TPU («TPU Pods») для обучения больших моделей, таких как GPT-4, BERT, T5, Stable Diffusion и др. TPU обеспечивают высокую скорость обучения: например, модель BERT-Large (340 миллионов параметров) обучается на 16 TPU v3 за 4 дня, тогда как на 8 GPU NVIDIA V100 — за 8 дней.

Внутренние сервисы Google

TPU используются в ядре многих продуктов Google:

• Поиск Google — обработка запросов с помощью нейронных сетей (RankBrain, BERT).
• Google Translate — машинный перевод с использованием моделей Transformer.
• Google Photos — распознавание объектов и лиц, автоматическая группировка снимков.
• YouTube — рекомендации видео, модерация контента.
• DeepMind — обучение моделей для AlphaFold (предсказание структуры белков) и AlphaGo.

Научные исследования

TPU применяются в академических и исследовательских проектах, например, в проекте Folding@home (распределённые вычисления для моделирования сворачивания белков) и в задачах климатического моделирования.

Преимущества и недостатки

Преимущества

• Высокая производительность — TPU превосходят GPU по скорости выполнения матричных операций на ватт потребляемой энергии.
• Энергоэффективность — TPU потребляют меньше энергии на одну операцию по сравнению с GPU общего назначения.
• Масштабируемость — возможность объединения тысяч процессоров в кластеры с низкой задержкой.
• Интеграция с TensorFlow — оптимизированная программная экосистема.

Недостатки

• Узкая специализация — TPU эффективны только для задач, сводимых к матричным операциям (нейронные сети). Для других типов вычислений (например, рекуррентные сети, графовые алгоритмы) они бесполезны.
• Привязка к экосистеме Google — TPU доступны только через Google Cloud, что создаёт vendor lock-in (зависимость от поставщика).
• Высокая стоимость — аренда TPU Pods обходится дороже, чем использование GPU в некоторых сценариях.
• Ограниченная поддержка фреймворков — хотя PyTorch и JAX частично поддерживаются, основная платформа — TensorFlow.

Критика

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

Будущее развитие

Google продолжает совершенствовать TPU, делая акцент на снижение энергопотребления и увеличение производительности. Ожидается, что будущие поколения (TPU v6 и далее) будут использовать более продвинутые технологии, такие как 3D-компоновка чипов (chiplet) и квантовые интерконнекты. Также ведутся работы по интеграции TPU с другими ускорителями (например, для обработки графов и рекуррентных сетей).

Источники

• Jouppi, N. P. et al. «In-Datacenter Performance Analysis of a Tensor Processing Unit». Proceedings of the 44th Annual International Symposium on Computer Architecture (ISCA), 2017.
• Google Cloud TPU Documentation. «Cloud TPU Overview». Google, 2024.
• Jouppi, N. P. et al. «TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings». Proceedings of the 50th Annual International Symposium on Computer Architecture (ISCA), 2023.
• Dean, J. «The Deep Learning Revolution and Its Implications for Computer Architecture and Chip Design». IEEE Micro, 2018.
• «Tensor Processing Unit (TPU)». Wikipedia, 2024.