TPU v6e

Ce document décrit l'architecture et les configurations compatibles de Cloud TPU v6e (Trillium).

Trillium est la dernière génération d'accélérateurs d'IA Cloud TPU. Sur toutes les surfaces techniques, telles que l'API et les journaux, et tout au long de ce document, Trillium sera désigné par le terme "v6e".

Avec une empreinte de 256 puces par pod, v6e présente de nombreuses similitudes avec v5e. Ce système est optimisé pour être le produit le plus performant pour l'entraînement, l'affinage et la mise en service des transformateurs, des conversions de texte en image et des réseaux de neurones convolutifs (CNN).

Architecture du système

Chaque puce v6e contient un TensorCore. Chaque TensorCore comporte deux unités de multiplication matricielle (MXU), une unité vectorielle et une unité scalaire. Le tableau suivant présente les principales spécifications et leurs valeurs pour les TPU v6e par rapport aux TPU v5e.

Configurations compatibles

Le tableau suivant présente les formes de tranche 2D compatibles avec v6e :

Les tranches avec huit puces (v6e-8) associées à une seule VM sont optimisées pour l'inférence, ce qui permet d'utiliser les huit puces dans une même charge de travail de diffusion. Vous pouvez effectuer une inférence multi-hôtes à l'aide de Pathways on Cloud. Pour en savoir plus, consultez Effectuer une inférence multihôte à l'aide de Pathways.

Pour en savoir plus sur le nombre de VM pour chaque topologie, consultez Types de VM.

Types de VM

Chaque VM TPU v6e peut contenir une, quatre ou huit puces. Les tranches de quatre puces ou moins ont le même nœud NUMA (Non-Uniform Memory Access). Pour en savoir plus sur les nœuds NUMA, consultez Non-uniform memory access sur Wikipédia.

Les tranches v6e sont créées à l'aide de VM à demi-hôte, chacune avec quatre puces TPU. Il existe deux exceptions à cette règle :

• v6e-1 : VM avec une seule puce, principalement destinée aux tests
• v6e-8 : VM à hôte complet optimisée pour un cas d'utilisation d'inférence avec l'intégralité des huit puces associées à une seule VM.

Le tableau suivant présente une comparaison des types de VM TPU v6e :

Spécifier la configuration v6e

Lorsque vous allouez une tranche TPU v6e à l'aide de l'API TPU, vous spécifiez sa taille et sa forme à l'aide du paramètre AcceleratorType.

Si vous utilisez GKE, utilisez le flag --machine-type pour spécifier un type de machine compatible avec le TPU que vous souhaitez utiliser. Pour en savoir plus, consultez Planifier des TPU dans GKE dans la documentation GKE.

Utiliser AcceleratorType

Lorsque vous allouez des ressources TPU, vous utilisez AcceleratorType pour spécifier le nombre de TensorCores dans une tranche. La valeur que vous spécifiez pour AcceleratorType est une chaîne au format v$VERSION-$TENSORCORE_COUNT. Par exemple, v6e-8 spécifie une tranche TPU v6e avec huit TensorCores.

L'exemple suivant montre comment créer une tranche TPU v6e avec 32 TensorCores à l'aide de AcceleratorType :

    $ gcloud compute tpus tpu-vm create tpu-name \
        --zone=zone \
        --accelerator-type=v6e-32 \
        --version=v2-alpha-tpuv6e

Étapes suivantes

• Exécuter l'entraînement et l'inférence à l'aide de TPU v6e