Présentation des instances multitranches Cloud TPU

Cloud TPU multi-tranche est une technologie de scaling des performances full stack qui permet à un job d'entraînement d'utiliser plusieurs tranches de TPU dans un seul pod ou sur des tranches de plusieurs pods avec un simple parallélisme des données. Avec les puces TPU v4, les tâches d'entraînement peuvent utiliser plus de 4 096 puces en une seule exécution. Pour les tâches d'entraînement nécessitant moins de 4 096 puces, une seule tranche peut offrir les meilleures performances. Toutefois, plusieurs tranches plus petites sont plus facilement disponibles, ce qui accélère le démarrage lorsque les tranches multitranches sont utilisées avec des tranches plus petites.

Plusieurs segments effectuent un scaling linéaire des performances

Lorsqu'elles sont déployées dans des configurations à plusieurs tranches, les puces TPU de chaque tranche communiquent via ICI (inter-chip-interconnect). Les puces TPU de différentes tranches communiquent en transférant des données aux processeurs (hôtes), qui à leur tour transmettent les données sur le réseau du centre de données (DCN).

Flux de données à plusieurs tranches

Les développeurs n'ont pas besoin d'écrire de code pour implémenter la communication entre segments DCN. Le compilateur XLA génère ce code pour vous et chevauche la communication avec le calcul pour des performances optimales.

Concepts

Type d'accélérateur
Forme de chaque tranche TPU comprenant une multitranche. Chaque segment d'une requête à plusieurs tranches est du même type d'accélérateur. Un type d'accélérateur est composé d'un type de TPU (v4 ou v5e) suivi du nombre de TensorCores. Par exemple, v4-128 spécifie un TPU v4 avec 128 TensorCores.
Réparation automatique
Lorsqu'une tranche rencontre un événement de maintenance, une préemption ou une défaillance matérielle, Cloud TPU crée une nouvelle tranche. Dans les rares cas où les ressources sont insuffisantes pour créer une tranche, la création ne se termine que lorsque le matériel devient disponible. Une fois la nouvelle tranche créée, toutes les autres tranches de l'environnement multisegments redémarrent pour que l'entraînement puisse se poursuivre.Avec un script de démarrage correctement configuré, le script d'entraînement peut se relancer automatiquement sans intervention de l'utilisateur, chargement et reprise à partir du dernier point de contrôle.
Ensemble de données
Données utilisées par un modèle à des fins d'entraînement ou d'inférence.
Mise en réseau de centres de données (DCN)
Réseau à faible latence et à faible débit (par rapport à ICI) qui connecte les tranches TPU dans une configuration à plusieurs tranches.
Planification de gangs
Lorsque toutes les tranches TPU sont provisionnées ensemble, en même temps, cela garantit que toutes les tranches ou aucune d'entre elles ne sont correctement provisionnées.
Organisateur
Un hôte est un ordinateur physique qui exécute des VM. Un hôte peut exécuter au maximum quatre VM à la fois. Chaque VM dispose d'un TPU dédié.
Inférence
Charger un modèle de machine learning pré-entraîné sur un hôte et réaliser des prédictions sur les données
Interchip Interconnect (ICI)
Liens internes haut débit et à faible latence qui connectent les TPU au sein d'un pod TPU.
Tranches multiples
Deux ou plusieurs tranches de puce TPU qui peuvent communiquer sur le réseau multichaîne.
Nœud
Dans le contexte à plusieurs tranches, le nœud fait référence à une seule tranche TPU. Chaque tranche TPU dans une tranche multislice se voit attribuer un ID de nœud.
Pod
Ensemble de puces TPU connectées par des interfaces réseau ICI dédiées. Un pod vous permet de répartir la charge de traitement sur plusieurs TPU.
Ressource en file d'attente (QR)
Représentation des ressources TPU, utilisée pour mettre en file d'attente et gérer une requête pour un environnement TPU à tranche unique ou à plusieurs tranches.
Script de démarrage
Script de démarrage Compute Engine standard qui s'exécute à chaque démarrage ou redémarrage d'une VM. Pour les multisegments, il est spécifié dans la requête de création de code QR. Pour en savoir plus sur les scripts de démarrage Cloud TPU, consultez la page Gérer les ressources TPU.
Tranche TPU
Sous-section logique d'un pod TPU composée de puces TPU. Tous les chips d'un segment communiquent entre eux via le réseau ICI.
VM TPU
Machine virtuelle exécutant Linux qui a accès aux TPU sous-jacents. Pour les TPU v4, chaque VM TPU dispose d'un accès direct à quatre puces. Parfois, nous appelons une VM TPU un nœud de calcul.
Tensor
Structure de données utilisée pour représenter des données multidimensionnelles dans un modèle de machine learning.
Tensor Processing Unit (TPU)
Puce d'accélération de ML développée en interne par Google. Ils sont conçus pour offrir des capacités de calcul rapides et économes en énergie pour les tâches de machine learning clés telles que la multiplication de matrices.
Types de capacité Cloud TPU

Les TPU peuvent être créés à partir de différents types de capacité (consultez la section "Options d'utilisation" dans la section Fonctionnement des tarifs des TPU) :

  • Réservation: cible un quota réservé. Pour utiliser un quota réservé, vous devez avoir conclu un accord de réservation avec Google. Utilisez l'option --reserved lors de la création de vos ressources.
  • Spot: cible un quota préemptif à l'aide de VM Spot. Vos ressources peuvent être préemptées pour faire de la place aux requêtes d'une tâche de priorité plus élevée. Utilisez l'option --spot lors de la création de vos ressources.
  • À la demande: cible le quota à la demande, qui ne nécessite pas de réservation et ne sera pas préempté. La requête TPU est mise en file d'attente dans une file d'attente de quota à la demande proposée par Cloud TPU. La disponibilité des ressources n'est pas garantie. Cette option est sélectionnée par défaut, aucun indicateur n'est nécessaire.

Démarrer

Si vous n'avez jamais utilisé de TPU auparavant, commencez par installer la Google Cloud CLI et configurer votre environnement Cloud TPU. Pour utiliser des tranches multiples, vos ressources TPU doivent être gérées en tant que ressources en file d'attente.

