Présentation de la multicouche 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, cela signifie que 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 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 tranches permettent d'améliorer les performances de façon linéaire.

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

Dataflow multicouche

Les développeurs n'ont pas besoin d'écrire de code pour implémenter la communication DCN inter-slice. Le compilateur XLA génère ce code pour vous et superpose la communication au calcul pour des performances maximales.

Concepts

Type d'accélérateur
Forme de chaque tranche de TPU qui compose une multitranche. Chaque tranche d'une requête multitranche 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'un segment rencontre un événement de maintenance, une préemption ou une défaillance matérielle, Cloud TPU crée un nouveau segment. Dans les rares cas où les ressources ne sont pas suffisantes pour créer une tranche, la création ne sera pas terminée tant que le matériel ne sera pas disponible. Une fois le nouveau segment créé, toutes les autres tranches de l'environnement multislice sont redémarrées pour que l'entraînement puisse se poursuivre.Avec un script de démarrage correctement configuré, le script d'entraînement peut être relancé automatiquement sans 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 pour l'entraînement ou l'inférence.
Réseaux de centres de données (DCN)
Réseau à latence plus élevée et à débit inférieur (par rapport à l'ICI) qui connecte les tranches TPU dans une configuration Multislice.
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 ne peut exécuter au maximum que quatre VM à la fois. 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.
Interconnexion entre puces (ICI)
Liens internes à haut débit et à faible latence qui connectent les TPU d'un pod TPU.
Multislice
Deux tranches de puces TPU ou plus pouvant communiquer via le DCN.
Nœud
Dans le contexte Multislice, le nœud fait référence à une seule tranche TPU. Chaque tranche de TPU d'une tranche multiple se voit attribuer un ID de nœud.
Pod
Ensemble de puces TPU connectées par des interfaces réseau ICI dédiées. Un pod vous permet de répartir la charge de traitement sur plusieurs TPU.
Ressource en file d'attente (QR)
Représentation des ressources TPU, utilisée pour mettre en file d'attente et gérer une requête pour un environnement TPU à tranche unique ou multitranches.
Script de démarrage
Script de démarrage Compute Engine standard exécuté à chaque démarrage ou redémarrage d'une VM. Pour Multislice, il est spécifié dans la requête de création du code QR. Pour en savoir plus sur les scripts de démarrage Cloud TPU, consultez Gérer les ressources TPU.
Tranche de TPU
Sous-section logique d'un pod TPU composée de puces TPU. Toutes les puces d'une tranche communiquent entre elles à l'aide du réseau ICI.
VM TPU
Une machine virtuelle exécutant Linux qui a accès aux TPU sous-jacents. Pour les TPU v4, chaque VM TPU a un accès direct à quatre puces. Nous appelons parfois une VM TPU un nœud de calcul.
Tensor
Structure de données utilisée pour représenter des données multidimensionnelles dans un modèle de machine learning.
TPU (Tensor Processing Unit)
Puce d'accélération du ML développée en interne par Google. Elles sont conçues pour offrir un calcul rapide et économe en énergie pour les tâches de machine learning clés, telles que la multiplication de matrices.
Types de capacités 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: pour utiliser une réservation, vous devez disposer d'un contrat de réservation avec Google. Utilisez l'indicateur --reserved lorsque vous créez vos ressources.
  • Spot: cible le quota préemptif à l'aide de VM Spot. Vos ressources peuvent être préemptées pour laisser de la place aux requêtes d'une tâche de priorité plus élevée. 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 sera pas préempté. La demande de 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. Sélectionnée par défaut, aucune balise n'est nécessaire.

Commencer

Si vous n'avez jamais utilisé de TPU, commencez par installer Google Cloud CLI, puis configurez votre environnement Cloud TPU. Pour utiliser Multislice, vos ressources TPU doivent être gérées en tant que ressources en file d'attente.

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

Exemple d'introduction

Ce tutoriel utilise le code du dépôt GitHub MaxText. MaxText est un LLM de base hautes performances, évolutif de manière arbitraire, Open Source et bien testé, écrit en Python et en Jax. MaxText a été conçu pour s'entraîner efficacement sur Cloud TPU.

Le code de shardings.py est conçu pour vous aider à commencer à tester différentes options de parallélisme. Par exemple, le parallélisme de données, le parallélisme de données entièrement segmenté (FSDP) et le parallélisme de tenseur. Le code passe d'une seule tranche à des environnements multicouches.

Parallélisme ICI

L'ICI fait référence à l'interconnexion à haut débit qui connecte les TPU d'une seule tranche. Le fractionnement ICI correspond au fractionnement 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 fragments pour chaque méthode de parallélisme.

Ces entrées doivent être limitées afin que ici_data_parallelism * ici_fsdp_parallelism * ici_tensor_parallelism soit égal au nombre de puces 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 les versions 4 à 8:

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

Notez que ici_data_parallelism doit être défini sur 1 dans la plupart des cas, car le réseau ICI est suffisamment rapide pour préférer presque toujours le FSDP au parallélisme de données.

Cet exemple suppose que vous savez exécuter du code sur une seule tranche TPU, comme dans 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é à remplacer le fichier google_compute_engine existant, faites-le.

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

  3. Provisionnez vos TPU:

    gcloud

    $ 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 du Cloud TPU que vous souhaitez créer. Pour en savoir plus sur les types d'accélérateurs compatibles avec chaque version de TPU, consultez la section Versions de TPU.
    runtime-version
    Version logicielle de Cloud TPU
    node-id
    ID des ressources TPU qui seront créées en réponse à la requête de code QR.
    reserved
    Utilisez une réservation lors de la création des tranches.
    spot
    Utilisez des VM Spot lors de la création des tranches.

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

    Console

    1. Dans la console Google Cloud, accédez à la page TPU:

      Accéder aux TPU

    2. Cliquez sur Créer un TPU.

    3. Dans le champ Nom, saisissez un nom pour votre TPU.

    4. Dans le champ Zone, sélectionnez la zone dans laquelle vous souhaitez créer le TPU.

    5. Dans le champ Type de TPU, sélectionnez un type d'accélérateur. Le type d'accélérateur spécifie la version et la taille du 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.

    6. Dans le champ Version logicielle du TPU, sélectionnez une version logicielle. Lorsque vous créez une VM Cloud TPU, la version logicielle TPU spécifie la version de l'environnement d'exécution TPU à installer. Pour en savoir plus, consultez la section Images de VM TPU.

    7. Cliquez sur le bouton Activer la mise en file d'attente.

    8. Dans le champ Nom de la ressource mise en file d'attente, saisissez un nom pour votre requête de ressource mise en file d'attente.

    9. Cliquez sur Créer pour créer votre requête de ressource mise en file d'attente.

  4. Attendez que la ressource en file d'attente soit à l'état ACTIVE, ce qui signifie que les nœuds de travail sont à l'état READY. Une fois le provisionnement de la ressource mise en file d'attente lancé, il peut prendre une à cinq minutes, en fonction de la taille de la ressource mise en file d'attente. Vous pouvez vérifier l'état d'une demande de ressources mise en file d'attente à l'aide de gcloud CLI ou de la console Google Cloud:

    gcloud

    $ gcloud compute tpus queued-resources \
    list --filter=YOUR_QR_ID

    Console

    1. Dans la console Google Cloud, accédez à la page TPU:

      Accéder aux TPU

    2. Cliquez sur l'onglet Ressources en file d'attente.

    3. Cliquez sur le nom de votre demande de ressources mise en file d'attente.

  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 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 TFLOPS par seconde, soit une utilisation impressionnante de plus de 90%des FLOPS. Dans ce cas, nous avons sélectionné approximativement le lot maximal qui s'adapte à la mémoire à haut débit (HBM) du TPU.

  9. N'hésitez pas à explorer d'autres stratégies de partitionnement sur 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. Supprimez la ressource en file d'attente et la tranche TPU lorsque vous avez terminé. Vous devez exécuter ces étapes de nettoyage à partir de l'environnement dans lequel vous avez configuré la tranche (exécutez d'abord exit pour quitter la session SSH). La suppression prend deux à cinq minutes. Si vous utilisez gcloud CLI, vous pouvez exécuter cette commande en arrière-plan avec l'option --async facultative.

    gcloud

    $ gcloud compute tpus queued-resources
    delete YOUR_QR_ID --force (--async)

    Console

    1. Dans la console Google Cloud, accédez à la page TPU:

      Accéder aux TPU

    2. Cliquez sur l'onglet Ressources en file d'attente.

    3. Cochez la case à côté de votre demande de ressources mise en file d'attente.

    4. Cliquez sur Supprimer.

Segmentation multicouche à 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 fragments 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 afin que dcn_data_parallelism * dcn_fsdp_parallelism * dcn_tensor_parallelism soit égal au nombre de tranches.

Pour deux tranches, utilisez --dcn_data_parallelism = 2.

dcn_data_parallelism dcn_fsdp_parallelism dcn_tensor_parallelism Nombre de tranches
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 partitionnement. Pour les charges de travail LLM typiques sur les puces v4, dcn_fsdp_parallelism doit également être défini sur 1 et dcn_data_parallelism doit donc ê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 maintenez la taille de tranche et le lot par tranche constants), vous augmentez le parallélisme des données.

Exécuter shardings.py dans un environnement Multislice

Vous pouvez exécuter shardings.py dans un environnement multislice à 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 similaires à celles de la section Premiers pas: tests rapides sur plusieurs segments du dépôt MaxText, sauf que 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 à plusieurs reprises les mêmes TPU. Étant donné que le script multihost_runner.py dépend de connexions SSH de longue durée, nous ne le recommandons pas pour les tâches de longue durée. Si vous souhaitez exécuter une tâche plus longue (par exemple, pendant des heures ou des jours), nous vous recommandons d'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 nœuds de calcul pour désigner les VM TPU qui constituent vos tranches. Vous pouvez exécuter multihost_runner.py sur une machine locale ou sur n'importe quelle VM Compute Engine du même projet que vos tranches. 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 travail TPU à l'aide de SSH.

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

Configurer votre environnement

  1. Clonez MaxText sur votre machine de course.

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

  3. 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é à indiquer que le fichier google_compute_engine existe déjà, sélectionnez l'option 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 multicouche à l'aide de la commande queued-resources create ou de Google Cloud Console.

    gcloud

    La commande suivante montre comment créer un TPU multislice v4. Pour utiliser v5e, spécifiez un accelerator-type v5e (par exemple, v5litepod-16) et le 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 du Cloud TPU que vous souhaitez créer. Multislice n'est compatible qu'avec les versions Cloud TPU v4 et ultérieures. Pour en savoir plus sur les types d'accélérateurs compatibles avec chaque version de TPU, consultez la section Versions de TPU.
    runtime-version
    Version logicielle de Cloud TPU
    node-count
    Nombre de tranches à créer.
    node-prefix
    Préfixe utilisé pour générer des noms pour chaque tranche. Un nombre est ajouté au préfixe pour chaque tranche. Par exemple, si vous définissez node-prefix sur mySlice, les tranches sont nommées mySlice-0, mySlice-1, et ainsi de suite de manière numérique pour chaque tranche.
    reserved
    Utilisez une réservation lors de la création des tranches.
    spot
    Utilisez des VM Spot lors de la création des tranches.

    Console

    1. Dans la console Google Cloud, accédez à la page TPU:

      Accéder aux TPU

    2. Cliquez sur Créer un TPU.

    3. Dans le champ Nom, saisissez un nom pour votre TPU.

    4. Dans le champ Zone, sélectionnez la zone dans laquelle vous souhaitez créer le TPU.

    5. Dans le champ Type de TPU, sélectionnez un type d'accélérateur. Le type d'accélérateur spécifie la version et la taille du Cloud TPU que vous souhaitez créer. Multislice n'est compatible qu'avec les versions v4 et ultérieures de Cloud TPU. Pour en savoir plus sur les types d'accélérateurs compatibles avec chaque version de TPU, consultez la section Versions de TPU.

    6. Dans le champ Version logicielle du TPU, sélectionnez une version logicielle. Lorsque vous créez une VM Cloud TPU, la version logicielle TPU spécifie la version de l'environnement d'exécution TPU à installer. Pour en savoir plus, consultez la section Images de VM TPU.

    7. Cliquez sur le bouton Activer la mise en file d'attente.

    8. Dans le champ Nom de la ressource mise en file d'attente, saisissez un nom pour votre requête de ressource mise en file d'attente.

    9. Cochez la case Transformer en TPU multitranches.

    10. Dans le champ Nombre de tranches, saisissez le nombre de tranches que vous souhaitez créer.

    11. Cliquez sur Créer pour créer votre requête de ressource mise en file d'attente.

  6. Lorsque le provisionnement de la ressource mise en file d'attente commence, il peut prendre jusqu'à cinq minutes, en fonction de la taille de la ressource mise en file d'attente. Attendez que la ressource mise en file d'attente soit dans l'état ACTIVE. Vous pouvez vérifier l'état d'une demande de ressources mise en file d'attente à l'aide de la gcloud CLI ou de la console Google Cloud:

    gcloud

    $ 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
...

    Console

    1. Dans la console Google Cloud, accédez à la page TPU:

      Accéder aux TPU

    2. Cliquez sur l'onglet Ressources en file d'attente.

    3. Cliquez sur le nom de votre demande de ressources mise en file d'attente.

    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/s de performances dans les fichiers journaux.

  9. Nettoyez les TPU et la ressource en file d'attente lorsque vous avez terminé. La suppression prend deux à cinq minutes. Si vous utilisez gcloud CLI, vous pouvez exécuter cette commande en arrière-plan avec l'option --async facultative.

Évoluer une charge de travail vers Multislice

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

Il s'agit des seules modifications de code nécessaires lors de la migration vers Multislice. Pour obtenir des performances élevées, le réseau de données de calcul doit être mappé sur des axes de données parallèles, de données parallèles entièrement segmentées ou de pipeline parallèle. Les considérations sur les performances et les stratégies de partitionnement sont abordées plus en détail dans la section Partitionnement avec multislice pour des performances maximales.

Pour vérifier que votre code peut accéder à tous les appareils, vous pouvez affirmer que len(jax.devices()) est égal au nombre de puces de votre environnement Multislice. 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 des tailles de tranche pour les environnements multicouches

Pour obtenir une accélération linéaire, ajoutez des tranches de la même taille que votre tranche existante. 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 la section Diviser en fragments avec Multislice pour des performances maximales.

Exécuter votre tâche 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 de test, multihost_runner.py
  2. À l'aide du script de l'exécuteur de production, multihost_job.py
  3. Utiliser une approche manuelle

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

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

Étant donné que multihost_runner.py maintient des connexions SSH persistantes, il n'est adapté qu'aux tests de petite envergure et relativement courts. Vous pouvez adapter les étapes du tutoriel multihost_runner.py à votre charge de travail et à votre configuration matérielle.

Script de l'exécuteur de production

Pour les tâches de production qui nécessitent une résilience contre les 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, nous fournissons multihost_job.py, qui déclenche l'appel d'API Created Queued Resource avec le script de démarrage approprié pour exécuter votre entraînement et reprendre en cas de préemption. Le script multihost_job.py est documenté dans le fichier README MaxText.

Étant donné que multihost_job.py doit provisionner des ressources pour chaque exécution, il ne fournit pas 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 dans votre configuration Multislice. Toutefois, si vous préférez provisionner et gérer directement votre environnement à l'aide de commandes QR, consultez la section Gérer un environnement multicouche.

Gérer un environnement Multislice

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 ressources en file d'attente

gcloud

  1. 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

    Descriptions des variables

    Entrée Description
    YOUR_QR_ID ID attribué par l'utilisateur à la ressource mise en file d'attente.
    PROJECT Nom du projet Google Cloud
    ZONE Spécifie la zone dans laquelle créer les ressources.
    NETWORK_NAME Nom des réseaux VPC.
    SUBNETWORK_NAME Nom du sous-réseau dans les réseaux VPC
    RUNTIME_VERSION Version logicielle de Cloud TPU
    ACCELERATOR_TYPE v4-16
    EXAMPLE_TAG_1, EXAMPLE_TAG_2 … Tags utilisés pour identifier les sources et les cibles valides pour les pare-feu de réseau
    SLICE_COUNT Nombre de tranches. Limité à 256 tranches au maximum.
    STARTUP_SCRIPT Si vous spécifiez un script de démarrage, il s'exécute lorsque la tranche de TPU est provisionnée ou redémarrée.

  2. Créez une requête de ressource mise en file d'attente à l'aide de la commande suivante:

    $ 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 ressource mise en file d'attente.
    project
    Projet Google Cloud dans lequel vous créez la requête de ressource mise en file d'attente.
    zone
    Zone Google Cloud dans laquelle créer la ressource mise en file d'attente.
    node-count
    Nombre de tranches à créer.
    accelerator-type
    Le type d'accélérateur spécifie la version et la taille du Cloud TPU que vous souhaitez créer. Multislice n'est compatible qu'avec les versions Cloud TPU v4 et ultérieures. Pour en savoir plus sur les types d'accélérateurs compatibles avec chaque version de TPU, consultez la section Versions de TPU.
    runtime-version
    Version logicielle de Cloud TPU
    network
    Nom d'un réseau VPC auquel associer la ressource TPU.
    subnetwork
    Nom d'un sous-réseau VPC auquel associer la ressource TPU.
    reserved
    Utilisez une réservation lors de la création des tranches.
    spot
    Utilisez des 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 par défaut sur demande. Pour en savoir plus sur les types de quotas, consultez la page Règles de quotas.

curl

  1. 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"
    }
    }
    ]
}
}

    Remplacez les valeurs suivantes :

    • your-project-number : numéro de votre projet Google Cloud
    • your-zone : zone dans laquelle vous souhaitez créer votre ressource mise en file d'attente
    • accelerator-type : version et taille d'une seule tranche. Multislice n'est compatible qu'avec les versions v4 et ultérieures de Cloud TPU.
    • tpu-vm-runtime-version : version d'exécution de la VM TPU que vous souhaitez utiliser.
    • your-network-name : facultatif, réseau auquel la ressource mise en file d'attente sera associée
    • your-subnetwork-name : facultatif, sous-réseau auquel la ressource mise en file d'attente sera associée
    • example-tag-1 : facultatif, chaîne de balise arbitraire
    • your-startup-script : script de démarrage qui sera exécuté lorsque la ressource mise en file d'attente sera allouée
    • slice-count : nombre de tranches de TPU dans votre environnement Multislice
    • YOUR_QR_ID : ID fourni par l'utilisateur pour la ressource mise en file d'attente

    Pour en savoir plus, consultez la documentation de 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.

  2. Envoyez la requête de création de ressources mise en file d'attente 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

    Remplacez les valeurs suivantes :

    • your-project-id : ID de votre projet Google Cloud
    • your-zone : zone dans laquelle vous souhaitez créer votre ressource mise en file d'attente
    • YOUR_QR_ID : ID fourni par l'utilisateur pour la ressource mise en file d'attente

