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, la v6e présente de nombreuses similitudes avec la v5e. Ce système est optimisé pour être le produit le plus intéressant pour l'entraînement, le réglage fin et le service des transformateurs, des réseaux de neurones convolutifs (CNN) et des modèles de conversion de texte en image.

Architecture du système

Chaque puce v6e contient un TensorCore. Chaque TensorCore comporte deux unités de multiplication de matrice (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 comportant 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ôte à 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 un seul processeur, principalement destinée aux tests
• v6e-8 : VM hôte complet optimisée pour un cas d'utilisation d'inférence avec les huit puces associées à une seule VM.

Le tableau suivant compare les 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 l'indicateur --machine-type pour spécifier un type de machine compatible avec le TPU que vous souhaitez utiliser. Pour en savoir plus, consultez la section Planifier les TPU dans GKE de 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 de 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