Si vous utilisez déjà TPU v4 et que vous disposez d'une réservation, vous devrez peut-être la migrer vers un nouveau système de réservation. Pour en savoir plus, contactez votre responsable de compte Google Cloud.

Exemple d'introduction

Ce tutoriel utilise du code du dépôt GitHub MaxText. MaxText est un LLM de base hautement performant, arbitrairement évolutif, Open Source et éprouvé, écrit en Python et en Jax. MaxText a été conçu pour permettre un entraînement efficace sur Cloud TPU.

Le code de shardings.py est conçu pour vous aider à tester différentes options de parallélisme. Exemples : parallélisme des données, parallélisme des données entièrement segmentées (FSDP) et parallélisme des Tensors. Le code passe d'une tranche unique à un environnement à plusieurs tranches.

Parallélisme ICI

ICI fait référence à l'interconnexion haut débit qui connecte les TPU en une tranche unique. La segmentation ICI correspond à la segmentation au sein d'une tranche. shardings.py fournit trois paramètres de parallélisme ICI:

  • ici_data_parallelism
  • ici_fsdp_parallelism
  • ici_tensor_parallelism

Les valeurs que vous spécifiez pour ces paramètres déterminent le nombre de segments pour chaque méthode de parallélisme.

Ces entrées doivent être limitées pour que ici_data_parallelism * ici_fsdp_parallelism * ici_tensor_parallelism soit égal au nombre de chips dans la tranche.

Le tableau suivant présente des exemples d'entrées utilisateur pour le parallélisme ICI pour les quatre puces disponibles dans la version v4-8:

ici_data_parallelism ici_fsdp_parallelism ici_tensor_parallelism
FSDP quadridirectionnel 1 4 1
Parallélisme du Tensor quadridirectionnel 1 1 4
FSDP bidirectionnel + parallélisme Tensor bidirectionnel 1 2 2

Notez que ici_data_parallelism doit rester défini sur 1 dans la plupart des cas, car le réseau ICI est suffisamment rapide pour presque toujours privilégier le FSDP au parallélisme des données.

Dans cet exemple, nous partons du principe que vous savez exécuter du code sur une seule tranche TPU, comme dans la section Exécuter un calcul sur une VM Cloud TPU à l'aide de JAX. Cet exemple montre comment exécuter shardings.py sur une seule tranche.

  1. Configurez l'environnement:

    $ gcloud auth login
$ gcloud config set project your-project-id
$ gcloud config set compute/zone your-zone

  2. Créez des clés SSH pour gcloud. Nous vous recommandons de ne pas saisir de mot de passe (appuyez deux fois sur Entrée après avoir exécuté la commande suivante). Si vous êtes invité à indiquer que le fichier google_compute_engine existe déjà, remplacez la version existante.

    $ ssh-keygen -f ~/.ssh/google_compute_engine

  3. Provisionnez vos TPU à l'aide de la commande suivante:

    $ gcloud alpha compute tpus queued-resources \
create your-qr-id \
--accelerator-type your-accelerator-type \
--runtime-version tpu-ubuntu2204-base \
--node-id qr-id \
[--reserved |--spot]

    Description des options de commande

    your-qr-id
    Chaîne définie par l'utilisateur qui identifie la requête de code QR.
    accelerator-type
    Le type d'accélérateur spécifie la version et la taille de la ressource Cloud TPU que vous souhaitez créer. Pour en savoir plus sur les types d'accélérateurs compatibles avec chaque version de TPU, consultez la section Versions de TPU.
    runtime-version
    [Version logicielle de Cloud TPU](/tpu/docs/supported-tpu-configurations#tpu_software_versions).
    node-id
    ID des ressources TPU qui seront créées en réponse à la requête QR.
    reserved
    Utilisez un quota réservé lors de la création des tranches.
    best-effort
    Utilisez le quota de la façon la plus optimale possible lors de la création des tranches [par défaut].

    La Google Cloud CLI ne prend pas en charge toutes les options de création de code QR, telles que les balises. Pour en savoir plus, consultez Créer des codes QR.

  4. Attendez que le code QR soit à l'état ACTIVE, ce qui signifie que les nœuds de calcul présentent l'état READY. Une fois que le provisionnement du code QR a commencé, il peut prendre une à cinq minutes en fonction de la taille du code QR. Vous pouvez vérifier l'état d'une requête QR à l'aide de la commande suivante:

    $ gcloud compute tpus queued-resources \
  list --filter=your-qr-id

  5. Une tranche v4-8 comporte une seule VM TPU. Connectez-vous à la VM TPU à l'aide de SSH:

    $ gcloud compute tpus tpu-vm ssh your-qr-id

  6. Clonez MaxText (qui inclut shardings.py) sur votre VM TPU.

  7. Dans le répertoire du dépôt MaxText, exécutez le script de configuration pour installer JAX et d'autres dépendances sur votre tranche TPU. L'exécution du script de configuration prend quelques minutes.

    $ bash setup.sh

  8. Exécutez la commande suivante pour exécuter shardings.py sur votre tranche TPU.

    $ python3 pedagogical_examples/shardings.py \
  --ici_fsdp_parallelism 4 \
  --batch_size 131072 \
  --embedding_dimension 2048

    Vous pouvez voir les résultats dans les journaux. Vos TPU doivent atteindre environ 260 TFLOP par seconde, soit une utilisation impressionnante de plus de 90%de FLOP. Dans ce cas, nous avons sélectionné approximativement le lot maximal qui correspond à la mémoire à haut débit (HBM) du TPU.

  9. N'hésitez pas à explorer d'autres stratégies de segmentation via ICI. Vous pouvez par exemple essayer la combinaison suivante:

    $ python3 pedagogical_examples/shardings.py \
  --ici_tensor_parallelism 4 \
  --batch_size 131072 \
  --embedding_dimension 2048

  10. Supprimez le code QR et la tranche TPU lorsque vous avez terminé. Vous devez exécuter ces étapes de nettoyage à partir de l'environnement dans lequel vous avez configuré la tranche (exécutez d'abord exit pour quitter la session SSH). La suppression prend deux à cinq minutes et peut être exécutée en arrière-plan avec l'option facultative --async.

    $ gcloud compute tpus queued-resources
  delete your-qr-id --force (--async)