La page devrait ressembler à l'exemple ci-dessous :

{
"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 la chaîne pour l'attribut name afin d'obtenir des informations sur la requête de ressource mise en file d'attente.

Console

  1. Dans la console Google Cloud, accédez à la page TPU:

    Accéder aux TPU

  2. Cliquez sur Créer un TPU.

  3. Dans le champ Nom, saisissez un nom pour votre TPU.

  4. Dans le champ Zone, sélectionnez la zone dans laquelle vous souhaitez créer le TPU.

  5. Dans le champ Type de TPU, sélectionnez un type d'accélérateur. Le type d'accélérateur spécifie la version et la taille du Cloud TPU que vous souhaitez créer. Multislice n'est compatible qu'avec les versions v4 et ultérieures de Cloud TPU. Pour en savoir plus sur les types d'accélérateurs compatibles pour chaque version de TPU, consultez la section Versions de TPU.

  6. Dans le champ Version logicielle du TPU, sélectionnez une version logicielle. Lorsque vous créez une VM Cloud TPU, la version logicielle TPU spécifie la version de l'environnement d'exécution TPU à installer. Pour en savoir plus, consultez la section Images de VM TPU.

  7. Cliquez sur le bouton Activer la mise en file d'attente.

  8. Dans le champ Nom de la ressource mise en file d'attente, saisissez un nom pour votre requête de ressource mise en file d'attente.

  9. Cochez la case Transformer en TPU multitranches.

  10. Dans le champ Nombre de tranches, saisissez le nombre de tranches que vous souhaitez créer.

  11. Cliquez sur Créer pour créer votre requête de ressource mise en file d'attente.

Récupérez l'état d'une ressource en file d'attente.

gcloud

$ gcloud compute tpus queued-resources describe ${YOUR_QR_ID}

Pour une ressource mise en file d'attente qui est dans l'état ACTIVE, la sortie se présente comme suit:

...
state:
    state: ACTIVE
...

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}

