Tensor Processing Unit

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

История

Предпосылки создания

К середине 2010-х годов компания Google столкнулась с резким ростом вычислительных потребностей своих сервисов, основанных на машинном обучении. Такие продукты, как Поиск Google, Google Фото, Google Ассистент и Google Translate, использовали всё более сложные нейронные сети. Использование стандартных CPU и GPU (например, от NVIDIA) для обучения и особенно для инференса (обработки запросов пользователей в реальном времени) приводило к высоким затратам на электроэнергию и охлаждение, а также к необходимости масштабирования дата-центров.

Первое поколение (TPUv1)

Первое поколение TPU было анонсировано в мае 2016 года на конференции Google I/O. Разработка велась с 2013 года. TPUv1 был спроектирован исключительно для инференса (вывода) нейронных сетей, а не для их обучения. Он использовал 8-битную целочисленную арифметику (INT8) для ускорения вычислений. Чип был установлен на специальной плате, подключаемой к серверам через интерфейс PCIe. По заявлениям Google, TPUv1 был в 15-30 раз быстрее и в 30-80 раз энергоэффективнее современных ему GPU (NVIDIA K80) при выполнении задач инференса.

Второе поколение (TPUv2)

В 2017 году Google представила TPUv2, который стал первым поколением, способным как к обучению, так и к инференсу. Он использовал плавающую запятую (bfloat16) и был оптимизирован для обучения больших моделей. TPUv2 поставлялся в виде «подов» (pods) — кластеров из 64 чипов, объединённых высокоскоростной сетью.

Третье поколение (TPUv3)

В 2018 году было анонсировано третье поколение — TPUv3. Оно обеспечило удвоение производительности по сравнению с TPUv2 за счёт увеличения тактовой частоты и использования жидкостного охлаждения. TPUv3 также поставлялся в виде подов (до 1024 чипов).

Четвёртое поколение (TPUv4)

В 2021 году Google объявила о TPUv4, который стал первым поколением, использующим оптическую коммутацию (OCS — Optical Circuit Switching) для соединения чипов внутри пода. Это позволило создавать более крупные и гибкие кластеры. TPUv4 также продемонстрировал значительный прирост производительности при обучении больших языковых моделей.

Пятое поколение (TPUv5e и TPUv5p)

В 2023 году Google представила TPUv5e (для задач с ограниченным бюджетом) и TPUv5p (флагманское поколение). TPUv5p обеспечил двукратное улучшение производительности по сравнению с TPUv4 и был оптимизирован для обучения моделей с сотнями миллиардов параметров.

Шестое поколение (Trillium)

В 2024 году было анонсировано шестое поколение TPU под кодовым названием Trillium. Оно обещает улучшение производительности в 4,7 раза по сравнению с TPUv5e и повышение энергоэффективности на 67%.

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

Матричный умножитель (Systolic Array)

Ключевой элемент архитектуры TPU — это систолический массив (systolic array). Это матрица из вычислительных ядер, которые выполняют операцию умножения и накопления (MAC — Multiply-Accumulate) за один такт. Данные (веса и активации) «протекают» через массив синхронно, что позволяет выполнять огромное количество параллельных операций. В TPUv1 массив имел размер 256×256, что позволяло выполнять до 65 536 операций MAC за такт.

Память

TPU имеют собственную высокоскоростную память (HBM — High Bandwidth Memory), которая расположена непосредственно на кристалле или рядом с ним. Это позволяет минимизировать задержки при передаче данных между памятью и вычислительными ядрами. Объём памяти в разных поколениях варьируется от 8 ГБ (TPUv1) до 95 ГБ (TPUv5p).

Программное обеспечение

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

Классификация и поколения

Применение

Обучение больших моделей

TPU активно используются Google для обучения своих самых крупных моделей, включая:

• BERT (Bidirectional Encoder Representations from Transformers) — модель для понимания естественного языка.
• PaLM (Pathways Language Model) — большая языковая модель с 540 миллиардами параметров.
• Gemini — мультимодальная модель, обученная на TPUv5p.
• AlphaFold — модель для предсказания структуры белков.

Инференс в реальном времени

TPUv1 и последующие поколения используются для обработки запросов пользователей в таких сервисах, как:

• Google Поиск — ранжирование результатов поиска.
• Google Фото — распознавание лиц и объектов.
• Google Translate — машинный перевод.
• YouTube — рекомендации видео.

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

Google Cloud Platform (GCP) предоставляет доступ к TPU в виде облачных ресурсов. Пользователи могут арендовать отдельные чипы или целые поды для обучения собственных моделей. Это позволяет компаниям и исследователям использовать специализированное оборудование без необходимости его покупки.

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

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

• Высокая производительность на ватт: TPU значительно энергоэффективнее GPU при выполнении матричных операций.
• Специализация: Оптимизация под конкретные задачи машинного обучения позволяет достичь максимальной скорости.
• Масштабируемость: Возможность объединения тысяч чипов в кластеры для обучения сверхбольших моделей.
• Программная экосистема: Глубокая интеграция с TensorFlow и JAX упрощает разработку.

Недостатки

• Узкая специализация: TPU неэффективны для задач, не связанных с машинным обучением (например, рендеринг, научные вычисления общего назначения).
• Привязка к экосистеме Google: Наиболее эффективно работают с фреймворками Google (TensorFlow, JAX). Поддержка PyTorch менее зрелая.
• Отсутствие на рынке: TPU не продаются отдельно, а доступны только через облачную платформу Google Cloud. Это ограничивает их использование для широкого круга разработчиков.
• Сложность программирования: Требуют глубокого понимания архитектуры и оптимизации кода для достижения максимальной производительности.

Критика и альтернативы

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

Альтернативами TPU являются:

• GPU (NVIDIA, AMD): Более универсальны, широко доступны, имеют зрелую экосистему (CUDA, ROCm).
• Другие ASIC: Например, Intel Habana Gaudi, Amazon Trainium, Cerebras Wafer-Scale Engine.
• FPGA (Xilinx, Intel): Программируемые логические интегральные схемы, которые можно настроить под конкретные задачи.

Интересные факты

• TPUv1 был разработан за 15 месяцев, что является рекордно коротким сроком для ASIC такого уровня сложности.
• Название «Tensor Processing Unit» происходит от термина «тензор» — математического объекта, который является основным типом данных в TensorFlow.
• Google использует TPU не только для своих сервисов, но и для внутренних исследовательских проектов, таких как квантовые вычисления (Sycamore).

Источники

• 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.
• "Google's Tensor Processing Units (TPUs) — A Deep Dive." The Next Platform, 2023.
• "TPU v5p: Google's latest AI accelerator." Google AI Blog, 2023.
• "Trillium: Google's sixth-generation TPU." Google Cloud Blog, 2024.