Segmentation de plusieurs tranches à l'aide du parallélisme DCN

Le script shardings.py utilise trois paramètres qui spécifient le parallélisme des réseaux multichaînes, correspondant au nombre de segments de chaque type de parallélisme des données:

  • dcn_data_parallelism
  • dcn_fsdp_parallelism
  • dcn_tensor_parallelism

Les valeurs de ces paramètres doivent être limitées de sorte que dcn_data_parallelism * dcn_fsdp_parallelism * dcn_tensor_parallelism soit égal au nombre de tranches.

Par exemple pour deux segments, utilisez --dcn_data_parallelism = 2.

dcn_data_parallelism dcn_fsdp_parallelism dcn_tensor_parallelism Nombre de segments
Parallélisme des données bidirectionnelle 2 1 1 2

dcn_tensor_parallelism doit toujours être défini sur 1, car le DCN n'est pas adapté à ce type de segmentation. Pour les charges de travail LLM classiques sur des puces v4, dcn_fsdp_parallelism doit également être défini sur 1. Par conséquent, dcn_data_parallelism doit être défini sur le nombre de tranches, mais cela dépend de l'application.

À mesure que vous augmentez le nombre de tranches (en supposant que la taille de tranche et le lot par tranche restent constants), vous augmentez la quantité de parallélisme des données.

Exécuter shardings.py dans un environnement à plusieurs tranches

Vous pouvez exécuter shardings.py dans un environnement à plusieurs tranches à l'aide de multihost_runner.py ou en exécutant shardings.py sur chaque VM TPU. Ici, nous utilisons multihost_runner.py. Les étapes suivantes sont très semblables à celles de la section Premiers pas: tests rapides sur plusieurs segments du dépôt MaxText, sauf qu'ici, nous exécutons shardings.py au lieu du LLM plus complexe dans train.py.

L'outil multihost_runner.py est optimisé pour des tests rapides, en réutilisant à plusieurs reprises les mêmes TPU. Étant donné que le script multihost_runner.py dépend de connexions SSH de longue durée, nous ne le recommandons pas pour les tâches de longue durée. Si vous souhaitez exécuter une tâche plus longue (par exemple, avec des heures ou des jours), nous vous recommandons d'utiliser multihost_job.py.

Dans ce tutoriel, nous utilisons le terme runner pour désigner la machine sur laquelle vous exécutez le script multihost_runner.py. Nous utilisons le terme nœuds de calcul pour désigner les VM TPU qui constituent vos tranches. Vous pouvez exécuter multihost_runner.py sur une machine locale ou sur n'importe quelle VM Compute Engine du même projet que vos tranches. Il n'est pas possible d'exécuter multihost_runner.py sur un nœud de calcul.

multihost_runner.py se connecte automatiquement aux nœuds de calcul TPU à l'aide de SSH.

Dans cet exemple, nous exécutons shardings.py sur deux tranches v4-16, soit un total de quatre VM et 16 puces TPU. Vous pouvez modifier l'exemple pour l'exécuter sur plus de TPU.

Configurer votre environnement

  1. Clonez MaxText sur votre ordinateur d'exécution.

  2. Accédez au répertoire du dépôt.

  3. Créez des clés SSH pour gcloud. Nous vous recommandons de ne pas saisir de mot de passe (appuyez deux fois sur Entrée après avoir exécuté la commande suivante). Si vous êtes invité à indiquer que le fichier google_compute_engine existe déjà, choisissez de ne pas conserver la version existante.

      $ ssh-keygen -f ~/.ssh/google_compute_engine

  4. Ajoutez une variable d'environnement pour définir le nombre de tranches TPU sur 2.

      $ export SLICE_COUNT=2

  5. Créez un environnement à plusieurs tranches à l'aide de queued-resources create.

    La commande suivante montre comment créer un TPU multislice v4. Pour utiliser la version v5e, spécifiez un accelerator-type v5e (par exemple, v5litepod-16) et un élément runtime-version v5e (v2-alpha-tpuv5-lite).

      $ gcloud alpha compute tpus queued-resources 
    
  create your-qr-id 
    
  --accelerator-type=your-accelerator-type 
    
  --runtime-version=tpu-vm-runtime-version 
    
  --node-count=node-count 
    
  --node-prefix=your-qr-id 
    
  [--reserved|--spot]

    Description des options de commande

    your-qr-id
    Chaîne définie par l'utilisateur qui identifie la requête de code QR.
    accelerator-type
    Le type d'accélérateur spécifie la version et la taille de la ressource Cloud TPU que vous souhaitez créer. Pour en savoir plus sur les types d'accélérateurs compatibles avec chaque version de TPU, consultez la section Versions de TPU.
    runtime-version
    Version logicielle de Cloud TPU.
    node-count
    Nombre de tranches à créer.
    node-prefix
    Préfixe utilisé pour générer les noms de chaque tranche. Un numéro est ajouté au préfixe pour chaque tranche. Par exemple, si vous définissez node-prefix sur mySlice, les tranches sont nommées mySlice-0, mySlice-1, etc.
    reserved
    Utilisez un quota réservé lors de la création des tranches.
    best-effort
    Utilisez le quota de la façon la plus optimale possible lors de la création des tranches [par défaut].

  6. Lorsque le provisionnement du code QR commence, il peut s'écouler jusqu'à cinq minutes selon la taille du code QR. Attendez que la ressource en file d'attente (QR) passe à l'état ACTIVE. Vous pouvez vérifier l'état d'une requête QR à l'aide de la commande suivante:

    $ gcloud compute tpus queued-resources list \
--filter=your-qr-id

    Vous devriez obtenir un résultat semblable à celui-ci:

    NAME        ZONE           NODE_COUNT  ACCELERATOR_TYPE  STATE
