TPU (Tensor Processing Unit)

Une unité de traitement tensoriel (TPU) est un circuit intégré à application spécifique (ASIC) conçu par Google pour accélérer l’entraînement et l’inférence de réseaux neuronaux. Une TPU s’articule autour de grandes unités de multiplication matricielle, car la multiplication matricielle représente l’essentiel du travail en apprentissage profond. Google utilise des TPU dans ses centres de données depuis 2015, entraîne ses modèles Gemini dessus et les loue à d’autres entreprises sous le nom de Cloud TPU. La huitième génération, annoncée en avril 2026, se décline en deux puces : TPU 8t pour l’entraînement et TPU 8i pour l’inférence.

Fonctionnement d’une TPU#

Le cœur de chaque TPU est une unité de multiplication matricielle (MXU) organisée en réseau systolique, c’est-à-dire une grille de cellules de multiplication-accumulation qui transmettent directement les données à leurs voisines. Les poids sont chargés une fois dans les cellules, les entrées parcourent la grille à raison d’une étape par cycle d’horloge et les sommes partielles descendent le long des colonnes ; chaque valeur lue en mémoire est ainsi réutilisée par de nombreuses cellules. La lecture en mémoire exige bien plus de temps et d’énergie que les calculs, ce qui explique le gain d’efficacité de cette conception. Dans les TPU v2 et v3 de Cloud TPU, chaque MXU était une matrice de 128×128 qui multipliait en bfloat16 et accumulait en FP32. Une courte simulation montre le flux des données sur une grille de 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)}")

Le programme affiche 96 multiply-adds in 12 cycles, equal to X @ W: True : les 16 cellules produisent le produit matriciel exact, alors que chaque poids est lu une seule fois en mémoire et réutilisé pour les six lignes d’entrée. La TPU de première génération appliquait ce principe sur une grille de cellules 8 bits de 256×256, avec 65 536 opérations de multiplication-accumulation par cycle et un pic de 92 billions d’opérations par seconde. Son exécution simple et déterministe répondait à l’objectif de Google : respecter une limite de temps de réponse au 99e centile pour l’inférence.

Les TPU exécutent du code compilé par XLA, qui prend l’ensemble du graphe de calcul issu de JAX, TensorFlow ou PyTorch et le répartit sur les unités matricielles. Des interconnexions dédiées relient les puces en pods de plusieurs milliers d’unités, ce qui permet d’entraîner simultanément les plus grands modèles fondamentaux sur de nombreuses TPU.

TPU, GPU et CPU#

• CPU : Quelques cœurs puissants qui gèrent les branchements et le code généraliste, ainsi que le chargement et le prétraitement des données dans la plupart des pipelines d’entraînement.
• GPU : Des milliers de cœurs programmables et des cœurs Tensor dédiés aux calculs matriciels. Les GPU exécutent presque tous les modèles et frameworks et sont proposés par de nombreux fournisseurs et services cloud.
• TPU : Elle échange cette flexibilité contre une meilleure efficacité sur les grandes opérations matricielles denses. Les charges de travail dont les formes de tenseur sont statiques s’exécutent le mieux, car tout changement de forme oblige XLA à recompiler. Les TPU conviennent mal aux opérations personnalisées, aux branchements complexes et aux calculs en haute précision.
• NPU : Terme générique désignant les accélérateurs de réseaux neuronaux présents dans les téléphones, les ordinateurs portables et les puces de périphérie ; cette catégorie comprend l’Edge TPU.

L’Edge TPU#

L’Edge TPU est une petite TPU réservée à l’inférence, présente dans l’accélérateur USB Coral de Google, les modules M.2 et PCIe, ainsi que le Dev Board. Elle effectue 4 billions d’opérations par seconde pour une puissance de 2 watts et n’exécute que des calculs entiers sur 8 bits. Les réseaux doivent donc être entièrement convertis en INT8 par quantification du modèle, puis compilés avec le compilateur Edge TPU de Google, qui place les opérations compatibles sur la puce et laisse les autres au CPU hôte. Ultralytics YOLO automatise les deux étapes avec format="edgetpu", qui active la quantification INT8 et lance le compilateur sur une machine Linux x86-64. Le modèle obtenu s’exécute sur des cartes comme un Raspberry Pi équipé d’un accélérateur USB Coral :

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}")

La plateforme Coral de Google s’articule désormais autour de Coral NPU, une conception d’accélérateur RISC-V libre annoncée en octobre 2025 pour les appareils portables, qui vise 512 milliards d’opérations par seconde pour une consommation de quelques milliwatts.