Créer une tranche TPU multi-hôtes

Créer un modèle d'instance

Pour créer un modèle d'instance pour les VM TPU v6e, utilisez la gcloud compute instance-templates create commande :

gcloud compute instance-templates create quickstart-tpu-instance-template \
    --machine-type=ct6e-standard-4t \
    --maintenance-policy=TERMINATE \
    --image-family=ubuntu-accel-2204-amd64-tpu-v5e-v5p-v6e \
    --image-project=ubuntu-os-accelerator-images \
    --region=us-east5

Créer une règle de charge de travail

Une règle de charge de travail définit les propriétés physiques de vos instances de calcul. Dans les tranches TPU, la topologie d'accélérateur définit la disposition physique des puces TPU. La spécification d'une topologie d'accélérateur est requise pour les tranches TPU multi-hôtes interconnectées.

Pour créer une règle de charge de travail pour une tranche TPU multi-hôtes, utilisez la gcloud compute resource-policies create workload-policy commande avec l'option --accelerator-topology. La commande suivante crée une règle de charge de travail avec une topologie 2x4 :

gcloud compute resource-policies create workload-policy quickstart-tpu-workload-policy \
    --type=high-throughput \
    --accelerator-topology=2x4 \
    --region=us-east5

Créer un MIG

Exécutez les commandes suivantes pour créer un MIG qui forme une tranche TPU multi-hôtes.

1. Pour créer un MIG qui forme une tranche TPU multi-hôtes, utilisez la gcloud compute instance-groups managed create commande :

   gcloud compute instance-groups managed create quickstart-tpu-mig \
       --size=2 \
       --target-size-policy-mode=bulk \
       --template=quickstart-tpu-instance-template \
       --region=us-east5 \
       --target-distribution-shape=any-single-zone \
       --instance-redistribution-type=none \
       --default-action-on-vm-failure=do-nothing \
       --workload-policy=projects/PROJECT_ID/regions/us-east5/resourcePolicies/quickstart-tpu-workload-policy
   

   Remplacez PROJECT_ID par l'ID du Cloud de Confiance by S3NS projet.

2. Vous pouvez également vérifier que les instances gérées sont en cours d'exécution à l'aide des commandes suivantes :

   • Pour afficher l'état global du MIG, utilisez la gcloud compute instance-groups managed describe commande :

     gcloud compute instance-groups managed describe quickstart-tpu-mig \
         --region=us-east5
     

   • Pour afficher l'état de chaque instance, utilisez la gcloud compute instance-groups managed list-instances commande :

     gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
         --region=us-east5
     

Installer JAX

Installez les dépendances et le framework JAX dans un environnement virtuel sur toutes les instances de VM TPU du MIG. Si vos VM TPU sont installées avec une version de Python antérieure à la version 3.11, vous devez installer Python 3.11 pour exécuter la dernière version de JAX.

1. Vérifiez quelle version de Python est exécutée sur vos VM TPU :

   gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
       --region=us-east5 \
       --uri \
   | xargs -I {} -P 0 gcloud compute ssh {} \
       --command='python3 --version'
   

   Si la version est antérieure à Python 3.11, installez Python 3.11 :

   gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
       --region=us-east5 \
       --uri \
   | xargs -I {} -P 0 gcloud compute ssh {} \
       --command='sudo apt update && \
       sudo apt install -y software-properties-common && \
       sudo add-apt-repository -y ppa:deadsnakes/ppa && \
       sudo apt update && \
       sudo apt install -y python3.11 python3.11-dev'
   

2. Créez un environnement virtuel :

   gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
       --region=us-east5 \
       --uri \
   | xargs -I {} -P 0 gcloud compute ssh {} \
       --command='sudo apt install -y python3.11-venv && \
       python3.11 -m venv ~/jax_venv'
   

3. Installez JAX dans l'environnement virtuel :

   gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
       --region=us-east5 \
       --uri \
   | xargs -I {} -P 0 gcloud compute ssh {} \
       --command='source ~/jax_venv/bin/activate && \
       pip install --upgrade pip -q && \
       pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html -q'
   

Exécuter du code JAX sur la tranche

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

Préparer le code

Créez un fichier nommé example.py sur chaque instance :

gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
    --region=us-east5 \
    --uri \
| xargs -I {} -P 0 gcloud compute ssh {} \
    --command="cat << 'EOF' > ~/example.py
import jax

# Initialize the slice
jax.distributed.initialize()

# The total number of TPU cores in the slice
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 slice
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)
EOF"

Exécuter le code sur la tranche

Exécutez le programme example.py sur chaque VM TPU de la tranche :

gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
    --region=us-east5 \
    --uri \
| xargs -I {} -P 0 gcloud compute ssh {} \
    --command='source ~/jax_venv/bin/activate && python3 ~/example.py'

La sortie devrait ressembler à ce qui suit :

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