Google TPU

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

История создания

Разработка TPU началась в 2013 году в ответ на растущие потребности Google в вычислительных мощностях для задач машинного обучения. К тому времени компания уже несколько лет использовала центральные процессоры (CPU) и графические процессоры (GPU) для обучения и выполнения нейросетей, однако их производительность и энергоэффективность не удовлетворяли требованиям растущих сервисов, таких как Google Поиск, Google Фото и Google Translate. Руководство Google приняло решение создать собственную специализированную микросхему, которая бы выполняла операции, характерные для нейронных сетей, с максимальной эффективностью.

Первое поколение TPU (TPUv1) было анонсировано в мае 2016 года на конференции Google I/O. Первоначально процессоры использовались внутри компании с 2015 года для ускорения инференса — выполнения уже обученных моделей. Второе поколение (TPUv2) было представлено в 2017 году и уже поддерживало как инференс, так и обучение нейронных сетей. Третье поколение (TPUv3) вышло в 2018 году, а четвёртое (TPUv4) — в 2021 году. В 2023 году было анонсировано пятое поколение (TPUv5e), а в 2024 году — TPUv5p, которое стало самым производительным на момент анонса.

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

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

TPU построен на архитектуре, принципиально отличающейся от CPU и GPU. Ключевым элементом является матричный умножитель (MXU — Matrix Multiply Unit), который выполняет операции умножения матриц и свёртки — основные вычислительные задачи в нейронных сетях. MXU представляет собой систолический массив, состоящий из тысяч умножителей и сумматоров, работающих параллельно.

В состав TPU входят:

• Блок управления — отвечает за загрузку инструкций и данных, управление потоком выполнения.
• Буфер весов — хранит обученные параметры нейронной сети (веса и смещения).
• Буфер активаций — хранит промежуточные результаты вычислений (активации нейронов).
• Матричный умножитель (MXU) — выполняет основную вычислительную работу.
• Блок активации — применяет функции активации (например, ReLU, сигмоида) к результатам.
• Блок нормализации — выполняет пакетную нормализацию (batch normalization) и другие операции.
• Встроенная память (HBM — High Bandwidth Memory) — высокоскоростная память для хранения данных, используемая в TPUv2 и более новых поколениях.

Поколения TPU

Примечание: производительность указана для одного чипа в режиме обучения с плавающей запятой (BF16). Для TPUv1 указана производительность в режиме инференса с целочисленными данными (INT8).

Применение

Внутреннее использование Google

Основное применение TPU — это ускорение сервисов Google, требующих машинного обучения:

• Google Поиск — ранжирование результатов поиска, обработка запросов.
• Google Translate — машинный перевод текстов.
• Google Фото — распознавание лиц, объектов, сцен на фотографиях.
• Google Assistant — обработка голосовых команд и распознавание речи.
• YouTube — рекомендации видео, модерация контента.
• Waymo — обработка данных с сенсоров беспилотных автомобилей (дочерняя компания Alphabet).

Облачный сервис Google Cloud TPU

С 2017 года Google предоставляет доступ к TPU через облачную платформу Google Cloud. Пользователи могут арендовать вычислительные ресурсы в виде отдельных чипов, целых подов (pod) — кластеров, объединяющих тысячи TPU, соединённых высокоскоростной сетью. Это позволяет исследователям и компаниям обучать крупные нейронные сети, не имея собственного оборудования.

Исследовательские проекты

TPU широко используются в академических и промышленных исследованиях для обучения больших языковых моделей (LLM), моделей компьютерного зрения, генеративных нейросетей и других задач. Например, модель BERT (Bidirectional Encoder Representations from Transformers) была обучена на TPUv3.

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

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

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

Недостатки

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

Сравнение с другими ускорителями

TPU vs GPU

Графические процессоры (GPU) традиционно использовались для ускорения машинного обучения до появления TPU. GPU (например, NVIDIA Tesla V100, A100, H100) также имеют высокую производительность для матричных операций, но они менее специализированы. TPU, как правило, обеспечивает более высокую производительность на ватт и на доллар для задач, оптимизированных под его архитектуру. Однако GPU остаются более универсальными и поддерживаются большим числом программных платформ.

TPU vs CPU

Центральные процессоры (CPU) общего назначения значительно уступают TPU в производительности для задач машинного обучения, особенно для операций с большими матрицами. CPU оптимизированы для последовательных вычислений и широкого спектра задач, в то время как TPU — для параллельных матричных операций.

Критика

Основная критика TPU связана с закрытостью технологии и зависимостью от одного поставщика — Google. Компания не продаёт чипы отдельно, а предоставляет их только через облачный сервис или использует внутри себя. Это ограничивает возможности независимых исследований и разработок. Кроме того, привязка к TensorFlow и экосистеме Google вызывает опасения по поводу вендор-лока (зависимости от одного поставщика). Некоторые эксперты также отмечают, что TPU может быть неоптимален для некоторых типов нейронных сетей, например, для рекуррентных сетей или моделей с нестандартными операциями.

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

• Первое поколение TPU (TPUv1) было способно выполнять 92 триллиона операций в секунду (TOPS) для целочисленных вычислений (INT8), что примерно в 30 раз быстрее, чем современные на тот момент GPU.
• TPUv2 и более новые поколения используют формат BF16 (bfloat16) — 16-битное представление чисел с плавающей запятой, разработанное Google специально для машинного обучения.
• TPUv4p — это первое поколение, в котором используется оптическая межсоединение (Optical Interconnect) для связи между чипами, что позволяет значительно увеличить пропускную способность сети.
• В 2021 году Google объявила, что TPUv4 используются для обучения модели Pathways Language Model (PaLM) — одной из крупнейших языковых моделей на тот момент.
• TPU не продаются в розницу; единственный способ получить доступ к ним — через облачный сервис Google Cloud или через внутреннее использование в Google.

Источники

• Google Cloud TPU documentation (Google LLC — организация признана иноагентом в РФ)
• 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.
• Jouppi, N. P., et al. «Ten Lessons From Three Generations Shaped Google’s TPUv4i.» Proceedings of the 48th Annual International Symposium on Computer Architecture (ISCA), 2021.
• «Tensor Processing Unit (TPU) — Google AI» (Google LLC — организация признана иноагентом в РФ)
• «TPU v5p: A New Generation of Google’s Tensor Processing Units» — Google Cloud Blog, 2024.