Cette page a été traduite par l'API Cloud Translation.

Exécuter du code JAX sur des tranches de pods Cloud TPU

Une fois que votre code JAX s'exécute sur une seule carte TPU, vous pouvez augmenter la capacité en l'exécutant sur une tranche de pod TPU. Les tranches de pod TPU sont des cartes de TPU interconnectées sur des connexions réseau haut débit dédiées. Ce document est une introduction à l'exécution de code JAX sur des tranches de pods TPU. Pour des informations plus détaillées, consultez la section Utiliser JAX dans des environnements multihôtes et multiprocessus.

Si vous souhaitez utiliser le stockage NFS installé pour le stockage de données, vous devez définir OS Login pour toutes les VM TPU de la tranche de pod. Pour en savoir plus, consultez la section Utiliser NFS pour le stockage de données.

Configurer votre environnement

Dans Cloud Shell, exécutez la commande suivante pour vous assurer d'exécuter la version actuelle de gcloud:
```
$ gcloud components update
```
Si vous devez installer gcloud, utilisez la commande suivante:
```
$ sudo apt install -y google-cloud-sdk
```

Créez des variables d'environnement:

$ export TPU_NAME=tpu-name
$ export ZONE=us-central2-b
$ export RUNTIME_VERSION=tpu-ubuntu2204-base
$ export ACCELERATOR_TYPE=v4-32

Créer une tranche de pod TPU

Avant d'exécuter les commandes de ce document, assurez-vous d'avoir suivi les instructions de la section Configurer un compte et un projet Cloud TPU. Exécutez les commandes suivantes sur votre ordinateur local.

Créez une tranche de pod TPU à l'aide de la commande gcloud. Par exemple, pour créer une tranche de pod v4-32, utilisez la commande suivante:

$ gcloud compute tpus tpu-vm create ${TPU_NAME}  \
  --zone=${ZONE} \
  --accelerator-type=${ACCELERATOR_TYPE}  \
  --version=${RUNTIME_VERSION}

Installer JAX sur la tranche de pod

Après avoir créé la tranche de pod TPU, vous devez installer JAX sur tous les hôtes de la tranche de pod TPU. Vous pouvez installer JAX sur tous les hôtes à l'aide d'une seule commande avec l'option --worker=all :

  gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
  --zone=${ZONE} \
  --worker=all \
  --command="pip install \
  --upgrade 'jax[tpu]>0.3.0' \
  -f https://storage.googleapis.com/jax-releases/libtpu_releases.html"

Exécuter du code JAX sur la tranche de pod

Pour exécuter du code JAX sur une tranche de pod TPU, vous devez exécuter le code sur chaque hôte de la tranche de pod TPU. L'appel jax.device_count() cesse de répondre jusqu'à ce qu'il soit appelé sur chaque hôte de la tranche du pod. L'exemple suivant montre comment exécuter un calcul JAX sur une tranche de pod TPU.

Préparer le code

Vous avez besoin de la version gcloud >= 344.0.0 (pour la commande scp). Utilisez gcloud --version pour vérifier votre version de gcloud et exécutez gcloud components upgrade, si nécessaire.

Créez un fichier nommé example.py avec le code suivant:

# The following code snippet will be run on all TPU hosts
import jax

# The total number of TPU cores in the Pod
device_count = jax.device_count()

# The number of TPU cores attached to this host
local_device_count = jax.local_device_count()

# The psum is performed over all mapped devices across the Pod
xs = jax.numpy.ones(jax.local_device_count())
r = jax.pmap(lambda x: jax.lax.psum(x, 'i'), axis_name='i')(xs)

# Print from a single host to avoid duplicated output
if jax.process_index() == 0:
    print('global device count:', jax.device_count())
    print('local device count:', jax.local_device_count())
    print('pmap result:', r)

Copiez `example.py` sur toutes les VM de nœud de calcul TPU de la tranche de pod.

$ gcloud compute tpus tpu-vm scp example.py ${TPU_NAME} \
  --worker=all \
  --zone=${ZONE}

Si vous n'avez jamais utilisé la commande scp, un message d'erreur semblable à celui-ci peut s'afficher:

ERROR: (gcloud.alpha.compute.tpus.tpu-vm.scp) SSH Key is not present in the SSH
agent. Please run `ssh-add /.../.ssh/google_compute_engine` to add it, and try
again.

Pour résoudre l'erreur, exécutez la commande ssh-add comme indiqué dans le message d'erreur, puis réexécutez la commande.

Exécuter le code sur la tranche du pod

Lancez le programme example.py sur chaque VM :

$ gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
  --zone=${ZONE} \
  --worker=all \
  --command="python3 example.py"

Résultat (généré avec une tranche de pod v4-32):

global device count: 16
local device count: 4
pmap result: [16. 16. 16. 16.]

Effectuer un nettoyage

Lorsque vous avez fini d'utiliser votre VM TPU, procédez comme suit pour nettoyer vos ressources.

Déconnectez-vous de l'instance Compute Engine, si vous ne l'avez pas déjà fait :
```
(vm)$ exit
```
Votre invite devrait maintenant être username@projectname, indiquant que vous êtes dans Cloud Shell.

Supprimez vos ressources Cloud TPU et Compute Engine.

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

Vérifiez que les ressources ont été supprimées en exécutant la commande gcloud compute tpus execution-groups list. La suppression peut prendre plusieurs minutes. Le résultat de la commande suivante ne doit inclure aucune des ressources créées dans ce tutoriel:
```
$ gcloud compute tpus tpu-vm list --zone=${ZONE}
```