...
que-res-id  us-central2-b  4           v4-16             ACTIVE
...

    Contactez votre responsable de compte Google Cloud si l'état du code QR est à l'état WAITING_FOR_RESOURCES ou PROVISIONING depuis plus de 15 minutes.

  7. Installez les dépendances.

    $ python3 multihost_runner.py \
  --TPU_PREFIX=your-qr-id \
  --COMMAND="bash setup.sh"

  8. Exécutez shardings.py sur chaque nœud de calcul à l'aide de multihost_runner.py.

    $ python3 multihost_runner.py \
  --TPU_PREFIX=your-qr-id \
  --COMMAND="python3 pedagogical_examples/shardings.py \
  --dcn_data_parallelism $SLICE_COUNT \
  --ici_fsdp_parallelism 8 \
  --batch_size 131072 \
  --embedding_dimension 2048"

    Vous constaterez environ 230 TFLOP par seconde de performances dans les fichiers journaux.

  9. Nettoyez les TPU et le code QR lorsque vous avez terminé. La suppression prend deux à cinq minutes et peut être exécutée en arrière-plan avec l'option facultative --async.

Procéder au scaling d'une charge de travail en multitranche

Avant d'exécuter votre modèle dans un environnement à plusieurs tranches, apportez les modifications de code suivantes:

Il doit s'agir des seules modifications de code nécessaires lors du passage à une tranche multislice. Pour obtenir des performances élevées, le DCN doit être mappé sur des données parallèles, entièrement segmentées ou sur des axes parallèles de pipeline. Les considérations de performances et les stratégies de segmentation sont abordées plus en détail dans la section Segmentation avec plusieurs tranches pour des performances maximales.

Pour vérifier que votre code peut accéder à tous les appareils, vous pouvez confirmer que len(jax.devices()) est égal au nombre de puces dans votre environnement multitranche. Par exemple, si vous utilisez quatre segments de v4-16, vous obtenez huit chips par tranche x 4 secteurs. len(jax.devices()) doit donc renvoyer 32.

Choisir la taille des tranches pour les environnements à plusieurs tranches

Pour accélérer de façon linéaire, ajoutez des tranches de la même taille que votre tranche existante. Par exemple, si vous utilisez une tranche v4-512, la multicouche atteindra environ deux fois plus de performances en ajoutant une deuxième tranche v4-512 et en doublant la taille de lot globale. Pour en savoir plus, consultez la page Segmentation avec plusieurs tranches pour des performances maximales.

Exécuter votre job sur plusieurs tranches

Il existe trois approches différentes pour exécuter une charge de travail personnalisée dans un environnement à plusieurs tranches:

  1. À l'aide du script du lanceur de tests, multihost_runner.py
  2. À l'aide du script d'exécuteur de production, multihost_job.py
  3. Utiliser une approche manuelle

Script de l'exécuteur d'expérimentation

Le script multihost_runner.py distribue le code dans un environnement multislice existant, exécute votre commande sur chaque hôte, copie vos journaux et suit l'état d'erreur de chaque commande. Le script multihost_runner.py est documenté dans le fichier README MaxText.

Étant donné que multihost_runner.py gère des connexions SSH persistantes, il n'est adapté qu'aux tests de taille modeste et de courte durée. Vous pouvez adapter les étapes du tutoriel multihost_runner.py à votre charge de travail et à votre configuration matérielle.

Script d'exécuteur de production

Pour les tâches de production nécessitant une résilience aux défaillances matérielles et autres préemptions, il est préférable de les intégrer directement à l'API Create Queued Resource. À titre d'exemple fonctionnel, nous fournissons multihost_job.py, qui déclenche l'appel d'API Created Queued Resource avec le script de démarrage approprié pour exécuter votre entraînement et reprendre l'entraînement en cas de préemption. Le script multihost_job.py est documenté dans le fichier README MaxText.

Étant donné que multihost_job.py doit provisionner des ressources pour chaque exécution, il ne fournit pas de cycle d'itération aussi rapide que multihost_runner.py.

Approche manuelle

Nous vous recommandons d'utiliser ou d'adapter multihost_runner.py ou multihost_job.py pour exécuter votre charge de travail personnalisée dans votre configuration multislice. Toutefois, si vous préférez provisionner et gérer directement votre environnement à l'aide de commandes QR, consultez la page Gérer un environnement à plusieurs tranches.

Gérer un environnement à plusieurs tranches

Pour provisionner et gérer manuellement les codes QR sans utiliser les outils fournis dans le dépôt MaxText, consultez les sections suivantes.

Créer des codes QR

Définissez les variables d'environnement suivantes avant de provisionner la capacité:

  $ export your-qr-id=your-queued-resource-id
  $ export PROJECT=your-project-name
  $ export ZONE=us-central2-b
  $ export NETWORK_NAME=your-network-name
  $ export SUBNETWORK_NAME=your-subnetwork-name
  $ export RUNTIME_VERSION=tpu-ubuntu2204-base
  $ export ACCELERATOR_TYPE=v4-16
  $ export SLICE_COUNT=4
  $ export STARTUP_SCRIPT="#!/bin/bash\n ..."
  $ gcloud config set project project-name
  $ gcloud config set compute/zone zone
Entrée Description
your-qr-id Identifiant du code QR attribué par l'utilisateur.
PROJET Nom du projet Google Cloud
ZONE us-central2-b
NETWORK_NAME Nom des réseaux VPC.
SUBNETWORK_NAME Nom du sous-réseau dans les réseaux VPC
RUNTIME_VERSION tpu-Ubuntu-2204-base
ACCELERATOR_TYPE v4-16
EXAMPLE_TAG_1, EXAMPLE_TAG_2... Tags permettant d'identifier les sources ou les cibles valides pour les pare-feu du réseau
SLICE_COUNT Nombre de tranches. Limité à 256 tranches au maximum.
STARTUP_SCRIPT S'il est ajouté à la requête de création, un script de démarrage peut s'exécuter chaque fois qu'une tranche TPU est provisionnée ou redémarrée, et si la tranche TPU est réparée ou réinitialisée.

Créer une demande de code QR à l'aide de gcloud

