Présentation de Cloud TPU Multislice

Cloud TPU multislice est une technologie de scaling des performances full stack qui permet à un job d'entraînement d'utiliser plusieurs tranches TPU dans une seule tranche ou sur des tranches dans plusieurs pods avec un parallélisme des données standard. Avec les puces TPU v4, les jobs d'entraînement peuvent utiliser plus de 4 096 puces en une seule exécution. Pour les tâches d'entraînement qui nécessitent moins de 4 096 puces, une seule tranche peut offrir les meilleures performances. Toutefois, plusieurs petites tranches sont plus facilement disponibles, ce qui permet un démarrage plus rapide lorsque Multislice est utilisé avec des tranches plus petites.

Le fait d'avoir plusieurs tranches permet de faire évoluer les performances de manière linéaire.

Lorsqu'elles sont déployées dans des configurations Multislice, les puces TPU de chaque tranche communiquent via une interconnexion entre puces (ICI). Les puces TPU de différentes tranches communiquent en transférant des données aux CPU (hôtes), qui à leur tour transmettent les données sur le réseau du centre de données (DCN). Pour en savoir plus sur le scaling avec Multislice, consultez Faire évoluer l'entraînement IA jusqu'à plusieurs dizaines de milliers de puces Cloud TPU grâce à Multislice.

Flux de données Multislice

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

Concepts

Type d'accélérateur

Forme de chaque tranche de TPU qui compose un environnement Multislice. 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 version ultérieure) suivi du nombre de TensorCores. Par exemple, v5litepod-128 spécifie un TPU v5e 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. Si les ressources sont insuffisantes pour créer une tranche, la création ne se terminera pas tant que le matériel ne sera pas disponible. Une fois la nouvelle tranche créée, toutes les autres tranches de l'environnement Multislice seront 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 se relancer automatiquement sans intervention de l'utilisateur, en chargeant et en reprenant l'entraînement à partir du dernier point de contrôle.

Réseau de centre de données (DCN)

