Présentation de Cloud TPU en multitranche

Cloud TPU multislice est une technologie de mise à l'échelle des performances full stack qui permet à une tâche d'entraînement d'utiliser plusieurs tranches TPU dans un seul pod ou sur des tranches de plusieurs pods avec un parallélisme des données simple. 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 chips, une seule tranche peut offrir les meilleures performances. Cependant, plusieurs tranches de petite taille sont plus facilement disponibles, ce qui accélère le démarrage lorsque l'application Multislice est utilisée avec des tranches plus petites.

Évolutivité linéaire des performances avec plusieurs secteurs

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

Flux de données à plusieurs tranches

Les développeurs n'ont pas à écrire de code pour implémenter la communication DCN entre tranches. 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 tranche d'une requête à plusieurs tranches est du même type d'accélérateur. Un type d'accélérateur se compose 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 sera pas terminée tant que le matériel ne sera pas disponible. Une fois la nouvelle tranche créée, toutes les autres tranches de l'environnement multitranche sont redémarrées pour que l'entraînement puisse se poursuivre.Avec un script de démarrage correctement configuré, ce script peut automatiquement être relancé automatiquement sans l'intervention de l'utilisateur, en chargeant et en reprenant à 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 des centres de données (DCN)
Réseau à latence plus élevée et à faible débit (par rapport à ICI) qui relie les tranches TPU dans une configuration à plusieurs tranches.
Planification de groupes
Lorsque toutes les tranches TPU sont provisionnées ensemble, en même temps, cela garantit que toutes les tranches ou aucune des tranches sont provisionnées correctement.
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 effectuer 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
Au moins deux tranches de puce TPU peuvent communiquer via DCN.
Nœud
Dans le contexte de plusieurs tranches, le nœud fait référence à une seule tranche TPU. Chaque tranche TPU d'une multitranche reçoit 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, qui permet de mettre en file d'attente et de gérer une requête pour un environnement TPU à une seule tranche ou à plusieurs tranches.
Script de démarrage
Script de démarrage Compute Engine standard qui s'exécute chaque fois qu'une VM est démarrée ou redémarrée. Pour les multitranches, il est spécifié dans la requête de création de code QR. Pour en savoir plus sur les scripts de démarrage de 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'une tranche communiquent entre eux via le réseau ICI.
VM TPU
Machine virtuelle exécutant Linux et ayant accès aux TPU sous-jacents. Pour les TPU v4, chaque VM TPU a un accès direct à quatre puces. Une VM TPU est parfois appelée 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)
Chip d'accélération du ML développé en interne par Google. Ils sont conçus pour offrir des calculs rapides et économes en énergie pour les tâches de machine learning clés telles que la multiplication matricielle.
Types de capacités de Cloud TPU