$ gcloud alpha compute tpus queued-resources \
  create ${your-qr-id} \
  --project your-project-id \
  --zone your-zone \
  --node-count ${SLICE_COUNT} \
  --accelerator-type ${ACCELERATOR_TYPE} \
  --runtime-version ${RUNTIME_VERSION} \
  --network ${NETWORK_NAME} \
  --subnetwork ${SUBNETWORK_NAME} \
  --tags ${EXAMPLE_TAG_1},${EXAMPLE_TAG_2} \ --metadata=startup-script='${STARTUP_SCRIPT}'
  [--reserved|--spot]

Description des options de commande

your-qr-id
Chaîne définie par l'utilisateur qui identifie la requête de code QR.
project
Chaîne définie par l'utilisateur qui identifie la requête de code QR.
zone
Zone Google Cloud dans laquelle créer le code QR.
node-count
Nombre de tranches à créer.
accelerator-type
Le type d'accélérateur spécifie la version et la taille de la ressource Cloud TPU que vous souhaitez créer. Pour en savoir plus sur les types d'accélérateurs compatibles avec chaque version de TPU, consultez la section Versions de TPU.
runtime-version
Version logicielle de Cloud TPU.
network
Nom d'un réseau VPC auquel associer la ressource TPU.
subnetwork
Nom d'un sous-réseau VPC auquel associer la ressource TPU.
reserved
Utilisez un quota réservé lors de la création des tranches.
spot
Utilisez le quota de VM Spot lors de la création des tranches.

Assurez-vous de disposer du quota correspondant avant de sélectionner --reserved, --spot ou le quota à la demande par défaut. Pour en savoir plus sur les types de quotas, consultez les Règles relatives aux quotas.

Créer une demande de code QR à l'aide de curl

Créez un fichier nommé queued-resource-req.json et copiez-y le code JSON suivant.

{
  "guaranteed": { "reserved": true },
  "tpu": {
    "node_spec": [
    {
      "parent": "projects/your-project-number/locations/your-zone",
        "node": {
          "accelerator_type": "accelerator-type",
          "runtime_version": "tpu-vm-runtime-version",
          "network_config": {
            "network": "your-network-name",
            "subnetwork": "your-subnetwork-name",
            "enable_external_ips": true
          },
          "tags" : ["example-tag-1"]
          "metadata": {
            "startup-script": "your-startup-script"
          }
      },
      "multi_node_params": {
        "node_count": slice-count,
        "node_id_prefix": "your-queued-resource-id"
      }
    }
    ]
  }
}
  • your-project-number : numéro de votre projet Google Cloud
  • your-zone : zone dans laquelle vous souhaitez créer votre code QR
  • accelerator-type : version et taille d'une tranche unique
  • tpu-vm-runtime-version : versions d'exécution des VM TPU
  • your-network-name : facultatif, réseau auquel le code QR sera associé
  • your-subnetwork-name : facultatif, sous-réseau auquel le code QR sera associé
  • example-tag-1 : chaîne de balise arbitraire facultative
  • your-startup-script : script de démarrage qui sera exécuté lorsque le code QR sera alloué
  • slice-count : nombre de tranches TPU dans votre environnement à plusieurs tranches
  • your-qr-id : ID fourni par l'utilisateur pour le code QR

Consultez la documentation de l'API REST en file d'attente pour connaître toutes les options disponibles.

Pour utiliser la capacité du Spot, remplacez:

"guaranteed": { "reserved": true } avec "spot": {}

Supprimez la ligne pour utiliser la capacité à la demande par défaut.

Envoyez la requête de création de code QR avec la charge utile JSON:

  $ curl -X POST -H "Authorization: Bearer $(gcloud auth print-access-token)" -H "Content-Type: application/json" -d @queuedresourcereq.json https://tpu.googleapis.com/v2alpha1/projects/your-project-id/locations/your-zone/queuedResources\?queued_resource_id\=your-qr-id
  • your-project-id : ID de votre projet Google Cloud
  • your-zone : zone dans laquelle vous souhaitez créer votre code QR
  • your-qr-id : ID fourni par l'utilisateur pour le code QR

La réponse doit se présenter sous la forme suivante:

{
  "name": "projects/<your-project-id>/locations/<your-zone>/operations/operation-<your-qr-guid>",
  "metadata": {
    "@type": "type.googleapis.com/google.cloud.common.OperationMetadata",
    "createTime": "2023-11-01T00:17:05.742546311Z",
    "target": "projects/<your-project-id>/locations/<your-zone>/queuedResources/<your-qa-id>",
    "verb": "create",
    "cancelRequested": false,
    "apiVersion": "v2alpha1"
  },
  "done": false
}

Indiquez la valeur GUID à la fin de la valeur de chaîne pour l'attribut name afin d'obtenir des informations sur la requête QR.

Récupérer l'état d'un code QR

Pour obtenir l'état de la requête QR, utilisez la commande suivante:

  $ curl -X GET -H "Authorization: Bearer $(gcloud auth print-access-token)" -H "Content-Type: application/json" https://tpu.googleapis.com/v2/projects/your-project-id/locations/your-zone/operations/operation-your-qr-guid
  • your-project-id : ID de votre projet Google Cloud
  • your-zone : zone dans laquelle créer le code QR
  • your-qr-guid : GUID suivant name dans le résultat de la requête de création de code QR.

La réponse de cette commande contient l'état de l'opération:

{
  "name": "projects/<your-project-id>/locations/<your-zone>/operations/operation-<your-qa-guid>,
  "metadata": {...},
  "done": true,
  "response": {
    "@type": "type.googleapis.com/google.cloud.tpu.v2.QueuedResource",
    ...
    "state": {
      "state": "WAITING_FOR_RESOURCES"
    }
  }
}

Si le code QR a bien été créé (("done = true")), l'état dans le champ response sera WAITING_FOR_RESOURCES ou FAILED. Si le code QR est à l'état WAITING_FOR_RESOURCES, il a été mis en file d'attente et commencera le provisionnement lorsque les ressources seront suffisantes. Si le code QR est à l'état FAILED, le motif de l'échec sera indiqué dans la sortie. Pour en savoir plus sur les autres états possibles, consultez le guide de l'utilisateur des ressources en file d'attente.