Réseau à latence plus élevée et à débit plus faible (par rapport à l'ICI) qui connecte les tranches de TPU dans une configuration Multislice.

Planification de groupe

Lorsque toutes les tranches de TPU sont provisionnées ensemble, en même temps, cela garantit que toutes les tranches sont provisionnées avec succès ou qu'aucune ne l'est.

Interconnexion entre puces (ICI)

Liens internes à haut débit et à faible latence qui connectent les TPU au sein d'un pod TPU.

Multitranches

Au moins deux tranches de puces TPU pouvant communiquer sur le DCN.

Nœud

Dans le contexte Multislice, le terme "nœud" fait référence à une seule tranche de TPU. Chaque tranche de TPU d'un environnement Multislice reçoit un ID de nœud.

Script de démarrage

Un 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 un environnement Multislice, il est spécifié dans la demande de création du code QR. Pour en savoir plus sur les scripts de démarrage Cloud TPU, consultez Gérer les ressources TPU.

Tensor

Structure de données utilisée pour représenter des données multidimensionnelles dans un modèle de machine learning.

Types de capacité Cloud TPU

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

Réservation : pour utiliser une réservation, vous devez avoir conclu un accord de réservation avec Google. Utilisez le flag --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 faire de la place aux requêtes d'un job de priorité plus élevée. Utilisez le flag --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é par défaut, aucun flag n'est nécessaire.

Commencer

Configurez votre environnement Cloud TPU.
In the Google Cloud console, activate Cloud Shell.

Activate Cloud Shell

At the bottom of the Google Cloud console, a Cloud Shell session starts and displays a command-line prompt. Cloud Shell is a shell environment with the Google Cloud CLI already installed and with values already set for your current project. It can take a few seconds for the session to initialize.

Pour utiliser Multislice, vos ressources TPU doivent être gérées en tant que ressources mises en file d'attente.

Exemple d'introduction

Ce tutoriel utilise le code du dépôt GitHub MaxText. MaxText est un LLM de base hautes performances, évolutif à volonté, Open Source et bien testé, écrit en Python et 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 à tester différentes options de parallélisme. Par exemple, le parallélisme des données, le parallélisme des données entièrement segmentées (FSDP) et le parallélisme des tenseurs. Le code s'adapte aux environnements à une ou plusieurs tranches.

Parallélisme ICI

ICI fait référence à l'interconnexion à haut débit qui connecte les TPU d'une même 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 contraintes de sorte 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 de saisies utilisateur pour le parallélisme ICI pour les quatre puces disponibles dans la version 4-8 :

	ici_data_parallelism	ici_fsdp_parallelism	ici_tensor_parallelism
FSDP à quatre voies	1	4	1
Parallélisme de tenseur à quatre voies	1	1	4
FSDP à deux voies + 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 que FSDP soit presque toujours préféré au parallélisme des données.

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

Configurer l'environnement :

$ gcloud auth login
$ export QR_ID=your-queued-resource-id
$ export TPU_NAME=your-tpu-name
$ export PROJECT=your-project-name
$ export ZONE=us-central1-a
$ export NETWORK_NAME=your-network-name
$ export SUBNETWORK_NAME=your-subnetwork-name
$ export RUNTIME_VERSION=v2-alpha-tpuv5-lite
$ export ACCELERATOR_TYPE=v5litepod-16
$ export EXAMPLE_TAG_1=your-tag-1
$ export EXAMPLE_TAG_2=your-tag-2
$ export SLICE_COUNT=4
$ export STARTUP_SCRIPT='#!/bin/bash\n'

Descriptions des variables

Entrée	Description
`QR_ID`	ID attribué par l'utilisateur à la ressource mise en file d'attente.
`TPU_NAME`	Nom attribué par l'utilisateur à votre TPU.
PROJET	Nom du projetGoogle 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 ou cibles valides pour les pare-feu de réseau
SLICE_COUNT	Nombre de tranches. Limité à 256 tranches 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.

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
```
Provisionnez vos TPU :
gcloud
```
$ gcloud compute tpus queued-resources \
    create ${QR_ID} \
    --accelerator-type=${ACCELERATOR_TYPE} \
    --runtime-version=${RUNTIME_VERSION} \
    --node-id=${TPU_NAME} \
    --zone=${ZONE} \
    [--reserved |--spot]
```
Google Cloud CLI ne prend pas en charge toutes les options de création de codes QR, comme les tags. Pour en savoir plus, consultez Créer des QR codes.
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 la zone 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 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 du TPU spécifie la version de l'environnement d'exécution TPU à installer. Pour en savoir plus, consultez Versions logicielles de 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 demande de ressource mise en file d'attente.
9. Cliquez sur Créer pour créer votre demande de ressource mise en file d'attente.
Attendez que la ressource mise en file d'attente soit à l'état ACTIVE, ce qui signifie que les nœuds de calcul sont à l'état READY. Une fois le provisionnement des ressources en file d'attente démarré, il peut prendre entre une et cinq minutes, selon la taille de la ressource en file d'attente. 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=${QR_ID} --zone=${ZONE}
```
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 ressource mise en file d'attente.

Connectez-vous à la VM TPU à l'aide de SSH :

$ gcloud compute tpus tpu-vm ssh ${TPU_NAME} --zone=${ZONE}

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

$ git clone https://github.com/AI-Hypercomputer/maxtext && cd maxtext

Installez Python 3.10 :

$ sudo apt-get update
$ sudo apt install python3.10
$ sudo apt install python3.10-venv

Créez et activez un environnement virtuel :

Remarque : Créer un environnement virtuel pour chaque projet est une bonne pratique pour gérer les dépendances Python. Pour en savoir plus, reportez-vous à la documentation venv.
```
$ python3 -m venv your-venv-name
$ source your-venv-name/bin/activate
```
Dans le répertoire du dépôt MaxText, exécutez le script d'installation pour installer JAX et d'autres dépendances sur votre tranche de TPU. L'exécution du script de configuration prend quelques minutes.
```
$ bash setup.sh
```
Exécutez la commande suivante pour exécuter shardings.py sur votre tranche de TPU.
```
$ python3 -m pedagogical_examples.shardings \
  --ici_fsdp_parallelism 4 \
  --batch_size 131072 \
  --embedding_dimension 2048
```
Vous pouvez consulter les résultats dans les journaux. Vos TPU devraient atteindre environ 260 TFLOP par seconde, soit une utilisation des FLOPS de plus de 90 % ! Dans ce cas, nous avons sélectionné approximativement la taille de lot maximale qui tient dans la mémoire à haut débit (HBM) du TPU.

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

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

Une fois l'opération terminée, supprimez la ressource en file d'attente et la tranche de TPU. Vous devez exécuter ces étapes de nettoyage à partir de l'environnement dans lequel vous avez configuré la tranche (exécutez d'abord exit pour quitter la session SSH). La suppression prend entre deux et cinq minutes. Si vous utilisez la gcloud CLI, vous pouvez exécuter cette commande en arrière-plan avec l'option facultative --async.

Important : Les ressources mises en file d'attente doivent être supprimées après utilisation. Sinon, elles continueront de consommer votre quota de ressources en file d'attente, qui est suivi séparément du quota de TPU.

gcloud

Présentation de Cloud TPU Multislice

Concepts

Commencer

Exemple d'introduction

Parallélisme ICI

Descriptions des variables

gcloud

Console

gcloud

Console

gcloud