Entraînement Cloud TPU v5p
Cloud TPU v5p est la cinquième génération de Cloud TPU de Google Cloud et le successeur du TPU v4. v5p est optimisé pour l'entraînement à grande échelle et constitue une plate-forme de pointe pour le développement de LLM de base, de modèles de diffusion et d'IA générative. À un niveau élevé, la v5p offre des performances jusqu'à deux fois supérieures à celles de la v4, tout en empaquetant deux fois plus de TPU dans un pod (6 000 tranche la plus grande contre 3 000 dans la v4), ce qui se traduit par des performances jusqu'à quatre fois supérieures au niveau du pod. Il fonctionne également à une fréquence d'horloge plus élevée (1,75 GHz contre 1,05 GHz), ajoute SparseCore pour les représentations vectorielles continues à grande échelle et triple la capacité de la mémoire à haut débit (HBM).
Concepts Cloud TPU v5p
Si vous ne connaissez pas les Cloud TPU, consultez la page d'accueil de la documentation TPU.
Les concepts Cloud TPU (par exemple, les tranches, les hôtes et TensorCores) et l'architecture du système Cloud TPU pour toutes les versions de Cloud TPU sont décrits sur la page Architecture du système Cloud TPU.
Chaque version de Cloud TPU nécessite des types d'accélérateurs spécifiques pour l'entraînement ou l'inférence. Ces types d'accélérateurs sont décrits dans les configurations v5p.
Gérer les ressources TPU
Toutes les commandes que vous pouvez utiliser pour gérer vos VM TPU sont décrites dans la section Gérer les TPU ou le guide de l'utilisateur des ressources en file d'attente pour gérer les ressources en file d'attente.
Configuration du framework
Cette section décrit le processus général de configuration de l'entraînement de modèle à l'aide de JAX ou PyTorch avec TPU v5p.
Configuration pour JAX
Si vous avez des formes de tranches supérieures à 4 chips, vous disposez de plusieurs VM dans une même tranche. Dans ce cas, vous devez utiliser l'option --worker=all
pour exécuter l'installation sur toutes les VM TPU à l'aide d'une seule commande:
gcloud compute tpus tpu-vm ssh ${TPU_NAME} --project ${PROJECT_ID} \ --zone ${ZONE} \ --worker=all \ --command='pip install "jax[tpu]==0.4.20" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
Vous pouvez exécuter la commande suivante pour vérifier le nombre d'appareils (les résultats présentés ici ont été produits avec une tranche v5p-32). Ce code vérifie que tout est correctement installé en vérifiant que JAX voit les TensorCores de Cloud TPU et peut exécuter des opérations de base:
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project ${PROJECT_ID} \ --zone ${ZONE} \ --worker=all \ --command='python3 -c "import jax; print(jax.device_count()); print(jax.local_device_count())"'
Le résultat ressemble à ce qui suit :
SSH: Attempting to connect to worker 0... SSH: Attempting to connect to worker 1... SSH: Attempting to connect to worker 2... SSH: Attempting to connect to worker 3... 16 4 16 4 16 4 16 4
jax.device_count()
indique le nombre total de chips dans la tranche donnée. jax.local_device_count()
indique le nombre de chips accessibles par une seule VM dans cette tranche.
# Check the number of chips in the given slice by summing the count of chips # from all VMs through the # jax.local_device_count() API call. gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project ${PROJECT_ID} \ --zone ${ZONE} \ --worker=all \ --command='python3 -c "import jax; xs=jax.numpy.ones(jax.local_device_count()); print(jax.pmap(lambda x: jax.lax.psum(x, \"i\"), axis_name=\"i\")(xs))"'
Le résultat ressemble à ce qui suit :
SSH: Attempting to connect to worker 0... SSH: Attempting to connect to worker 1... SSH: Attempting to connect to worker 2... SSH: Attempting to connect to worker 3... [16. 16. 16. 16.] [16. 16. 16. 16.] [16. 16. 16. 16.] [16. 16. 16. 16.]
Utilisez --node=all
pour exécuter la commande sur tous les nœuds de calcul à plusieurs tranches.
gcloud compute tpus queued-resources ssh ${QUEUED_RESOURCE_ID} \ --project ${PROJECT_ID} --zone ${ZONE} --node=all --worker=all \ --command='python3 -c "import jax; print(jax.device_count()); print(jax.local_device_count())"'
Consultez les tutoriels JAX de ce document pour faire vos premiers pas avec l'entraînement v5p à l'aide de JAX.
Configuration pour PyTorch
L'environnement d'exécution PJRT est le seul compatible avec la version v5p. PyTorch 2.1 et versions ultérieures utilisent PJRT comme environnement d'exécution par défaut pour toutes les versions de TPU. Cette section explique comment commencer à utiliser PJRT sur les pods v5p avec PyTorch/XLA 2.2.0 pour tous les nœuds de calcul.
Installer des dépendances
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project ${PROJECT_ID} \ --zone ${ZONE} \ --worker=all \ --command=' sudo apt-get update sudo apt-get install libopenblas-dev -y pip3 install numpy pip install torch~=2.2.0 torch_xla[tpu]~=2.2.0 -f https://storage.googleapis.com/libtpu-releases/index.html '
Utilisez un script Python avec PJRT pour valider l'installation et afficher les appareils TPU disponibles (les résultats présentés ici ont été produits avec une tranche v5p-32).
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project ${PROJECT_ID} --zone ${ZONE} --worker=all \ --command=' PJRT_DEVICE=TPU python3 -c "import torch_xla.core.xla_model as xm; print(xm.get_xla_supported_devices(\"TPU\"))" '
SSH: Attempting to connect to worker 0... SSH: Attempting to connect to worker 1... SSH: Attempting to connect to worker 2... SSH: Attempting to connect to worker 3... ['xla:0', 'xla:1', 'xla:2', 'xla:3'] ['xla:0', 'xla:1', 'xla:2', 'xla:3'] ['xla:0', 'xla:1', 'xla:2', 'xla:3'] ['xla:0', 'xla:1', 'xla:2', 'xla:3']
Utilisez --node=all
pour exécuter la commande sur tous les nœuds de calcul à plusieurs tranches.
gcloud compute tpus queued-resources ssh ${QUEUED_RESOURCE_ID} \ --project ${PROJECT_ID} --zone ${ZONE} --node=all --worker=all \ --command=' PJRT_DEVICE=TPU python3 -c "import torch_xla.core.xla_model as xm; print(xm.get_xla_supported_devices(\"TPU\"))" '
Consultez les tutoriels PyTorch de ce document pour faire vos premiers pas avec l'entraînement v5p à l'aide de PyTorch.
Surveiller et profiler
Cloud TPU v5p prend en charge la surveillance et le profilage à l'aide des mêmes méthodes que les générations précédentes de Cloud TPU. Consultez la page Profiler votre modèle avec les outils Cloud TPU pour en savoir plus sur le profilage, et la section Surveiller des VM Cloud TPU pour en savoir plus sur la surveillance.
Tutoriels de formation
Cette section se concentre sur les tutoriels d'entraînement pour une seule tranche.
Vous pouvez adapter ces tutoriels à l'entraînement à plusieurs tranches en ajoutant l'option --node=all
aux commandes SSH.
Pour en savoir plus et connaître les bonnes pratiques, consultez la page Introduction aux multitranches.
Tutoriels JAX
Train Diffusion 2.1
Ce tutoriel explique comment entraîner le modèle Stable Diffusion à partir de HuggingFace à l'aide de l'ensemble de données Pokémon sur Cloud TPU v5p.
Le modèle Stable Diffusion est un modèle texte-image latent qui génère des images photoréalistes à partir de n'importe quelle entrée de texte. Pour en savoir plus, consultez les ressources suivantes :
Préparation
Créez des variables d'environnement :
export PROJECT_ID=your_project_ID export ACCELERATOR_TYPE=v5p-32 export ZONE=us-east5-a export RUNTIME_VERSION=v2-alpha-tpuv5 export SERVICE_ACCOUNT=your_service_account export TPU_NAME=your_tpu_name export QUEUED_RESOURCE_ID=queued_resource_id export QUOTA_TYPE=quota_type export VALID_UNTIL_DURATION=1d
Description des options de commande
Variable Description PROJECT_ID Nom du projet Google Cloud ACCELERATOR_TYPE Consultez la page Versions de TPU pour connaître votre version de TPU. ZONE Consultez le document Régions et zones TPU pour connaître les zones disponibles. RUNTIME_VERSION Pour la version v5p, utilisez v2-alpha-tpuv5 pour RUNTIME_VERSION. SERVICE_ACCOUNT Il s'agit de l'adresse de votre compte de service, que vous pouvez trouver dans la console Google Cloud -> IAM -> Comptes de service. Exemple: tpu-service-account@myprojectID.iam.gserviceaccount.com TPU_NAME ID textuel du TPU attribué par l'utilisateur, qui est créé lors de l'allocation de la demande de ressource en file d'attente. QUEUED_RESOURCE_ID ID du texte attribué par l'utilisateur dans la requête de ressource en file d'attente. Consultez le document Ressources en file d'attente pour plus d'informations sur ces ressources. QUOTA_TYPE Il peut s'agir de reserved
oubest-effort
. Si aucune de ces options n'est spécifiée, la valeur par défaut du QUOTA_TYPE eston-demand
. Consultez la section quotas pour en savoir plus sur les différents types de quotas compatibles avec Cloud TPU.VALID_UNTIL_DURATION Durée de validité de la requête. Consultez la section Ressources en file d'attente pour en savoir plus sur les différentes durées valides. -
gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \ --node-id ${TPU_NAME} \ --project ${PROJECT_ID} \ --zone ${ZONE} \ --accelerator-type ${ACCELERATOR_TYPE} \ --runtime-version ${RUNTIME_VERSION} \ --valid-until-duration ${VALID_UNTIL_DURATION} \ --service-account ${SERVICE_ACCOUNT} \ --${QUOTA_TYPE}
Vous pourrez vous connecter en SSH à votre VM TPU une fois que votre ressource en file d'attente passe à l'état
ACTIVE
. Vérifiez l'état de votre ressource en file d'attente en exécutant la commande suivante:gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \ --project ${PROJECT_ID} --zone ${ZONE}
Lorsque la ressource en file d'attente est à l'état
ACTIVE
, le résultat ressemble à ce qui suit:state: ACTIVE
Installez JAX et ses dépendances.
# compatible with v5p: only jax version 0.4.19 and later \ # jax 0.4.19 requires py 3.10 \ gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} --zone=${ZONE} --worker=all \ --command='pip install "jax[tpu]==0.4.20" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
Téléchargez le dépôt HuggingFace et la configuration requise pour l'installation.
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --worker=all \ --command='git clone https://github.com/huggingface/diffusers.git && cd diffusers && pip install . && pip install tensorflow clu && pip install -U -r examples/text_to_image/requirements_flax.txt'
Entraîner le modèle
Entraînez le modèle avec un tampon prémappé de 4 Go.
gcloud compute tpus tpu-vm ssh ${TPU_NAME} --project=${PROJECT_ID} \ --zone=${ZONE} \ --worker=all \ --command='export PATH=$PATH:$HOME/.local/bin && cd diffusers/examples/text_to_image && JAX_PLATFORMS=tpu,cpu python3 train_text_to_image_flax.py --pretrained_model_name_or_path=stabilityai/stable-diffusion-2-1 --dataset_name=lambdalabs/pokemon-blip-captions --resolution=256 --center_crop --random_flip --train_batch_size=1 --mixed_precision=bf16 --max_train_steps=150 --learning_rate=1e-05 --max_grad_norm=1 --output_dir=sd-pokemon-model --from_pt'
Effectuer un nettoyage
Supprimez votre TPU et la demande de ressources en file d'attente à la fin de votre session ou pour supprimer les demandes de ressources en file d'attente à l'état "ÉCHEC". Pour supprimer une ressource en file d'attente, supprimez les tranches, puis la requête de ressource en file d'attente en deux étapes:
gcloud compute tpus tpu-vm delete ${TPU_NAME} --project=${PROJECT_ID} --zone=${ZONE} --quiet
gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} --project ${PROJECT_ID} --zone ${ZONE} --quiet
Vous pouvez également utiliser --force
pour supprimer la ou les tranches et la requête de ressource en file d'attente en une seule étape:
# With --force gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} --project ${PROJECT_ID} --zone ${ZONE} --quiet --force
Résultats du benchmark
Le script d'entraînement de Stable Diffusion a été exécuté sur les versions v5p-8, v5p-32 et v5p-128. Le tableau suivant indique le débit.
V5p-8 |
V5p-32 |
V5P-128 |
|
---|---|---|---|
Étape d'entraînement |
150 |
150 |
150 |
Taille du lot global |
32 |
64 |
64 |
Débit (exemples/s) |
12.10 |
18,08 |
19,10 |
MaxText
Ce tutoriel explique comment entraîner le modèle MaxText à l'aide d'un ensemble de données synthétique sur Cloud TPU.
MaxText est un LLM Open Source hautement performant, arbitrairement évolutif et éprouvé, écrit en Python/JAX pur et ciblant les Cloud TPU. MaxText offre aux chercheurs et aux développeurs un outil accessible et adaptable pour repousser les frontières de la recherche et du développement du traitement du langage naturel (TLN).
Avant d'exécuter ce tutoriel, vous devez configurer votre environnement Cloud TPU.
Configurer des variables d'environnement
export PROJECT_ID=your_project_ID export TPU_NAME=your_tpu_name # user defined TPU name export ACCELERATOR_TYPE=v5p-256 export ZONE=us-east5-a export RUNTIME_VERSION=v2-alpha-tpuv5 export RUN_NAME=your_experiment_run_name # user defined name for this run export GCS_BUCKET_NAME=your_bucket_name # Output cloud folder. Should start with gs:// export MAXTEXT_OUTPUT_PATH=${GCS_BUCKET_NAME}/your_experiment_output_path export NUM_SLICES=1 # Update the value to a number >1 for Multislice.
Description des options de commande
Variable Description PROJECT_ID Nom du projet Google Cloud TPU_NAME Nom défini par l'utilisateur pour votre TPU. ACCELERATOR_TYPE Consultez la page Versions de TPU pour connaître votre version de TPU. ZONE Consultez le document Régions et zones TPU pour connaître les zones disponibles. RUNTIME_VERSION Pour v5p, utilisez v2-alpha-tpuv5 pour la version d'exécution. RUN_NAME Nom de l'exécution du test fourni par l'utilisateur. Configuration facultative recommandée pour les tranches multitranches:
export NETWORK_NAME=your_network_name export FIREWALL_RULE_NAME=your_firewall_rule_name
Si vous exécutez des charges de travail multitranches et souhaitez obtenir des performances réseau optimales, envisagez de créer un réseau dédié avec une unité de transmission maximale (MTU) de 8 896 octets et de configurer des règles de pare-feu appropriées. Bien que facultative, cette étape peut considérablement améliorer les performances, en particulier lors du scaling à la hausse du nombre de tranches sur le réseau du centre de données (DCN). Notez que la création d'un réseau nécessite l'autorisation
compute.networks.create
dans le projet. Les exemples suivants montrent comment créer une règle de réseau et de pare-feu dédiée.Créez un réseau dédié:
gcloud compute networks create ${NETWORK_NAME} \ --mtu=8896 \ --project=${PROJECT_ID} \ --subnet-mode=auto \ --bgp-routing-mode=regional
Créez une règle de pare-feu :
gcloud compute firewall-rules create ${FIREWALL_RULE_NAME} \ --network ${NETWORK_NAME} --allow tcp,icmp,udp --project=${PROJECT_ID}
Cloner le dépôt MaxText
git clone https://github.com/google/maxtext.git
Entraîner le modèle
Les sections suivantes décrivent deux options pour l'entraînement de MaxText.
Option 1
Si vous souhaitez qu'un script gère l'ensemble du workflow, du provisionnement des Cloud TPU et de l'installation des dépendances à l'exécution de votre modèle et à la suppression des ressources, vous pouvez utiliser
multihost_job.py
.cd maxtext && python3 multihost_job.py --PROJECT=${PROJECT_ID} --ZONE=${ZONE} \ --NUM_SLICES=${NUM_SLICES} --TPU_TYPE=${ACCELERATOR_TYPE} \ --VERSION=${RUNTIME_VERSION} --RUN_NAME=${RUN_NAME} #user defined run name \ --BUCKET_NAME=${GCS_BUCKET_NAME} \ #used to store logs and configs --COMMAND="bash setup.sh && bash MaxText/configs/experimental/64b.sh RUN_NAME=${RUN_NAME} OUTPUT_PATH=${MAXTEXT_OUTPUT_PATH} PLATFORM=gce"
Une fois le script lancé, un message semblable au suivant doit s'afficher dans le journal. L'emplacement du journal est référencé dans le message de sortie. Cliquez sur le premier lien pour accéder aux journaux de tous les nœuds de calcul une fois le provisionnement du TPU terminé.
------------------------------------ multihost_job finished running, TPUs are starting up to run your job remotely. Logs for your job are displayed here: https://console.cloud.google.com/logs/query;query=resource.type%3D%22gce_instance%22%20AND%0Alog_id%2528%22
_log%22%2529;?project=PROJECT_ID To see the output of a single host, you may edit the slice and worker number in the `log_file_path` property here: https://console.cloud.google.com/logs/query;query=resource.type%3D%22gce_instance%22%20AND%0Alog_id%2528%22RUN_NAME_log%22%2529%20AND%0Alabels.%22agent.googleapis.com%2Flog_file_path%22%3D%20%22%2FRUN_NAME%2Fmain_command_log_slice_0_worker_0%22;?project=PROJECT_ID When your job is finished, the main command log is in your Cloud Storage bucket: https://console.cloud.google.com/storage/browser/YOUR_BUCKET_NAME/RUN_NAME?project=PROJECT_ID View the status of the created TPUs using: gcloud compute tpus queued-resources list --filter=RUN_NAME --zone=ZONE --project=PROJECT_ID
Option 2
Pour exécuter le script d'entraînement plusieurs fois sur un Cloud TPU provisionné, utilisez le script multihost_runner.py
pour utiliser la ressource.
Configurez des variables pour créer un TPU.
export SERVICE_ACCOUNT=your_service_account export TPU_NAME=your_tpu_name export QUEUED_RESOURCE_ID=your_queued_resource_id export VALID_DURATION=1d export QUOTA_TYPE=quota_type
--node-count ${NODE_COUNT} \ --node-prefix ${NODE_PREFIX} # optional, the default is QUEUED_RESOURCE_ID
Créez une ressource TPU.
gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \ --node-id ${TPU_NAME} \ --project ${PROJECT_ID} \ --zone ${ZONE} \ --accelerator-type ${ACCELERATOR_TYPE} \ --runtime-version ${RUNTIME_VERSION} \ --valid-until-duration ${VALID_DURATION} \ --service-account ${SERVICE_ACCOUNT} \ --${QUOTA_TYPE}
Vous pourrez vous connecter à vos VM TPU à l'aide de SSH une fois que votre
QueuedResource
aura l'étatACTIVE
:Utilisez la commande
describe
pour interroger l'état de votre ressource en file d'attente.gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} --project ${PROJECT_ID} --zone ${ZONE}
Lorsque votre ressource en file d'attente est à l'état ACTIVE, le résultat ressemble à ce qui suit:
state: ACTIVE
Se connecter à votre TPU à l'aide de SSH
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project ${PROJECT_ID} \ --zone ${ZONE}
Installer des dépendances
export TPU_NAME=your_tpu_name export MAXTEXT_OUTPUT_PATH=output-path
cd maxtext && python3 multihost_runner.py --TPU_PREFIX=${TPU_NAME} \ --COMMAND='bash setup.sh'
Exécutez le modèle avec divers scripts de configuration, tels que 32b.sh et 64b.sh. Si vous exécutez le script à partir d'une VM TPU, vous devez ajouter l'option
--INTERNAL_IP=true
.python3 multihost_runner.py --TPU_PREFIX=${TPU_NAME} \ --COMMAND="bash MaxText/configs/experimental/64b.sh RUN_NAME=${RUN_NAME} OUTPUT_PATH=${MAXTEXT_OUTPUT_PATH} PLATFORM=gce"
Effectuer un nettoyage
Supprimez votre TPU et les ressources en file d'attente.
Résultats du benchmark
Le script d'entraînement MaxText a été exécuté de 32 à 1 160 milliards avec une précision en bf16. Les résultats de ces exécutions sont présentés dans le tableau suivant.
Nbre de paramètres |
Type d'accélérateur |
TFLOP/puce/s |
Utilisation des FLOPS du modèle (MFU) |
---|---|---|---|
32 Mrds |
V5P-128 |
3,28E+02 |
71,47% |
64 Mrds |
V5P-128 |
3,23E+02 |
70,31% |
128 Mrds |
V5P-256 |
3,15E+02 |
68,68% |
128 Mrds |
V5P-512 |
3,15E+02 |
68,53% |
256 Mrds |
V5P-1024 |
3,16E+02 |
68,82% |
512 Mrds |
V5P-1024 |
2,94E+02 |
63,99% |
1024B |
V5P-2048 |
2,49E+02 |
64,05% |
1024B |
V5P-4096 |
2,97E+02 |
64,80% |
1 160 Mrds |
V5P-7680 |
2,95E+02 |
64,27% |
1 160 Mrds |
V5P-12288 |
3,04E + 02 |
66,23% |
Le modèle de paramètres 256B a été testé sur v5p-512 et v5p-1024 en utilisant à la fois des pondérations bf16 et int8. Le tableau suivant affiche les résultats de ces tests.
V5P-512 |
V5P-512 |
V5P-1024 |
V5P-1024 |
|
---|---|---|---|---|
Taille du lot global (jetons) |
5,24E+05 |
5,24E+05 |
1,05E+06 |
1,05E+06 |
Précision |
bf16 |
int8 |
bf16 |
int8 |
TFLOP/puce/s |
307 |
408 |
308 |
414 |
Utilisation des FLOPS du modèle (MFU) |
66,98% |
88,85% |
67,09% |
90,23% |
Tutoriels TensorFlow
Entraîner ResNet sur un seul hôte v5p
Ce tutoriel explique comment entraîner ImageNet sur un TPU v5p-8
à l'aide d'un ensemble de données factice. Si vous souhaitez utiliser un autre ensemble de données, consultez la section Préparer l'ensemble de données.
Préparation
Créez des variables d'environnement :
export PROJECT_ID=your-project-ID export ACCELERATOR_TYPE=v5p-8 export ZONE=us-east1-c export RUNTIME_VERSION=tpu-vm-tf-2.16.1-pjrt export TPU_NAME=your-tpu-name export QUEUED_RESOURCE_ID=your-queued-resource-id export QUOTA_TYPE=quota-type
Pour ce tutoriel, utilisez
v5p-8
commeACCELERATOR_TYPE
.-
gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \ --node-id ${TPU_NAME} \ --project ${PROJECT_ID} \ --zone ${ZONE} \ --accelerator-type ${ACCELERATOR_TYPE} \ --runtime-version ${RUNTIME_VERSION} \ --${QUOTA_TYPE}
Vous pourrez vous connecter à votre VM TPU via SSH une fois que votre ressource en file d'attente passe à l'état
ACTIVE
. Pour vérifier l'état de votre ressource en file d'attente, exécutez la commande suivante:gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \ --project ${PROJECT_ID} \ --zone ${ZONE}
Se connecter à votre TPU à l'aide de SSH
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project ${PROJECT_ID} \ --zone ${ZONE}
Définir des variables d'environnement
export MODELS_REPO=/usr/share/tpu/models export PYTHONPATH="${MODELS_REPO}:${PYTHONPATH}" export MODEL_DIR=gcp-directory-to-store-model export DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet export NEXT_PLUGGABLE_DEVICE_USE_C_API=true export TF_PLUGGABLE_DEVICE_LIBRARY_PATH=/lib/libtpu.so
Accédez au répertoire du dépôt de modèles et installez les conditions requises.
cd ${MODELS_REPO} && git checkout r2.15.0 pip install -r official/requirements.txt
Entraîner le modèle
Exécutez le script d'entraînement.
python3 official/vision/train.py \ --tpu=local \ --experiment=resnet_imagenet \ --mode=train_and_eval \ --config_file=official/vision/configs/experiments/image_classification/imagenet_resnet50_tpu.yaml \ --model_dir=${MODEL_DIR} \ --params_override="runtime.distribution_strategy=tpu,task.train_data.input_path=${DATA_DIR}/train*,task.validation_data.input_path=${DATA_DIR}/validation*,task.train_data.global_batch_size=2048,task.validation_data.global_batch_size=2048,trainer.train_steps=100"
Effectuer un nettoyage
Supprimez votre TPU et les ressources en file d'attente.
Entraîner ResNet sur un serveur v5p à hôtes multiples
Ce tutoriel explique comment entraîner ImageNet sur v5p-16
ou version ultérieure à l'aide d'un ensemble de données factice. Si vous souhaitez utiliser un autre ensemble de données, consultez la section Préparer l'ensemble de données.
Créez des variables d'environnement :
export PROJECT_ID=your_project_ID export TPU_NAME=your_tpu_name export ZONE=us-east1-c export ACCELERATOR_TYPE=v5p-16 export RUNTIME_VERSION=tpu-vm-tf-2.16.1-pod-pjrt export QUEUED_RESOURCE_ID=your-queued-resource-id export QUOTA_TYPE=quota-type
ACCELERATOR_TYPE
peut être de taillev5p-16
ou supérieure.-
gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \ --node-id ${TPU_NAME} \ --project ${PROJECT_ID} \ --zone ${ZONE} \ --accelerator-type ${ACCELERATOR_TYPE} \ --runtime-version ${RUNTIME_VERSION} \ --${QUOTA_TYPE}
Vous pourrez vous connecter à votre VM TPU via SSH une fois que votre ressource en file d'attente passe à l'état
ACTIVE
.Utilisez la commande
describe
pour interroger l'état de votre ressource en file d'attente:gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \ --project ${PROJECT_ID} \ --zone ${ZONE}
Se connecter à votre TPU (nœud de calcul zéro) à l'aide de SSH
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project ${PROJECT_ID} \ --zone ${ZONE}
Définir des variables d'environnement
export TPU_NAME=your_tpu_name export MODELS_REPO=/usr/share/tpu/models export PYTHONPATH="${MODELS_REPO}:${PYTHONPATH}" export MODEL_DIR=gcp-directory-to-store-model export DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet export TPU_LOAD_LIBRARY=0
Accédez au répertoire du dépôt de modèles et installez les conditions requises.
cd $MODELS_REPO && git checkout r2.15.0 pip install -r official/requirements.txt
Entraîner le modèle
Exécutez le script d'entraînement.
python3 official/vision/train.py \ --tpu=${TPU_NAME} \ --experiment=resnet_imagenet \ --mode=train_and_eval \ --config_file=official/vision/configs/experiments/image_classification/imagenet_resnet50_tpu.yaml \ --model_dir=${MODEL_DIR} \ --params_override="runtime.distribution_strategy=tpu,task.train_data.input_path=${DATA_DIR}/train*,task.validation_data.input_path=${DATA_DIR}/validation*,task.train_data.global_batch_size=2048,task.validation_data.global_batch_size=2048,trainer.train_steps=100"
Effectuer un nettoyage
Supprimez votre TPU et les ressources en file d'attente.
PyTorch/XLA
Llama 2
Ce tutoriel explique comment entraîner le modèle Llama 2 7B sur v5p à l'aide d'une copie du dépôt HuggingFace sur PyTorch/XLA avec une parallélisation générale et évolutive pour les graphiques de calcul ML (GSPMD, General and Scalable Parallelization for ML Computation Graphs).
Préparation
Créez des variables pour l'ID de projet, le type d'accélérateur, la zone, la version d'exécution et le nom du TPU.
export PROJECT_ID=your_project_ID export ACCELERATOR_TYPE=v5p-8 export ZONE=us-east5-a export RUNTIME_VERSION=v2-alpha-tpuv5 export SERVICE_ACCOUNT=your_service_account export TPU_NAME=your_tpu_name export QUEUED_RESOURCE_ID=your_queued_resource_id export QUOTA_TYPE=quota_type export VALID_DURATION=1d
Créer une ressource TPU
gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \ --node-id ${TPU_NAME} \ --project ${PROJECT_ID} \ --zone ${ZONE} \ --accelerator-type ${ACCELERATOR_TYPE} \ --runtime-version ${RUNTIME_VERSION} \ --valid-until-duration ${VALID_DURATION} \ --service-account ${SERVICE_ACCOUNT} \ --${QUOTA_TYPE}
Vous pourrez vous connecter à votre VM TPU à l'aide de SSH une fois que votre
QueuedResource
sera à l'étatACTIVE
:Utilisez la commande
describe
pour interroger l'état de votre ressource en file d'attente.gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \ --project ${PROJECT_ID} \ --zone ${ZONE}
Lorsque votre ressource en file d'attente est à l'état ACTIVE, le résultat ressemble à ce qui suit :
state: ACTIVE
Installez Pytorch/XLA et les dépendances requises.
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project ${PROJECT_ID} \ --zone ${ZONE} \ --worker=all \ --command=' sudo apt-get update sudo apt-get install libopenblas-dev -y pip3 install numpy pip3 install typing-extensions pip install torch~=2.2.0 torch_xla[tpu]~=2.2.0 -f https://storage.googleapis.com/libtpu-releases/index.html '
Téléchargez le dépôt HuggingFace et les conditions d'installation requises.
gcloud compute tpus tpu-vm ssh ${TPU_NAME} --project=${PROJECT_ID} \ --zone=${ZONE} \ --worker=all \ --command=' git clone -b llama2-google-next-training https://github.com/pytorch-tpu/transformers.git cd transformers pip3 install git+file://$PWD pip3 install datasets accelerate evaluate scikit-learn'
Téléchargez la configuration du modèle 7B.
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --worker=all \ --command="curl https://huggingface.co/TheBloke/Llama-2-7B-fp16/raw/main/config.json --output ~/config.json"
Entraîner le modèle
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --worker=all \ --command=' export PJRT_DEVICE=TPU export XLA_USE_BF16=1 export XLA_IR_DEBUG=1 export XLA_HLO_DEBUG=1 export LIBTPU_INIT_ARGS="--xla_enable_async_collective_permute=true --xla_tpu_enable_async_collective_fusion_multiple_steps=true --xla_tpu_enable_async_collective_fusion=true --xla_tpu_overlap_compute_collective_tc=true --xla_enable_async_all_gather=true --xla_jf_spmd_threshold_for_windowed_einsum_mib=0" export PROFILE_EPOCH=0 export PROFILE_STEP=3 export PROFILE_DURATION_MS=20000 export PROFILE_LOGDIR=/tmp/home/ cd transformers python examples/pytorch/language-modeling/run_clm.py \ --tokenizer_name hf-internal-testing/llama-tokenizer \ --dataset_name wikitext \ --dataset_config_name wikitext-2-raw-v1 \ --per_device_train_batch_size 96 \ --per_device_eval_batch_size 8 \ --num_train_epochs 1 \ --do_train \ --output_dir /tmp/output \ --overwrite_output_dir \ --config_name ~/config.json \ --save_strategy no \ --logging_strategy no \ --remove_unused_columns no \ --optim adafactor \ --torch_dtype bfloat16 \ --dataloader_drop_last yes \ --block_size 2048 \ --spmd_2d_sharding 1 \ --spmd_grad_chkpt '
Si vous exécutez l'application dans un environnement à plusieurs tranches, vous devez définir l'option --spmd_dcn_parallelism
sur le nombre de tranches.
SPMD_USER_GUIDE fournit un guide utilisateur plus détaillé qui explique les différentes variables d'environnement et les boutons d'activation/de désactivation du script HF. Notez que LIBTPU_INIT_ARGS sera intégré à PyTorch/XLA et activé par défaut dans les prochaines versions.
Effectuer un nettoyage
Supprimez votre TPU et les ressources en file d'attente.
Résultats du benchmark
Le débit des trois tailles de modèle Llama 2 est inclus dans le tableau suivant.
V5p-8 |
V5P-128 |
V5P-128 |
|
---|---|---|---|
Taille du modèle |
7 Mrds |
13 Mrds |
70 Mrds |
Taille du lot global |
96 |
1 024 |
128 |
Forme du maillage de segmentation |
(4, 1) |
(64, 1) |
(16, 4) |
Utilisation des FLOPS du modèle (MFU) |
56,67% |
55,80% |
51,85% |
Assistance et commentaires
Tous vos commentaires sont les bienvenus. Pour partager vos commentaires ou demander de l'aide, remplissez le formulaire de commentaires ou d'assistance Cloud TPU.