Pour une ressource mise en file d'attente qui est dans l'état ACTIVE, la sortie se présente comme suit:

{
"name": your-queued-res,
"tpu": {
    "nodeSpec": [
    {
        ... // node 1
    },
    {
        ... // node 2
    },
    ...
    ]
},
...
"state": "ACTIVE"
}

Console

  1. Dans la console Google Cloud, accédez à la page TPU:

    Accéder aux TPU

  2. Cliquez sur l'onglet Ressources en file d'attente.

  3. Cliquez sur le nom de votre demande de ressources mise en file d'attente.

Une fois votre TPU provisionné, vous pouvez également consulter les détails de votre demande de ressources mise en file d'attente en accédant à la page TPU, en recherchant votre TPU, puis en cliquant sur le nom de la demande de ressources mise en file d'attente correspondante.

Dans de rares cas, vous pouvez constater que votre ressource mise en file d'attente est à l'état FAILED, tandis que certaines tranches sont à l'état ACTIVE. Dans ce cas, supprimez les ressources créées, puis réessayez dans quelques minutes ou contactez l'assistance Google Cloud.

SSH et installation des dépendances

Exécuter du code JAX sur des tranches de pods TPU décrit comment vous 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 et travailleurs de QR à l'aide de SSH. La commande est groupée par quatre et envoyée simultanément. Le prochain lot de commandes est envoyé lorsque l'exécution du lot actuel est terminée. 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 la documentation de référence de l'API des 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 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 divisera les tranches en lots. Il exécute le script d'entraînement sur le premier lot de 64 tranches et attend que les scripts soient terminés avant d'exécuter le script sur le lot restant de 36 tranches. Toutefois, le premier lot de 64 tranches ne peut pas être terminé tant que les 36 tranches restantes ne commencent pas à 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 un esperluète (&) à la commande de script que vous spécifiez avec l'indicateur --command. Dans ce cas, 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 restant de 36 tranches. Vous devez canaliser 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 réseau

