TPU (Tensor Processing Unit)

Тензорный процессор (TPU) — это специализированная интегральная схема (ASIC), которую Google разрабатывает для ускорения обучения и инференса нейронных сетей. TPU построен вокруг крупных блоков матричного умножения, поскольку матричное умножение составляет основную часть вычислений при глубоком обучении. Google использует TPU в своих дата-центрах с 2015 года, обучает на них модели Gemini и сдает их в аренду другим компаниям как Cloud TPU. Восьмое поколение, представленное в апреле 2026 года, разделено на два чипа: TPU 8t для обучения и TPU 8i для инференса.

Как работает TPU#

В основе каждого TPU лежит блок матричного умножения (MXU), построенный по принципу систолического массива: это сетка ячеек умножения с накоплением, которые передают данные напрямую соседям. Веса загружаются в ячейки один раз, входные данные проходят по сетке на один шаг за такт, а частичные суммы передаются вниз по столбцам, поэтому каждое значение, полученное из памяти, повторно используется многими ячейками. Чтение памяти требует намного больше времени и энергии, чем арифметические операции, — в этом и заключается эффективность такой конструкции. В Cloud TPU v2 и v3 каждый MXU представлял собой массив 128×128, выполнявший умножение в формате bfloat16 и накопление в FP32. Короткая симуляция показывает поток данных в сетке 4×4:

import numpy as np

rng = np.random.default_rng(0)
W = rng.integers(-5, 5, (4, 4))  # weights, preloaded one per cell of a 4x4 grid
X = rng.integers(-5, 5, (6, 4))  # six input rows stream through the grid
Y = np.zeros((6, 4), dtype=int)

# Cell (k, j) holds W[k, j]. Input X[i, k] enters row k from the left and reaches
# cell (k, j) at cycle i + k + j, where it is multiplied by the stored weight and
# added to the partial sum for Y[i, j] passing down column j
cycles = 6 + 4 + 4 - 2
for t in range(cycles):
    for k in range(4):
        for j in range(4):
            i = t - k - j
            if 0 <= i < 6:
                Y[i, j] += X[i, k] * W[k, j]

print(f"{X.size * 4} multiply-adds in {cycles} cycles, equal to X @ W: {np.array_equal(Y, X @ W)}")

Выводится 96 multiply-adds in 12 cycles, equal to X @ W: True: все 16 ячеек вычисляют точное матричное произведение, при этом каждое значение веса считывается из памяти один раз и повторно используется для всех шести строк входных данных. TPU первого поколения выполнял такие вычисления в сетке 256×256 из 8-битных ячеек, совершая 65 536 операций умножения с накоплением за такт и достигая пиковой производительности 92 триллиона операций в секунду. Его простой и детерминированный механизм выполнения соответствовал цели Google — обеспечить целевое время отклика, не превышаемое в 99% случаев, при выполнении вывода.

TPU выполняют код, скомпилированный с помощью XLA: компилятор получает весь граф вычислений из JAX, TensorFlow или PyTorch и распределяет его по матричным блокам. Чипы соединяются специальными высокоскоростными интерфейсами в поды, насчитывающие тысячи устройств; так крупнейшие базовые модели обучают сразу на множестве TPU.

TPU, GPU и CPU#

• CPU: несколько мощных ядер для обработки ветвлений и кода общего назначения; в большинстве конвейеров обучения они загружают и предварительно обрабатывают данные.
• GPU: тысячи программируемых ядер и тензорные ядра для матричных вычислений. GPU работают практически с любыми моделями и фреймворками, а также доступны у многих поставщиков оборудования и облачных сервисов.
• TPU: в обмен на гибкость обеспечивает эффективность при работе с крупными плотными матричными операциями. Лучше всего он подходит для рабочих нагрузок со статическими формами тензоров, поскольку при изменении формы XLA приходится компилировать код заново; TPU плохо справляется с пользовательскими операциями, интенсивным ветвлением и арифметикой высокой точности.
• NPU: общее название ускорителей нейронных сетей в телефонах, ноутбуках и граничных устройствах; к этой категории относится Edge TPU.

Edge TPU#

Edge TPU — компактный TPU только для вывода, который используется в Google Coral USB Accelerator, модулях M.2 и PCIe, а также Dev Board. Он выполняет 4 триллиона операций в секунду при мощности 2 Вт и поддерживает только 8-битные целочисленные вычисления, поэтому сети полностью преобразуются в INT8 посредством квантования модели, а затем компилируются с помощью компилятора Edge TPU от Google. Компилятор размещает поддерживаемые операции на чипе, а остальные оставляет на хостовом CPU. Ultralytics YOLO автоматизирует оба этапа с помощью format="edgetpu": эта настройка включает квантование INT8 и запускает компилятор на компьютере с x86-64 и Linux, создавая модель для таких устройств, как Raspberry Pi с Coral USB Accelerator:

from ultralytics import YOLO

# Load the smallest YOLO26 detection model
model = YOLO("yolo26n.pt")

# Quantize to INT8 and compile for the Edge TPU (x86-64 Linux only)
path = model.export(format="edgetpu")
print(f"Edge TPU model: {path}")

Теперь платформа Coral сосредоточена на Coral NPU — анонсированной в октябре 2025 года разработке ускорителя на открытой архитектуре RISC-V для носимых устройств, рассчитанной на производительность 512 миллиардов операций в секунду при энергопотреблении в несколько милливатт.