Les TPU peuvent être créés à partir de trois types de capacité (voir les 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 lorsque vous créez vos ressources.
  • Préemptifs: cible un quota préemptif Vos ressources peuvent être préemptées pour laisser la place aux requêtes pour une tâche de priorité supérieure. Utilisez l'option --best-effort lorsque vous créez 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 sera 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.

Premiers pas

Si vous n'avez jamais utilisé de TPU, commencez par installer Google Cloud CLI et configurer votre environnement Cloud TPU. Pour utiliser les multitranches, 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 migrer votre réservation 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 le code du dépôt GitHub MaxText. MaxText est un LLM de base hautes performances, arbitrairement évolutif, Open Source et bien testé, écrit en Python et en Jax. MaxText a été conçu pour un entraînement efficace sur Cloud TPU.

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

Parallélisme ICI

ICI fait référence à l'interconnexion à haut débit qui connecte les TPU en une seule tranche. La segmentation ICI correspond au 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 de sorte que ici_data_parallelism * ici_fsdp_parallelism * ici_tensor_parallelism soit égal au nombre de chips de 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 4-8:

ici_data_parallelism ici_fsdp_parallelism ici_tensor_parallelism
FSDP quadridirectionnel 1 4 1
Parallélisme Tensor à quatre voies 1 1 4
Parallélisme FSDP bidirectionnel + parallélisme Tensor bidirectionnel 1 2 2

Notez que, dans la plupart des cas, ici_data_parallelism doit rester égal à 1, car le réseau ICI est suffisamment rapide pour préférer presque toujours 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 indiqué dans la section Effectuer un calcul sur une VM Cloud TPU à l'aide de JAX. Cet exemple montre comment exécuter shardings.py sur un seul secteur.

  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 laisser un mot de passe vide (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 |--best-effort]

    Description des options de commande

    your-qr-id
    Chaîne définie par l'utilisateur qui identifie la requête 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
    La [version du logiciel 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 d'optimisation lors de la création des tranches [par défaut].

    La Google Cloud CLI n'est pas compatible avec toutes les options de création de code QR, telles que les tags. Pour en savoir plus, consultez Créer des codes QR.

  4. Attendez que le code QR passe à l'état ACTIVE, ce qui signifie que les nœuds de calcul passent à 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 alpha 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) vers 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 consulter les résultats dans les journaux. Vos TPU devraient atteindre environ 260 TFLOP par seconde, soit une utilisation FLOP impressionnante de plus de 90 % ! Dans ce cas, nous avons sélectionné approximativement le nombre maximal de lots pouvant être compris dans la mémoire à haut débit (HBM) du TPU.

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

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

  10. Lorsque vous avez terminé, supprimez le code QR et la tranche TPU. 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'indicateur --async facultatif.

    $ gcloud alpha 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 DCN, correspondant au nombre de segments de chaque type de parallélisme de 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 secteurs.

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

dcn_data_parallelism dcn_fsdp_parallelism dcn_tensor_parallelism Nbre de segments
Parallélisme des données bidirectionnel 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 exécutées 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 vous conservez la taille constante de la taille de tranche et du lot par tranche), 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 du document Getting Started: Quick Experiments on Multiple segments du dépôt MaxText, si ce n'est qu'ici, nous exécutons shardings.py au lieu du LLM plus complexe dans train.py.

L'outil multihost_runner.py est optimisé pour les tests rapides, en réutilisant de manière répétée 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, plusieurs heures ou 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 composent vos tranches. Vous pouvez exécuter multihost_runner.py sur une machine locale ou sur n'importe quelle VM Compute Engine dans le 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 qu'il s'exécute sur d'autres TPU.

Configurer votre environnement

  1. Clonez MaxText sur votre exécuteur.

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

  3. Créez des clés SSH pour gcloud. Nous vous recommandons de laisser un mot de passe vide (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 multitranche v4. Pour utiliser la version v5e, spécifiez un accelerator-type v5e (par exemple, v5litepod-16) et la version v5e runtime-version (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|--best-effort]

    Description des options de commande

    your-qr-id
    Chaîne définie par l'utilisateur qui identifie la requête 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 Cloud TPU.
    node-count
    Nombre de secteurs à créer.
    node-prefix
    Préfixe utilisé pour générer les noms de chaque secteur. Un nombre est ajouté au préfixe de chaque secteur. 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 d'optimisation lors de la création des tranches [par défaut].

  6. Au début du provisionnement du code QR, l'opération peut prendre jusqu'à cinq minutes en fonction de 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 alpha compute tpus queued-resources list \
--filter=your-qr-id

    Un résultat semblable aux lignes suivantes doit s'afficher:

    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 verrez environ 230 TFLOPs par seconde de performances dans les fichiers journaux.

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

Effectuer le scaling d'une charge de travail pour passer à une multitranche

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

Il s'agit des seules modifications de code nécessaires lorsque vous passez à la multitranche. Pour obtenir des performances élevées, DCN doit être mappé sur des données parallèles de données, entièrement segmentées ou sur des axes parallèles du pipeline. Les considérations de performances et les stratégies de segmentation sont traité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 chips dans votre environnement à plusieurs tranches. Par exemple, si vous utilisez quatre tranches de v4-16, vous avez huit chips par tranche * 4 tranches. Par conséquent, len(jax.devices()) doit renvoyer 32.

Choisir la taille des tranches pour les environnements à plusieurs tranches

Pour accélérer l'affichage d'une ligne linéaire, ajoutez des tranches de la même taille que votre tranche existante. Par exemple, si vous utilisez une tranche v4-512, la multitranche obtiendra 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 Partitionner avec plusieurs tranches pour des performances maximales.

Exécuter votre tâche sur plusieurs tranches

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

  1. Avec le script d'exécution de tests multihost_runner.py
  2. À l'aide du script d'exécution de production multihost_job.py
  3. Approche manuelle

Script d'exécution de tests

Le script multihost_runner.py distribue le code dans un environnement multislice existant, exécute la 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 de 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 durée relativement courte. 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 face aux défaillances matérielles et à d'autres préemptions, il est préférable d'intégrer directement l'API Create Queued Resource. À titre d'exemple fonctionnel, nous fournissons multihost_job.py, qui déclenche l'appel de l'API Created Queued Resource avec le script de démarrage approprié pour exécuter l'entraînement et reprendre la préemption. Le script multihost_job.py est documenté dans le fichier README de 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 à secteurs multiples. 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 ID 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-Ubuntu2204-base
ACCELERATOR_TYPE v4-16
EXAMPLE_TAG_1, EXAMPLE_TAG_2... Tags utilisés pour identifier les sources ou les cibles valides pour les pare-feu de réseau
SLICE_COUNT Nombre de secteurs. Limité à un maximum de 256 secteurs.
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 elle est réparée ou réinitialisée.

Créer une requête QR avec 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|--best-effort]

Description des options de commande

your-qr-id
Chaîne définie par l'utilisateur qui identifie la requête QR.
project
Chaîne définie par l'utilisateur qui identifie la requête QR.
zone
Zone Google Cloud dans laquelle créer le code QR.
node-count
Nombre de secteurs à 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 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.
best-effort
Utilisez le quota d'optimisation lors de la création des tranches [par défaut].

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

Créer une requête QR avec 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 seule tranche
  • tpu-vm-runtime-version : versions d'exécution de la 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 (facultatif)
  • your-startup-script : script de démarrage qui s'exécutera 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 en savoir plus sur toutes les options disponibles.

Pour utiliser la capacité préemptive, remplacez:

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

Vous pouvez également supprimer 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 comme suit:

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

Utilisez la valeur GUID à la fin de la valeur de chaîne de l'attribut name pour 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/v2alpha1/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.v2alpha1.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 le provisionnement commencera lorsque les ressources seront suffisantes. Si le code QR est à l'état FAILED, le motif de l'échec est 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 la commande describe QRs pour surveiller les étapes du code QR.

Dans de rares cas, votre code QR peut présenter l'état FAILED, tandis que certaines tranches peuvent indiquer ACTIVE. Dans ce cas, supprimez les ressources qui ont été créées, puis réessayez dans quelques minutes ou contactez l'équipe Cloud TPU pour résoudre le problème.

Se connecter en SSH et installer les dépendances

L'article Exécuter du code JAX sur des tranches de pod TPU explique comment se connecter à vos VM TPU à l'aide de SSH sur une seule tranche. Pour vous connecter à toutes les VM TPU de votre environnement à plusieurs tranches via SSH et installer des dépendances, utilisez 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 tous les nœuds en QR à l'aide de SSH. La commande est regroupée par groupes de quatre et envoyée simultanément. Le lot de commandes suivant 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'interrompt et aucun autre lot n'est envoyé. Pour en savoir plus, consultez la documentation de référence de l'API pour les ressources en file d'attente. Si le nombre de tranches que vous utilisez dépasse la limite de thread 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 ordinateur local 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 les divisera en lots. Elle exécute le script d'entraînement sur le premier lot de 64 tranches et attend la fin des scripts avant de les exécuter sur le lot de 36 tranches restant. 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 effectuez cette opération, après avoir démarré le script d'entraînement sur le premier lot de tranches, le contrôle revient immédiatement à la commande SSH. La commande SSH peut ensuite commencer à exécuter le script d'entraînement sur le lot de 36 tranches restant. Vous devez acheminer 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 à plusieurs tranches. 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 tranches de la version 4-16, il y a huit chips par tranche et quatre tranches. Le jax.devices() devrait renvoyer un total de 32 puces (appareils).

Lister les codes QR

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

$ gcloud alpha 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 détaillée et l'état d'un code QR, utilisez l'API QR describe. Vous pouvez appeler cette API à l'aide de gcloud ou curl.

En utilisant gcloud :

$ gcloud alpha 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/v2alpha1/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 votre 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éinitialiser les codes QR

L'API ResetQueuedResource peut être utilisée pour réinitialiser toutes les VM via un code QR ACTIVE. La réinitialisation forcée des VM efface la mémoire des machines et rétablit leur état initial. 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 alpha 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'indicateur --async facultatif.

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

Reprise automatique après échec

En cas de perturbation, la fonctionnalité Multislice permet de réparer sans intervention de la tranche affectée, puis de réinitialiser toutes les tranches par la suite. La tranche affectée est remplacée par une nouvelle tranche, et les tranches restantes saines sont reset. Si aucune capacité n'est disponible pour allouer une tranche de remplacement, l'entraînement s'arrête.

Pour reprendre l'entraînement automatiquement après une interruption, vous devez spécifier un script de démarrage qui vérifie et charge les derniers points de contrôle enregistrés. Votre script de démarrage s'exécute automatiquement 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 les échecs et de reprendre l'entraînement à partir de points de contrôle stockés dans un bucket Cloud Storage pendant l'entraînement 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 d'essayer cette solution.

Profilage et débogage

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

Optimiser la formation

Segmentation avec multitranche pour des performances optimales

Pour atteindre des performances optimales dans les environnements à plusieurs tranches, vous devez déterminer comment 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é et parallélisme des pipelines. Nous vous déconseillons de segmenter les activations entre les dimensions du modèle (parfois appelée "parallélisme de Tensor"), car cela nécessite trop de bande passante entre les tranches. Pour toutes ces stratégies, vous pouvez conserver la même stratégie de segmentation dans une tranche qui a fonctionné pour vous par le passé.

Nous vous recommandons de commencer par le parallélisme pur des données. L'utilisation du parallélisme des données entièrement segmenté est utile pour libérer 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 que lorsque cela est nécessaire en fonction de la taille de lot (comme analysé ci-dessous).

Quand utiliser le parallélisme des données

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

Pour obtenir un scaling efficace sur plusieurs tranches, la durée nécessaire pour effectuer une réduction complète sur DCN doit être inférieure au temps nécessaire pour effectuer une rétrogradation. Le DCN est utilisé pour la communication entre les tranches et constitue un facteur limitant le débit des charges de travail.

Chaque puce TPU v4 offre un pic de 275 * 1012 FLOPS par seconde.

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

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

Votre modèle utilisera 32 à 64 bits de bande passante DCN pour chaque paramètre et par étape. Si vous utilisez deux tranches, votre modèle utilisera 32 bits de bande passante DCN. Si vous utilisez plus de deux tranches, le compilateur effectue une opération de brassage complet de réduction complète, et vous utilisez jusqu'à 64 bits de bande passante DCN pour chaque paramètre et par étape. La quantité de FLOPS nécessaire pour chaque paramètre varie en fonction de votre modèle. Plus précisément, pour les modèles de langage basés sur Transformer, le nombre de FLOPS requis pour une propagation 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 de la rétropropagation est de 4 * B.

Pour garantir un scaling efficace sur plusieurs tranches, assurez-vous que l'intensité opérationnelle dépasse l'intensité arithmétique du matériel TPU. Pour calculer l'intensité de fonctionnement, divisez le nombre de FLOPS par paramètre lors de la rétropropagation par la bande passante réseau (en bits) par paramètre et par pas : 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 rapidement supprimer des paquets, il est préférable de maintenir une marge d'erreur significative, en déployant le parallélisme des données uniquement si la taille de lot par pod est comprise entre 350 000 (deux pods) et 700 000 (nombreux pods).

Pour les autres architectures de modèle, vous devez estimer la durée d'exécution de votre rétrogradation par tranche (soit en la chronométrant à l'aide d'un profileur, soit en comptant les FLOPS). Vous pouvez ensuite comparer ce nombre à la durée d'exécution prévue pour réduire le volume de trafic continu et obtenir une bonne estimation de la pertinence du parallélisme des données.

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

Le parallélisme des données entièrement segmenté (FSDP) combine le parallélisme des données (partitionnement des données entre les nœuds) et la segmentation des pondérations entre les nœuds. Pour chaque opération dans les passes avant et arrière, les pondérations sont toutes collectées afin que chaque tranche possède les pondérations dont elle a besoin. Au lieu de synchroniser les gradients à l'aide de "all-reduce", les gradients sont "réduits" et dispersés au fur et à mesure de leur production. De cette manière, chaque tranche n'obtient que les gradients des pondérations dont elle est responsable.

Comme pour le parallélisme des données, FSDP nécessite un scaling linéaire de la taille de lot globale en fonction du nombre de tranches. 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 d'optimisation par tranche diminue, mais au prix d'une augmentation du trafic réseau et d'une plus grande possibilité de blocage en raison d'un collectif retardé.

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

Les opérations tout-gather et all-reduce du FSDP fonctionnent de la même manière que celles de la 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 celle décrite dans la section précédente.

Quand utiliser le parallélisme du pipeline ?

Le parallélisme des pipelines 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. Le parallélisme des pipelines permet aux tranches composant un pipeline de "partager" un lot. Cependant, le parallélisme des pipelines présente deux inconvénients importants:

  1. Elle entraîne une "bulle de pipeline" dans laquelle les chips sont inactifs, car ils attendent des données.
  2. Elle nécessite un microtraitement par lot qui réduit la taille effective du lot, l'intensité arithmétique et, au final, l'utilisation de FLOP.

Le parallélisme des pipelines 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'effectuer des tests pour voir empiriquement si la convergence par échantillon ralentit au niveau de la taille de lot nécessaire pour atteindre une FSDP hautes performances. FSDP tend à atteindre une utilisation plus élevée de FLOP pour le 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 des pipelines est nécessaire, nous vous recommandons de l'associer au parallélisme des données (FSDP). Cela vous permettra de réduire la profondeur du pipeline tout en augmentant la taille de lot par pipeline jusqu'à ce que la latence DCN devienne moins importante pour le débit. Concrètement, si vous avez N tranches, considérez les pipelines avec une profondeur de 2 et N/2 instances répliquées de parallélisme de données, puis les pipelines de profondeur 4 et N/4 du parallélisme des données, et ainsi de suite, jusqu'à ce que le lot par pipeline soit suffisamment volumineux pour que les collectifs DCN puissent être cachés derrière l'arithmétique dans la rétrogradation. Cela minimisera le ralentissement introduit par le parallélisme du pipeline tout en vous permettant de dépasser la limite globale de taille de lot.

Bonnes pratiques concernant les secteurs multitranches

Chargement des données…

Pendant l'entraînement, nous chargeons de manière répétée des lots à partir d'un ensemble de données pour alimenter le modèle. Il est important de disposer d'un chargeur de données asynchrone et efficace, qui segmente le lot entre les hôtes, afin d'éviter de priver les TPU de travail. Dans le chargeur de données actuel de MaxText, chaque charge d'hôte correspond à un sous-ensemble égal d'exemples. Cette solution est adaptée pour le texte, mais nécessite une segmentation du modèle. De plus, MaxText n'offre pas encore de création d'instantanés 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 de point de contrôle de PyTrees JAX vers un stockage local ou Google Cloud. 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 secteurs 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

La multitranche n'est compatible qu'avec les charges de travail JAX et PyTorch.

Parallélisme

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

Assistance et commentaires

Tous vos commentaires sont les bienvenus ! Pour partager vos commentaires ou demander de l'aide, contactez-nous via le formulaire d'assistance ou de commentaires pour Cloud TPU.