Pour vous assurer que les tranches TPU peuvent communiquer entre elles, procédez comme suit : Installez JAX sur chacune des tranches. Pour en savoir plus, consultez la section Exécuter du code JAX sur des tranches de pod TPU. Vérifiez que len(jax.devices()) est égal au nombre de puces dans votre environnement Multislice. Pour ce faire, exécutez la commande suivante sur chaque tranche:

  $ 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 tranche et quatre tranches, soit un total de 32 puces (appareils) qui doivent être renvoyées par jax.devices().

Répertorier les ressources en file d'attente

gcloud

Vous pouvez afficher l'état de vos ressources en file d'attente à l'aide de la commande queued-resources list:

$ gcloud compute tpus queued-resources list

La sortie ressemble à ceci :

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

Console

  1. Dans la console Google Cloud, accédez à la page TPU:

    Accéder aux TPU

  2. Cliquez sur l'onglet Ressources en file d'attente.

Démarrer votre tâche 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 des codes QR

L'API ResetQueuedResource permet de réinitialiser toutes les VM d'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 localement restent intactes, et le script de démarrage est 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 simple que le débogage.

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

$ gcloud compute tpus queued-resources reset YOUR_QR_ID

Supprimer des ressources en file d'attente

Pour libérer des ressources à la fin de votre session de formation, supprimez la ressource mise en file d'attente. La suppression prend deux à cinq minutes. Si vous utilisez gcloud CLI, vous pouvez exécuter cette commande en arrière-plan avec l'option --async facultative.