Une fois l'opération terminée, utilisez les codes QR de description pour surveiller les étapes du code QR.

Dans de rares cas, votre code QR peut afficher l'état FAILED alors que certaines tranches sont à l'état ACTIVE. Dans ce cas, supprimez les ressources qui ont été créées et réessayez dans quelques minutes. Vous pouvez également contacter l'équipe Cloud TPU pour résoudre le problème.

SSH et installation des dépendances

La section Exécuter le code JAX sur des tranches de pod TPU explique comment se connecter à vos VM TPU à l'aide de SSH dans une seule tranche. Pour vous connecter à toutes les VM TPU de votre environnement à plusieurs tranches via SSH et installer les dépendances, exécutez la commande gcloud suivante:

  $ gcloud compute tpus queued-resources ssh ${your-qr-id} \
    --zone your-zone \
    --node=all \
    --worker=all \
    --command="command-to-run"
    --batch-size=4

Cette commande gcloud envoie la commande spécifiée à tous les nœuds de calcul et les nœuds dans le fichier QR à l'aide de SSH. La commande est répartie par groupes de quatre et est envoyée simultanément. Le lot suivant de commandes est envoyé lorsque l'exécution du lot en cours est terminée. En cas d'échec de l'une des commandes, le traitement s'arrête et aucun autre lot n'est envoyé. Pour en savoir plus, consultez la documentation de référence de l'API de ressources en file d'attente. Si le nombre de tranches que vous utilisez dépasse la limite de threads de votre ordinateur local (également appelée limite de traitement par lot), vous rencontrerez un interblocage. Par exemple, supposons que la limite de traitement par lot sur votre machine locale soit de 64. Si vous essayez d'exécuter un script d'entraînement sur plus de 64 tranches (par exemple 100), la commande SSH décomposera les tranches en lots. Il exécute le script d'entraînement sur le premier lot de 64 tranches et attend la fin des scripts avant de l'exécuter sur le lot restant de 36 tranches. Toutefois, le premier lot de 64 tranches ne peut pas se terminer tant que les 36 tranches restantes n'ont pas commencé à exécuter le script, ce qui entraîne un interblocage.

Pour éviter ce scénario, vous pouvez exécuter le script d'entraînement en arrière-plan sur chaque VM en ajoutant une esperluette (&) à la commande de script que vous spécifiez avec l'option --command. Lorsque vous procédez ainsi, après avoir démarré le script d'entraînement sur le premier lot de tranches, le contrôle est immédiatement renvoyé à la commande SSH. La commande SSH peut alors lancer l'exécution du script d'entraînement sur le lot restant de 36 tranches. Vous devez diriger vos flux stdout et stderr de manière appropriée lorsque vous exécutez les commandes en arrière-plan. Pour augmenter le parallélisme dans le même code QR, vous pouvez sélectionner des tranches spécifiques à l'aide du paramètre --node.

Configuration du réseau

Assurez-vous que les tranches TPU peuvent communiquer entre elles en exécutant les étapes suivantes. Installez JAX sur chacune des tranches. Pour en savoir plus, consultez la section Exécuter du code JAX sur des tranches de pod TPU. Déclarez que len(jax.devices()) est égal au nombre de chips dans votre environnement à tranches multiples. Pour ce faire, sur chaque tranche, exécutez la commande suivante:

  $ python3 -c 'import jax; print(jax.devices())'

Si vous exécutez ce code sur quatre segments des versions 4 à 16, il y a huit puces par tranche et quatre tranches, un total de 32 puces (périphériques) doit être renvoyé par jax.devices().

Lister les codes QR

Vous pouvez afficher l'état de vos codes QR à l'aide de la commande queued-resources list:

$ gcloud compute tpus queued-resources list

NAME        ZONE           NODE_COUNT  ACCELERATOR_TYPE  STATE
...
que-res-id  us-central2-b  4           v4-16             ACTIVE
...

Décrire les codes QR

Pour afficher la configuration et l'état détaillés d'un code QR, utilisez l'API QR de description. Vous pouvez appeler cette API à l'aide de gcloud ou curl.

En utilisant gcloud :

$ gcloud compute tpus queued-resources describe ${your-qr-id}
...state:
 state: ACTIVE
...

En utilisant curl :

$ curl -X GET -H "Authorization: Bearer $(gcloud auth print-access-token)" -H "Content-Type: application/json" https://tpu.googleapis.com/v2/projects/your-project-id/locations/your-zone/queuedResources/${your-qr-id}
{
  "name": your-queued-res,
  "tpu": {
    "nodeSpec": [
      {
        ... // node 1
      },
      {
        ... // node 2
      },
      ...
    ]
  },
  ...
  "state": "ACTIVE"
}

state représente l'état d'un code QR. Pour en savoir plus sur les états possibles des codes QR, consultez la section Ressources en file d'attente.

Démarrer un job dans un environnement provisionné

Vous pouvez exécuter manuellement des charges de travail en vous connectant à tous les hôtes de chaque tranche via SSH et en exécutant la commande suivante sur tous les hôtes.

$ gcloud compute tpus tpu-vm ssh your-qr-id \
  --zone=your-zone \
  --worker=all \
  --node=all \
  --command="command-to-run"

Réinitialisation des codes QR

L'API ResetQueuedResource peut être utilisée pour réinitialiser toutes les VM dans un code QR ACTIVE. La réinitialisation forcée des VM efface la mémoire de la machine et rétablit l'état initial de la VM. Toutes les données stockées localement resteront intactes et le script de démarrage sera appelé après une réinitialisation. L'API ResetQueuedResource peut être utile lorsque vous souhaitez redémarrer tous les TPU. Par exemple, lorsque l'entraînement est bloqué et que la réinitialisation de toutes les VM est plus facile que le débogage.

Les réinitialisations de toutes les VM sont effectuées en parallèle et une opération ResetQueuedResource prend une à deux minutes. Pour appeler l'API, utilisez la commande suivante:

$ gcloud compute tpus queued-resources reset your-qr-id

Suppression des codes QR

