Types et topologies de TPU

Configurations TPU

Lorsque vous créez un TPU, vous devez spécifier une configuration TPU. Vous pouvez spécifier une configuration TPU en termes de TensorCores ou de puces TPU. La spécification d'une configuration TPU en termes de puces n'est nécessaire que dans les scénarios avancés où vous entraînez un modèle volumineux et ajustez ses performances.

Configuration basée sur TensorCore

Vous pouvez utiliser une configuration basée sur TensorCore avec les TPU v2 à v4. Vous décrivez les configurations basées sur TensorCore avec une version de TPU et le nombre de TensorCore que vous souhaitez utiliser. Par exemple, v4-128 spécifie une configuration avec des TPU v4 et 128 TensorCore. Spécifiez une configuration basée sur TensorCore avec l'option --accelerator-type de la commande gcloud compute TPU TPU vm:

  $ gcloud compute tpus tpu-vm create tpu-name
    --zone=zone
    --accelerator-type=v3-128
    --version=tpu-vm-tf-2.12.0

Configuration d'un TPU basé sur une puce

Vous pouvez utiliser une configuration basée sur une puce TPU avec les TPU v4 et versions ultérieures. Vous devez décrire les configurations basées sur des puces avec une version de TPU et une topologie (disposition physique de puces TPU). Spécifiez une version de TPU avec une chaîne de type v4. Vous spécifiez une topologie avec un 3-tuple décrivant la façon dont les puces sont organisées en trois dimensions. Spécifiez une configuration basée sur un chip avec les indicateurs --version et --topology de la commande gcloud compute TPU TPU vm. Par exemple, la commande suivante spécifie TPU v4 comme version de TPU et 2x2x4 (16 chips, chacun avec 2 TensorCore) comme topologie: 

  $ gcloud compute tpus tpu-vm create tpu-name
    --zone=zone
    --type=v4
    --topology=2x2x4
    --version=tpu-vm-tf-2.12.0

Pour en savoir plus sur les différentes versions et topologies de Cloud TPU, consultez Architecture système.

La représentation d'une topologie v4 est AxBxC, où A<=B<=C et A, B, C sont tous <= 4 ou tous les multiples entiers de 4. Les valeurs A, B et C correspondent au nombre de puces dans chacune des trois dimensions. Les topologies dans lesquelles 2A=B=C ou 2A=2B=C comportent également des variantes de topologie optimisées pour la communication globale, par exemple 4 × 4 × 8, 8 × 8 × 16 et 12 × 12 × 24.

Pour en savoir plus sur les topologies v4, consultez Configurations de TPU v4.

Variantes de topologie

La topologie par défaut associée à un nombre de puces donné est la plus semblable à un cube (voir Formes de topologie). Cette forme est probablement la meilleure option pour l'entraînement de ML avec parallélisme des données. D'autres topologies peuvent être utiles pour les charges de travail comportant plusieurs types de parallélisme (par exemple, le parallélisme des modèles et des données, ou le partitionnement spatial d'une simulation). Ces charges de travail sont plus performantes si la topologie est mise en correspondance avec le parallélisme utilisé. Par exemple, le placement du parallélisme des modèles quadridirectionnels sur la dimension X et du parallélisme des données à 256 voies sur les dimensions Y et Z correspond à une topologie de 4x16x16.

Les modèles comportant plusieurs dimensions de parallélisme sont plus performants avec leurs dimensions de parallélisme mappées à celles de la topologie TPU. Il s'agit généralement de modèles big data (LLM) parallèles data+model. Par exemple, pour une tranche de TPU v4 avec une topologie de 8x16x16, ses dimensions sont 8, 16 et 16. Il est plus performant d'utiliser le parallélisme des modèles à 8 ou 16 voies (mappé sur l'une des dimensions de la topologie physique du TPU). Un parallélisme de modèle à quatre voies serait moins optimal avec cette topologie, car il n'est aligné avec aucune des dimensions de la topologie TPU, mais serait optimal avec une topologie de 4 x 16 x 32 sur le même nombre de puces.

Petites topologies v4

Cloud TPU est compatible avec les tranches TPU v4 suivantes, de taille inférieure à 64 puces, soit un cube 4x4x4. Vous pouvez créer ces topologies v4 de petite taille à l'aide de leur nom basé sur TensorCore (par exemple, v4-32) ou de leur topologie (par exemple, 2x2x4):

Nom (basé sur le nombre de TensorCore) Nombre de chips Topology
v4-8 4 2x2x1
V4-16 8 2x2x2
v4-32 16 2x2x4
v4-64 32 2 x 4

Grandes topologies v4

Les tranches de TPU v4 sont disponibles par incréments de 64 chips, avec des formes de multiples de 4 sur les trois dimensions. Les dimensions doivent également être dans l'ordre croissant. Le tableau suivant présente plusieurs exemples. Certaines de ces topologies sont des topologies "personnalisées" qui ne peuvent être lancées qu'à l'aide de l'API de topologie, car elles ont le même nombre de puces qu'une topologie nommée plus communément utilisée.

Nom (basé sur le nombre de TensorCore) Nombre de chips Topology
v4-128 64 4 x 4
v4-256 128 4 x 4
V4-512 256 4x8x8
Voir API Topology 256 4x4x16
V4-1024 512 8x8x8
V4-1536 768 8x8x12
V4-2048 1 024 8x8x16
Voir API Topology 1 024 4x16x16
V4-4096 2 048 8x16x16
… … …

API Topology

Pour créer des tranches de pod Cloud TPU avec une topologie personnalisée, vous pouvez utiliser l'API TPU gcloud comme suit:

     $ gcloud compute tpus tpu-vm create tpu-name \
        --zone=us-central2-b \
        --subnetwork=tpusubnet \
        --type=v4 \
        --topology=4x4x16 \
        --version=runtime-version

Formes de topologie

L'image suivante montre trois formes de topologie possibles et les paramètres d'indicateur topology associés.

Image

Compatibilité avec les TPU

Vous pouvez remplacer le type de TPU par un autre type ayant le même nombre de TensorCore ou de chips (par exemple, v3-128 et v4-128). Vous pouvez également exécuter votre script d'entraînement sans modifier le code. Toutefois, si vous passez à un type de TPU comportant un nombre plus élevé ou plus faible de puces ou de TensorCores, vous devrez effectuer un réglage et une optimisation importants. Pour plus d'informations, consultez la page Entraînement sur pods TPU.

Étapes suivantes