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.
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).
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.
- 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
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.
Configurez l'environnement:
$ gcloud auth login $ gcloud config set project your-project-id $ gcloud config set compute/zone your-zone
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
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.
Attendez que le code QR soit à l'état
ACTIVE
, ce qui signifie que les nœuds de calcul dont l'état estREADY
. 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
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
Clonez MaxText (qui inclut
shardings.py
) sur votre VM TPU.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
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).
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
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
Clonez MaxText sur votre exécuteur. machine.
Accédez au répertoire du dépôt.
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 fichiergoogle_compute_engine
existe déjà, choisissez de ne pas conserver votre version existante.$ ssh-keygen -f ~/.ssh/google_compute_engine
Ajoutez une variable d'environnement pour définir le nombre de tranches TPU sur
2
.$ export SLICE_COUNT=2
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
) etruntime-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.
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
ouPROVISIONING
pendant plus de 15 minutes.Installez les dépendances.
$ python3 multihost_runner.py \ --TPU_PREFIX=your-qr-id \ --COMMAND="bash setup.sh"
Exécutez
shardings.py
sur chaque nœud de calcul à l'aide demultihost_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 .
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 :
- Utilisez jax.experimental.mesh_utils.create_hybrid_device_mesh au lieu de jax.experimental.mesh_utils.create_device_mesh lorsque vous créez votre réseau maillé.
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 :
- À l'aide du script du lanceur de tests,
multihost_runner.py
- À l'aide du script d'exécuteur de production,
multihost_job.py
- 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 :
- Cela entraîne une "bulle de pipeline" où les puces sont inactives parce qu'elles attendent pour les données.
- 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.