Pour libérer des ressources à la fin de votre session d'entraînement, supprimez la ressource en file d'attente avec l'option --force. La suppression prend deux à cinq minutes et peut être exécutée en arrière-plan avec l'option facultative --async.

$ gcloud compute tpus queued-resources \
delete your-qr-id --force (--async)

Reprise automatique après échec

En cas d'interruption, le multislice offre une réparation sans intervention de la tranche concernée et la réinitialisation de toutes les tranches par la suite. La tranche concernée est remplacée par une nouvelle tranche et les tranches encore saines sont réinitialisées. Si aucune capacité n'est disponible pour allouer une tranche de remplacement, l'entraînement s'arrête.

Pour reprendre automatiquement l'entraînement après une interruption, vous devez spécifier un script de démarrage qui recherche et charge les derniers points de contrôle enregistrés. Votre script de démarrage est automatiquement exécuté chaque fois qu'une tranche est réaffectée ou qu'une VM est réinitialisée. Vous spécifiez un script de démarrage dans la charge utile JSON que vous envoyez à l'API de création de requête QR.

Le script de démarrage suivant (utilisé dans la section Créer des codes QR) vous permet de récupérer automatiquement des données après des échecs et de reprendre l'entraînement à partir de points de contrôle stockés dans un bucket Cloud Storage pendant l'entraînement de MaxText:

{
 "tpu": {
   "node_spec": [
     {
      ...
         "metadata": {
               "startup-script": "#! /bin/bash \n pwd \n runuser -l user1 -c 'cd /home/user1/MaxText && python3 MaxText/train.py MaxText/configs/base.yml run_name=run_test_failure_recovery dcn_data_parallelism=4 ici_fsdp_parallelism=8 steps=10000 save_period=10 base_output_directory='gs://user1-us-central2'' EOF"
         }
     ...
     }
   ]
 }
}

Clonez le dépôt MaxText avant de commencer.

Profilage et débogage

Le profilage est identique dans les environnements à tranche unique et à plusieurs tranches. Pour en savoir plus, consultez la page Profiler des programmes JAX.

Optimiser la formation

Segmenter avec plusieurs tranches pour des performances optimales

Pour atteindre des performances optimales dans les environnements à plusieurs tranches, vous devez réfléchir à la façon de segmenter les différentes tranches. Il existe généralement trois options (parallélisme des données, parallélisme des données entièrement segmentées et parallélisme des pipelines). Nous vous déconseillons de segmenter les activations en fonction des dimensions du modèle (parfois appelé "parallélisme de Tensor"), car cela nécessite trop de bande passante entre segments d'application. Pour toutes ces stratégies, vous pouvez conserver la même stratégie de segmentation au sein d'un segment qui a fonctionné pour vous par le passé.

Nous vous recommandons de commencer par un seul parallélisme des données. L'utilisation du parallélisme des données entièrement segmenté est utile pour libérer de l'utilisation de la mémoire. L'inconvénient est que la communication entre les tranches utilise le réseau DCN et ralentit votre charge de travail. N'utilisez le parallélisme du pipeline qu'en cas de nécessité, en fonction de la taille de lot (comme analysé ci-dessous).

Quand utiliser le parallélisme des données ?

Le parallélisme des données purement fonctionne bien dans les cas où une charge de travail fonctionne correctement, mais que vous souhaitez améliorer ses performances en effectuant un scaling sur plusieurs tranches.

Pour obtenir un scaling intensive sur plusieurs tranches, le temps nécessaire pour effectuer une réduction globale sur le réseau de contenu national doit être inférieur au temps nécessaire pour effectuer un retour arrière. DCN est utilisé pour la communication entre les tranches et constitue un facteur limitant le débit de la charge de travail.

Chaque puce TPU v4 fonctionne avec un pic de 275 x 1012 FLOPS par seconde.

Chaque hôte TPU dispose de quatre puces, et chaque hôte dispose d'une bande passante réseau maximale de 50 Gbit/s.

Cela signifie que l'intensité arithmétique est égale à 4 * 275 * 1012 FLOPS/50 Gbit/s = 22 000 FLOPS/bit.

Votre modèle utilisera entre 32 et 64 bits de bande passante DCN pour chaque paramètre et et pour chaque étape. Si vous utilisez deux tranches, votre modèle utilisera 32 bits de bande passante du réseau de neurones convolutif. Si vous utilisez plus de deux tranches, le compilateur effectue une opération de lecture aléatoire intégrale et de réduction globale. Vous utiliserez jusqu'à 64 bits de bande passante du réseau multichaîne pour chaque paramètre et chaque étape. La quantité de FLOPS nécessaire pour chaque paramètre varie en fonction du modèle. Plus précisément, pour les modèles de langage basés sur Transformer, le nombre de FLOPS requis pour un transfert avant et arrière est d'environ 6 * B * P, où:

  • B est la taille de lot en jetons
  • "P" correspond au nombre de paramètres

Le nombre de FLOPS par paramètre est de 6 * B et le nombre de FLOPS par paramètre lors du retour arrière est de 4 * B.

Pour garantir un scaling élevé sur plusieurs tranches, assurez-vous que l'intensité opérationnelle dépasse l'intensité arithmétique du matériel TPU. Pour calculer l'intensité opérationnelle, divisez le nombre de FLOPS par paramètre lors du retour arrière par la bande passante réseau (en bits) par paramètre et par étape : Operational Intensity = FLOPSbackwards_pass / DCN bandwidth

Par conséquent, pour un modèle de langage basé sur Transformer, si vous utilisez deux tranches : Operational intensity = 4 * B / 32

Si vous utilisez plus de deux segments: Operational intensity = 4 * B/64

Cela suggère une taille de lot minimale comprise entre 176 000 et 352 000 pour les modèles de langage basés sur Transformer. Étant donné que le réseau DCN peut supprimer brièvement des paquets, il est préférable de maintenir une marge d'erreur importante. Ne déployez le parallélisme des données que si la taille de lot par pod est d'au moins 350 000 (deux pods) à 700 000 (plusieurs pods).