gcloud

$ gcloud compute tpus queued-resources \
    delete YOUR_QR_ID --force (--async)

Console

  1. Dans la console Google Cloud, accédez à la page TPU:

    Accéder aux TPU

  2. Cliquez sur l'onglet Ressources en file d'attente.

  3. Cochez la case à côté de votre demande de ressources mise en file d'attente.

  4. Cliquez sur Supprimer.

Reprise après échec automatique

En cas de perturbation, Multislice propose une réparation sans intervention de la tranche concernée et une réinitialisation de toutes les tranches par la suite. La tranche concernée est remplacée par une nouvelle tranche, et les autres tranches saines sont réinitialisées. Si aucune capacité n'est disponible pour allouer un segment 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 recherche 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 un script de démarrage dans la charge utile JSON que vous envoyez à l'API de création de requêtes QR.

Le script de démarrage suivant (utilisé dans 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 méthode.

Profilage et débogage

Le profilage est le même dans les environnements à tranche unique et multicouches. Pour en savoir plus, consultez la page Profiler des programmes JAX.

Optimiser la formation

Partitionnement avec Multislice 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 choix (parallélisme des données, parallélisme des données entièrement segmenté et parallélisme de pipeline). Nous vous déconseillons de fractionner les activations sur les dimensions du modèle (parfois appelé parallélisme tensoriel), car cela nécessite trop de bande passante inter-tranches. Pour toutes ces stratégies, vous pouvez conserver la même stratégie de partitionnement dans une tranche qui a fonctionné pour vous par le passé.

Nous vous recommandons de commencer par le parallélisme de données pur. 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 de pipeline que si nécessaire en fonction de la taille de lot (comme analysé ci-dessous).

Quand utiliser le parallélisme de données ?

Le parallélisme de données pur fonctionne bien lorsque vous disposez d'une charge de travail qui fonctionne correctement, mais que vous souhaitez améliorer ses performances en la faisant évoluer sur plusieurs tranches.

Pour obtenir une forte évolutivité sur plusieurs tranches, le temps nécessaire pour effectuer une réduction totale sur le DCN doit être inférieur au temps nécessaire pour effectuer une passe arrière. Le DCN est utilisé pour la communication entre les tranches et constitue un facteur limitant du débit de la charge de travail.

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

Il existe 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 chaque é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 mélange all-reduce complète et vous utiliserez jusqu'à 64 bits de bande passante DCN pour chaque paramètre par étape. La quantité de FLOPS requise 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 passe avant et une passe arrière est d'environ 6 * B * P, où:

  • B correspond à la taille du lot en jetons.
  • P est le nombre de paramètres.

Le nombre de FLOPS par paramètre est 6 * B et le nombre de FLOPS par paramètre lors de la passe arrière est 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é opérationnelle, divisez le nombre de FLOPS par paramètre lors de la rétropropagation par la bande passante du réseau (en bits) par paramètre 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 tranches: Operational intensity = 4 * B/64

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

Pour d'autres architectures de modèle, vous devrez estimer le temps d'exécution de votre pass en arrière par tranche (en le chronométrant à l'aide d'un profileur ou en comptant les FLOPS). Vous pouvez ensuite comparer cela au temps d'exécution prévu pour réduire tout le DCN et obtenir une bonne estimation de la pertinence du parallélisme de données pour vous.

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

Le parallélisme des données entièrement partitionné (FSDP) combine le parallélisme des données (partitionnement des données entre les nœuds) et le partitionnement des poids entre les nœuds. Pour chaque opération des passes avant et arrière, les poids sont tous rassemblés afin que chaque tranche dispose des poids dont elle a besoin. Au lieu de synchroniser les gradients à l'aide de la réduction globale, les gradients sont dispersés à mesure qu'ils sont produits. De cette manière, chaque tranche ne reçoit que les gradients des poids dont elle est responsable.

Comme pour le parallélisme de données, le parallélisme de données distribuées nécessite d'adapter la taille de lot globale de manière linéaire au nombre de tranches. La FSDP réduit la pression sur la mémoire à mesure que vous augmentez le nombre de tranches. En effet, le nombre de poids et l'état de l'optimiseur par tranche diminue, mais cela se fait 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 les tranches est préférable si vous augmentez le lot par tranche, stockez plus d'activations pour minimiser la rematérialisation lors de la passe arrière ou augmentez 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 de pipeline devient pertinent lorsque vous obtenez 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 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. Il génère la "bulle de pipeline", où les puces sont inactives, car elles attendent des données.
  2. Il nécessite un microtraitement qui réduit la taille de lot effective, l'intensité arithmétique et, en fin de compte, l'utilisation des FLOP du modèle.

Le parallélisme de pipeline ne doit être utilisé que si les autres stratégies de parallélisme nécessitent une taille de lot globale trop importante. Avant d'essayer le parallélisme de pipeline, il est utile de tester de manière empirique si la convergence par échantillon ralentit à la taille de lot nécessaire pour obtenir un FSDP hautes performances. Le FSDP tend à obtenir une utilisation plus élevée des FLOP du modèle, mais si la convergence par échantillon ralentit à mesure que la taille de lot augmente, le parallélisme de 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 de pipeline, mais votre charge de travail peut être différente.

Si le parallélisme de pipeline est nécessaire, nous vous recommandons de le combiner au parallélisme de données ou au 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 des pipelines de profondeur 2 et des réplicas de parallélisme de données N/2, puis des pipelines de profondeur 4 et des réplicas de parallélisme de données N/4, et ainsi de suite, jusqu'à ce que le lot par pipeline soit suffisamment important pour que les collectifs DCN puissent être masqués derrière l'arithmétique dans la passe inverse. Cela permet de minimiser le ralentissement introduit par le parallélisme de pipeline, tout en vous permettant de dépasser la limite de taille de lot globale.

Bonnes pratiques concernant les multicouches

Chargement des données…

Lors de l'entraînement, nous chargeons de manière répétée des lots à partir d'un ensemble de données pour les transmettre au modèle. Il est important de disposer d'un chargeur de données asynchrone et efficace qui divise le lot entre les hôtes pour éviter de priver les TPU de travail. Le chargeur de données actuel de MaxText permet à chaque hôte de charger un sous-ensemble égal des exemples. Cette solution est adaptée au texte, mais nécessite un rééclatement dans le modèle. De plus, MaxText n'offre pas encore de création d'instantanés déterministe, ce 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 point de contrôle Orbax fournit des primitives pour le point de contrôle des PyTrees JAX dans un stockage local ou Google Cloud. Nous fournissons une intégration de référence avec le point de contrôle synchrone dans MaxText dans checkpointing.py.

Configurations compatibles

Formes

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

Orchestration

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

Frameworks

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

Parallélisme

Nous recommandons aux utilisateurs de tester Multislice avec le parallélisme de 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

N'hésitez pas à nous faire part de vos commentaires. Pour nous faire part de vos commentaires ou demander de l'aide, contactez-nous à l'aide du formulaire d'assistance ou de commentaires Cloud TPU.