Présentation des instances multitranches Cloud TPU

Cloud TPU Multislice est une technologie de scaling des performances complète qui permet à une tâche d'entraînement d'utiliser plusieurs tranches de TPU dans un seul pod ou sur des tranches de plusieurs pods avec un parallélisme de données simple. Avec les puces TPU v4, signifie que les tâches d'entraînement peuvent utiliser plus de 4 096 puces en une seule exécution. Pour l'entraînement emplois nécessitant moins de 4 096 chips, une seule tranche peut offrir le meilleur des performances. Cependant, plusieurs tranches plus petites sont plus facilement disponibles, ce qui permet d'accélérer le démarrage lorsque Multislice est utilisé avec des tranches plus petites.

Plusieurs segments effectuent un scaling linéaire des performances

En cas de déploiement dans des configurations multitranches, les puces TPU de chaque tranche de communiquer via ICI (Inter-Chip-Interconnect). les puces TPU de différentes les tranches communiquent en transférant des données aux CPU (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 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 pour des performances optimales.

Concepts

Type d'accélérateur
Forme de chaque tranche TPU comprenant une multitranche. Chaque dans 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 tranche. Dans les rares cas où ressources insuffisantes pour créer une tranche, la création ne se terminera pas jusqu'à ce que le matériel soit disponible. Une fois la tranche créée, toutes les autres segments dans l'environnement à plusieurs tranches sont redémarrés. L'entraînement peut continuer.Avec un script de démarrage correctement configuré, le script d'entraînement peut redémarrer automatiquement sans intervention de l'utilisateur, sans chargement ni 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.
Réseaux de centres de données (DCN)
Un réseau à latence plus élevée et à bas débit (comparé à ICI) connecte les tranches TPU dans une configuration à plusieurs tranches.
Planification en groupe
Lorsque toutes les tranches TPU sont provisionnées ensemble, en même temps, toutes les tranches sont provisionnées avec succès ou aucune ne l'est.
Organisateur
Un hôte est un ordinateur physique qui exécute des VM. Un hôte peut exécuter quatre VM au maximum en même temps. Chaque VM dispose d'un TPU dédié.
Inférence
Chargez un modèle de machine learning pré-entraîné sur un hôte et effectuez des prédictions sur des données.
Interchip Interconnect (ICI)
Liens internes haut débit et à faible latence qui connectent les TPU au sein d'un pod TPU.
Multislice
Deux tranches de puces TPU ou plus pouvant communiquer via le DCN.
Nœud
Dans le contexte à plusieurs tranches, le nœud fait référence à une seule tranche TPU. Chaque La tranche TPU d'une tranche multicouche 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 permettant de mettre en file d'attente et de gérer une requête pour un un environnement TPU à une 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 Multislice, il est spécifié dans la requête de création du code QR. Pour plus d'informations 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 dans une tranche d'accès communiquent entre elles via le réseau ICI.
VM TPU
Machine virtuelle exécutant Linux qui a accès aux TPU sous-jacents. Pour TPU v4, chaque VM TPU dispose d'un accès direct à quatre puces. Parfois, on appelle un TPU VM en tant que 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.
TPU (Tensor Processing Unit)
Puce d'accélération de ML développée en interne par Google. Ils sont conçus pour offrent des fonctionnalités de calcul rapides et économes en énergie pour les tâches de machine learning clés telles que la multiplication matricielle.
Types de capacité Cloud TPU

Les TPU peuvent être créés à partir de différents types de capacités (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 disposer d'un contrat de réservation avec Google. Utilisez l'option --reserved lors de la création vos ressources.
  • Spot: cible un quota préemptif à l'aide de VM Spot. Votre ressources peuvent être préemptées pour faire de la place aux demandes d'augmentation une tâche prioritaire. Utilisez l'option --spot lorsque vous créez vos ressources.
  • À la demande: cible le quota à la demande, qui ne nécessite pas de réservation. et ne seront pas préemptées. La requête TPU sera mise en file d'attente dans un pipeline proposée par Cloud TPU, la disponibilité des ressources n'est pas sont garantis. Cette option est sélectionnée par défaut, aucun indicateur n'est nécessaire.

Commencer

Si vous n'avez jamais utilisé de TPU, commencez par installer la Google Cloud CLI. et configurer votre environnement Cloud TPU. Pour utiliser Sur une tranche multitranche, 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 dans 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 une solution Open Source hautement performante, arbitrairement évolutive et testée un LLM de base écrit en Python et en Jax. MaxText a été conçu pour s'entraîner efficacement Cloud TPU.

Le code dans shardings.py est conçu pour vous aider à tester différents parallélismes options. par exemple, le parallélisme des données, le FSDP (entièrement segmenté) et le parallélisme des Tensors. Le code passe d'une tranche à une autre de l'infrastructure.

Parallélisme ICI

ICI désigne l'interconnexion haut débit qui connecte les TPU dans une tranche. 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 soumises à des contraintes ici_data_parallelism * ici_fsdp_parallelism * ici_tensor_parallelism est égal(e) à 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 chips disponibles dans les versions 4-8:

ici_data_parallelism ici_fsdp_parallelism ici_tensor_parallelism
FSDP quadridirectionnel 1 4 1
Parallélisme du Tensor quadridirectionnel 1 1 4
Parallélisme FSDP à deux voies et parallélisme de tenseur à deux voies 1 2 2

Notez que ici_data_parallelism doit rester défini sur 1 dans la plupart des cas, car le ICI est assez 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 tranche TPU unique. 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 laisser le mot de passe vide (appuyez deux fois sur Entrée après avoir exécuté la commande suivante). Si vous êtes invité google_compute_engine fichier 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 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 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 de code QR.
    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.

    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 dont l'état est READY. Une fois le provisionnement du code QR lancé, 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 ne comporte qu'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 d'installation pour installer JAX et autres dépendances sur votre tranche TPU. Le 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 adapté aux TPU Mémoire à bande passante (HBM).

  9. N'hésitez pas à explorer d'autres stratégies de segmentation. par exemple, vous pouvez 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 prendra deux à cinq minutes, Elle peut s'exécuter 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 du réseau multichaîne

Le script shardings.py utilise trois paramètres qui spécifient le parallélisme du réseau multichaîne : correspondant au nombre de segments pour 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 dcn_data_parallelism * dcn_fsdp_parallelism * dcn_tensor_parallelism est égal(e) à le nombre de tranches.

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

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

dcn_tensor_parallelism doit toujours être défini sur 1, car le DCN est mauvais sont parfaitement adaptés à ce type de segmentation. Pour les charges de travail LLM classiques sur les 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 segments, mais c'est de l'application.

Au fur et à mesure que vous augmentez le nombre de tranches (en supposant que vous conservez la taille de tranche par constante de tranche), vous augmentez la quantité de parallélisme des données.

Exécuter shardings.py dans un environnement Multislice

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 Premiers pas: tests rapides sur plusieurs segments du dépôt MaxText, sauf ici, nous exécutons shardings.py au lieu de LLM plus complexes dans train.py.

L'outil multihost_runner.py est optimisé pour des tests rapides, et ce de manière répétée à réutiliser les mêmes TPU. Comme le script multihost_runner.py dépend de des connexions SSH de longue durée, nous ne le recommandons pas pour les jobs de longue durée. Si vous souhaitez exécuter un job plus long (plusieurs heures ou jours, par exemple), nous vous recommandons utiliser multihost_job.py.

Dans ce tutoriel, nous utilisons le terme exécuteur pour désigner la machine sur laquelle vous exécutez le script multihost_runner.py. Nous utilisons le terme workers pour indiquer des VM TPU constituant vos tranches. Vous pouvez exécuter multihost_runner.py sur une instance ou n'importe quelle VM Compute Engine dans le même projet que vos tranches. L'exécution de multihost_runner.py sur un worker n'est pas prise en charge.

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 des 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 exécuteur. machine.

  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 sur saisissez deux fois après avoir exécuté la commande suivante). Si vous êtes invité le fichier google_compute_engine existe déjà, choisissez de ne pas conserver votre 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 Multislice à l'aide de queued-resources create.

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

      $ gcloud 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 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 de chaque tranche. Par exemple, si vous définissez node-prefix à mySlice, les tranches sont nommées: mySlice-0, mySlice-1, en continuant numériquement pour chaque tranche.
    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.

  6. Lorsque le provisionnement du code QR commence, cela peut prendre jusqu'à cinq minutes en fonction de la taille du code QR. Attendre que la ressource en file d'attente (QR) soit dans l'état ACTIVE. Vous pouvez vérifier l'état d'une demande de code QR à l'aide du 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 WAITING_FOR_RESOURCES ou PROVISIONING pendant 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 TFLOP par seconde de performances dans le journal .

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

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

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

Il s'agit des seules modifications de code nécessaires lors de la migration vers Multislice. Pour obtenir de hautes performances, le réseau de contenu doit être mappé sur des données en parallèle, entièrement des données segmentées ou des axes parallèles de pipeline. Les considérations relatives aux performances les stratégies de segmentation sont abordées plus en détail dans Segmentez vos données à l'aide de plusieurs tranches pour des performances optimales.

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

Choisir la taille des tranches pour les environnements à plusieurs tranches

Pour accélérer linéairement, ajoutez des tranches de la même taille que tranche. Par exemple, si vous utilisez une tranche v4-512, Multislice doublera environ les performances en ajoutant une deuxième tranche v4-512 et en doublant votre taille de lot globale. Pour en savoir plus, consultez Segmentez vos données à l'aide de plusieurs tranches pour des performances optimales.

Exécuter votre job sur plusieurs tranches

Il existe trois approches différentes pour exécuter votre charge de travail personnalisée dans un environnement multislice :

  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 de test

multihost_runner.py distribue le code dans un environnement multisegment existant et exécute votre commande sur chaque hôte, copie vos journaux et suit les erreurs de chaque commande état. Le script multihost_runner.py est documenté dans Fichier README MaxText.

Comme multihost_runner.py gère des connexions SSH persistantes, il n'est utilisé convient aux tests de petite taille et relativement courts. 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éemptives, il est préférable de procéder à une intégration directe API. À titre d'exemple, nous fournissons multihost_job.py, ce qui déclenche l'appel d'API Created Queued Resource avec le démarrage approprié pour exécuter l'entraînement et le reprendre en cas de préemption. multihost_job.py est décrite dans Fichier README MaxText.

Étant donné que multihost_job.py doit provisionner des ressources pour chaque exécution, il ne fournit pas un 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 votre configuration multi-tranche. Toutefois, si vous préférez provisionner et gérer votre environnement directement à l'aide de commandes QR, consultez 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 la dans 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 attribué par l'utilisateur au code QR.
PROJECT 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 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 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 le 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 ci-dessous.

{
  "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 : facultatif, chaîne de balise arbitraire
  • 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

Pour en savoir plus, consultez la page sur l'API REST Queued Resource. pour connaître toutes les options disponibles.

Pour utiliser la capacité Spot, remplacez les éléments suivants :

"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 que l'attribut name permette d'obtenir sur la demande de code QR.

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

Pour obtenir l'état de la requête de code 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 la sortie du Demande 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 Le champ response doit être WAITING_FOR_RESOURCES ou FAILED. Si le code QR est à l'état WAITING_FOR_RESOURCES, cela signifie qu'il a été mis en file d'attente et commencera le provisionnement dès qu'il y aura suffisamment de ressources. Si le code QR est à l'état FAILED, le motif de l'échec est indiqué dans le résultat. Pour plus des informations sur les autres états possibles, consultez la Guide de l'utilisateur des ressources en file d'attente

Une fois l'opération terminée, utilisez la commande describe QR pour surveiller les étapes du code QR.

Dans de rares cas, votre code QR peut être à l'état FAILED alors que certains segments sont ACTIVE. Dans ce cas, supprimez les ressources créées, puis réessayez dans quelques minutes ou contactez l'équipe Cloud TPU pour résoudre le problème.

SSH et installation des dépendances

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 Multislice 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 les nœuds de QR à l'aide de SSH. La commande est regroupée en groupes de quatre et est envoyée simultanément. Le lot suivant de commandes est envoyé lorsque le lot actuel termine l'exécution. En cas de défaillance de l'une des commandes, le traitement s'arrête et aucun autre lot n'est envoyé. Pour en savoir plus, consultez les 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 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 est de 64. Si vous essayez d'exécuter d'entraînement sur plus de 64 tranches, disons 100, la commande SSH va diviser les tranches en lots. Il exécutera le script d'entraînement sur le premier lot de 64 segments et attendre que les scripts soient terminés avant de les exécuter sur le le lot restant de 36 tranches. Cependant, le premier lot de 64 tranches ne peut pas jusqu'à ce que les 36 tranches restantes exécutent le script, ce qui provoque une 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'indicateur --command. Lorsque vous effectuez cette opération, après avoir lancé le script d'entraînement sur le premier lot de tranches, le contrôle reviendra immédiatement 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 devrez rediriger vos stdout et stderr est diffusé de manière appropriée lors de l'exécution des commandes en arrière-plan. Pour augmenter parallélisme dans le même code QR, vous pouvez sélectionner des segments spécifiques à l'aide de --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 Exécutez le code JAX sur des tranches de pod TPU. Déclarez-vous len(jax.devices()) est égal au nombre de chips dans votre multitranche. environnement. 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 v4-16, il y a huit puces par et quatre tranches, un total de 32 chips (périphériques) doit être renvoyé d'ici le 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 le de décrire l'API QR. 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 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 des VM efface de force la mémoire de la machine et rétablit l'état initial de la VM. Toutes les données stockées en local restent intactes et le script de démarrage est appelé après une réinitialisation. La L'API ResetQueuedResource peut être utile lorsque vous souhaitez redémarrer tous les TPU. Pour 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 un ResetQueuedResource l'opération 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 de formation, supprimez les ressources en file d'attente ressource avec l'option --force. La suppression prendra deux à cinq minutes terminée 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 de perturbation, Multislice vous propose la réparation de la tranche concernée et la réinitialisation de toutes les tranches par la suite. Les tranche 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 vérifie et charge les derniers points de contrôle enregistrés. Votre script de démarrage est exécuté automatiquement chaque fois qu'une tranche est réallouée ou qu'une VM est réinitialisée. Vous spécifiez une start-up 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 après des défaillances et de reprendre l'entraînement des 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 tester cette méthode.

Profilage et débogage

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

Optimiser la formation

Segmenter avec plusieurs tranches pour des performances optimales

Pour obtenir des performances maximales dans les environnements Multislice, vous devez réfléchir à la répartition des données sur les différentes tranches. Il existe généralement trois (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 en fonction des dimensions du modèle (parfois, appelé parallélisme des Tensors), car il 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 de 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 qu'en cas de nécessité, en fonction de la taille de lot (tel qu'analyse ci-dessous).

Quand utiliser le parallélisme des données ?

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

Pour obtenir un scaling efficace sur plusieurs segments, le temps nécessaire pour réaliser une réduction globale sur le réseau de contenu ne doit pas dépasser le temps nécessaire pour effectuer une rétrogradation. DCN est utilisé pour la communication entre les tranches et est un facteur limitant le débit d'une charge de travail.

Chaque puce TPU v4 offre une puissance de calcul maximale de 275 * 1012 FLOPS par seconde.

Chaque hôte TPU a quatre puces, et chaque hôte a 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 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 DCN. Si vous Utiliser plus de deux tranches, le compilateur effectuera un brassage complet de réductions intégrales et vous utiliserez jusqu'à 64 bits de bande passante DCN pour chaque paramètre et étape. La quantité de FLOPS nécessaire pour chaque paramètre varie en fonction de votre du modèle. Plus précisément, pour les modèles de langage basés sur Transformer, le nombre de FLOPS nécessaires pour une propagation avant et arrière sont d'environ 6 * B * P, où:

  • B est la taille de lot en jetons
  • P est le 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 4 * B.

Pour garantir un scaling puissant sur plusieurs tranches, assurez-vous que les composants opérationnels dépasse l'intensité arithmétique du matériel TPU. Pour calculer le intensité opérationnelle, divisez le nombre de FLOPS par paramètre pendant 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 et 352 Ko pour Transformer des modèles de langage basés sur les données. Comme le réseau DCN peut supprimer brièvement des paquets, Il est préférable de conserver une marge d'erreur importante, en déployant uniquement le parallélisme des données 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èles, vous devez estimer la durée d'exécution par tranche (en chronométrant à l'aide d'un profileur ou en comptant FLOPS). Vous pouvez ensuite comparer cela à la durée d'exécution prévue pour réduire et obtenez 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 (segmente le données entre les nœuds) en segmentant les pondérations entre les nœuds. Pour chaque opération dans avant et après, les pondérations sont rassemblées de sorte que chaque a le poids nécessaire. Au lieu de synchroniser les gradients all-reduce, les gradients sont dispersés en réduction à mesure qu'ils sont produits. De cette façon, chaque tranche n'obtient que les gradients pour les pondérations dont elle est responsable.

Comme pour le parallélisme des données, le FSDP nécessite de mettre à l'échelle la taille de lot globale linéairement avec le nombre de secteurs. Le FSDP réduit la pression sur la mémoire au fur et à mesure vous augmentez le nombre de segments. En effet, le nombre de pondérations l'état de l'optimiseur par tranche diminue, mais cela entraîne une augmentation le trafic réseau et le risque de blocage en raison d'une collective.

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

Les opérations de collecte et de réduction complètes dans le FSDP fonctionnent de manière similaire à celles du DP. Vous pouvez donc déterminer si votre charge de travail FSDP est limitée par les performances du 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 des pipelines devient intéressant pour atteindre des performances élevées avec d'autres des stratégies de parallélisme nécessitant une taille de lot globale supérieure la taille de lot maximale souhaitée. Le parallélisme de pipeline permet aux tranches composant un pipeline de "partager" un lot. Toutefois, le parallélisme de pipeline présente deux inconvénients importants :

  1. Cela entraîne une "bulle de pipeline" où les puces sont inactives parce qu'elles attendent pour les données.
  2. Elle nécessite un microtraitement par lot qui réduit la taille effective du lot, l'intensité arithmétique et, finalement, modélisent l'utilisation du FLOP.

N'utilisez le parallélisme de pipeline que si les autres stratégies de parallélisme nécessitent une taille de lot globale trop élevée. Avant d'essayer le parallélisme des pipelines, vaut la peine d'effectuer des tests pour voir empiriquement si la convergence par échantillon ralentit la taille de lot nécessaire pour obtenir des FSDP hautes performances. Le FSDP a tendance à atteindre d'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 rester le meilleur choix. La plupart peuvent tolérer des tailles de lot suffisamment importantes pour ne pas bénéficier parallélisme du pipeline, mais votre charge de travail peut être différente.

Si le parallélisme des pipelines est nécessaire, nous vous recommandons de le combiner avec des données parallélisme, ou FSDP. Vous pourrez ainsi réduire la profondeur du pipeline tout en augmentant la taille de lot par pipeline jusqu'à ce que la latence du DCN devienne moins importante pour le débit. Concrètement, si vous avez N tranches, envisagez d'utiliser des pipelines profondeur 2 et N/2 instances répliquées du parallélisme des données, puis les pipelines de profondeur 4 et N/4 des instances répliquées du parallélisme des données, etc., jusqu'à ce que le lot par pipeline suffisamment pour que les collectifs DCN puissent être cachés derrière l'arithmétique dans en arrière. Cela réduira le ralentissement induit par le pipeline parallélisme 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 façon répétée à partir d'un ensemble de données pour alimenter du modèle. Avec un chargeur de données asynchrone efficace qui segmente le lot des hôtes afin d'éviter de priver les TPU de travail. Chargeur de données actuel dans MaxText, chaque hôte charge un sous-ensemble égal des exemples. Cette solution est convenable pour le texte, mais nécessite une nouvelle segmentation dans le modèle. En outre, MaxText ne propose pas encore de création d'instantanés déterministe, qui permettrait à l'itérateur de données pour 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 les primitives permettant de créer des points de contrôle des PyTrees JAX vers un espace de stockage local ou Google Cloud. Nous fournissons une intégration de référence avec les 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 plus d'informations, 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 l'implémentation du parallélisme de pipeline avec Multislice, contactez votre responsable de compte Google Cloud.

Assistance et commentaires

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