Pour les autres architectures de modèle, vous devez estimer l'environnement d'exécution de votre pass inverse par tranche (en le chronométrant à l'aide d'un profileur ou en comptant les FLOPS). Vous pouvez ensuite comparer cette durée à la durée d'exécution prévue afin de réduire le surnombre de points de terminaison du réseau et d'obtenir une bonne estimation de l'intérêt du parallélisme des données pour vous.

Quand utiliser le parallélisme des données entièrement segmentées (FSDP) ?

Le parallélisme des données entièrement segmenté (FSDP) combine ce parallélisme des données (segmentation des données entre les nœuds) et la segmentation des pondérations entre les nœuds. Pour chaque opération des passes avant et arrière, les pondérations sont toutes rassemblées de sorte que chaque tranche dispose des pondérations dont elle a besoin. Au lieu de synchroniser les gradients à l'aide de la fonction "all-reduce", les gradients sont dispersés en fonction du "réduit" à mesure qu'ils sont produits. De cette manière, chaque tranche ne reçoit que les gradients correspondant aux pondérations dont elle est responsable.

Comme pour le parallélisme des données, le FSDP nécessitera un scaling de la taille de lot globale de manière linéaire avec le nombre de tranches. Le FSDP réduit la pression sur la mémoire à mesure que vous augmentez le nombre de tranches. En effet, le nombre de pondérations et d'états de l'optimiseur par tranche diminue, mais cela entraîne une augmentation du trafic réseau et une plus grande probabilité de blocage en raison d'un retard collectif.

En pratique, la FSDP sur les tranches est préférable si vous augmentez le lot par tranche, stockez davantage d'activations afin de minimiser la rematérialisation lors de la rétrogradation ou si vous augmentez le nombre de paramètres dans votre réseau de neurones.

Les opérations all-collecter et all-reduce dans FSDP fonctionnent de la même manière que dans DP, vous pouvez donc déterminer si votre charge de travail FSDP est limitée par les performances DCN de la même manière que décrit dans la section précédente.

Quand utiliser le parallélisme de pipeline ?

Le parallélisme du pipeline devient pertinent lorsque vous atteignez des performances élevées avec d'autres stratégies de parallélisme qui nécessitent une taille de lot globale supérieure à votre taille de lot maximale préférée. Ce parallélisme permet aux tranches d'un pipeline de "partager" un lot. Cependant, le parallélisme des pipelines présente deux inconvénients majeurs:

  1. Cela entraîne une "bulle de pipeline" où les puces sont inactives, car elles attendent des données.
  2. Il nécessite un microtraitement par lot, qui diminue la taille effective du lot et l'intensité arithmétique, et, en fin de compte, modélise l'utilisation du FLOP.

Le parallélisme de pipeline ne doit être utilisé que si les autres stratégies de parallélisme nécessitent une taille de lot globale trop importante. Avant d'essayer le parallélisme des pipelines, il est utile d'expérimenter de manière empirique si la convergence par échantillon ralentit au niveau de la taille de lot nécessaire pour atteindre une FSDP hautes performances. Le FSDP tend à atteindre une utilisation plus élevée du FLOP du modèle, mais si la convergence par échantillon ralentit à mesure que la taille de lot augmente, le parallélisme du pipeline peut toujours être le meilleur choix. La plupart des charges de travail peuvent tolérer des tailles de lot suffisamment importantes pour ne pas bénéficier du parallélisme des pipelines, mais votre charge de travail peut être différente.

Si le parallélisme du pipeline est nécessaire, nous vous recommandons de le combiner au parallélisme des données ou FSDP. Cela vous permettra de minimiser la profondeur du pipeline tout en augmentant la taille de lot par pipeline jusqu'à ce que la latence du DCN devienne un facteur de débit inférieur. Concrètement, si vous avez N tranches, envisagez d'utiliser des pipelines de profondeur 2 et N/2 d'instances répliquées de parallélisme des données, puis des pipelines de profondeur 4 et N/4 d'instances répliquées de parallélisme des données, et ainsi de suite, jusqu'à ce que le lot par pipeline devienne suffisamment volumineux pour que les collectifs DCN puissent être cachés derrière l'arithmétique dans la passe arrière. Cela minimisera le ralentissement induit par le parallélisme du pipeline, tout en vous permettant d'évoluer au-delà de la limite globale de taille de lot.

Bonnes pratiques concernant les tranches d'âge

Chargement des données…

Pendant l'entraînement, nous chargeons des lots de manière répétée à partir d'un ensemble de données pour alimenter le modèle. Il est important de disposer d'un chargeur de données asynchrone efficace qui segmente le lot sur plusieurs hôtes afin d'éviter de priver les TPU de travail. Dans le chargeur de données actuel de MaxText, chaque hôte charge un sous-ensemble égal des exemples. Cette solution est adaptée au texte, mais nécessite une nouvelle segmentation dans le modèle. De plus, MaxText ne propose pas encore d'instantané déterministe, qui permettrait à l'itérateur de données de charger les mêmes données avant et après la préemption.

Points de contrôle

La bibliothèque de points de contrôle Orbax fournit des primitives permettant de créer des points de contrôle des PyTrees JAX vers un espace de stockage local ou Google Cloud Storage. Nous fournissons une intégration de référence avec des points de contrôle synchrones dans MaxText dans checkpointing.py.

Configurations compatibles

Formes

Toutes les tranches doivent avoir la même forme (par exemple, le même AcceleratorType). Les formes de tranches hétérogènes ne sont pas acceptées.

Orchestration

L'orchestration est compatible avec GKE. Pour en savoir plus, consultez la section TPU dans GKE.

Frameworks

Le multislice n'est compatible qu'avec les charges de travail JAX et PyTorch.

Parallélisme

Nous recommandons aux utilisateurs de tester les tranches multitranches avec le parallélisme des données. Pour en savoir plus sur la mise en œuvre du parallélisme de pipelines avec des tranches multislices, contactez votre responsable de compte Google Cloud.

Assistance et commentaires

Tous vos commentaires sont les bienvenus. Pour nous faire part de vos commentaires ou demander de l'aide, contactez-nous à l'aide du formulaire de commentaires ou d'assistance Cloud TPU.