TPU-Slice mit mehreren Hosts erstellen

Instanzvorlage erstellen

Verwenden Sie den gcloud compute instance-templates create Befehl, um eine Instanzvorlage für TPU v6e-VMs zu erstellen:

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

Arbeitslastrichtlinie erstellen

Eine Arbeitslastrichtlinie definiert die physischen Eigenschaften Ihrer Compute-Instanzen. In TPU-Slices definiert die Beschleunigertopologie die physische Anordnung der TPU-Chips. Die Angabe einer Beschleunigertopologie ist für miteinander verbundene TPU-Slices mit mehreren Hosts erforderlich.

Verwenden Sie den gcloud compute resource-policies create workload-policy Befehl mit dem --accelerator-topology Flag, um eine Arbeitslastrichtlinie für einen TPU-Slice mit mehreren Hosts zu erstellen. Der folgende Befehl erstellt eine Arbeitslastrichtlinie mit einer 2x4-Topologie:

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

Eine MIG erstellen.

Führen Sie die folgenden Befehle aus, um eine MIG zu erstellen, die einen TPU-Slice mit mehreren Hosts bildet.

1. Verwenden Sie den gcloud compute instance-groups managed create Befehl, um eine MIG zu erstellen, die einen TPU-Slice mit mehreren Hosts bildet:

   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
   

   Ersetzen Sie PROJECT_ID durch die Cloud de Confiance by S3NS Projekt-ID Ihres Projekts.

2. Optional können Sie mit den folgenden Befehlen prüfen, ob die verwalteten Instanzen ausgeführt werden:

   • Verwenden Sie den gcloud compute instance-groups managed describe Befehl, um den Gesamtstatus der MIG aufzurufen:

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

   • Verwenden Sie den gcloud compute instance-groups managed list-instances Befehl, um den Status jeder Instanz aufzurufen:

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

JAX installieren

Installieren Sie Abhängigkeiten und das JAX-Framework in einer virtuellen Umgebung auf allen TPU-VM-Instanzen in der MIG. Wenn auf Ihren TPU-VMs eine Python-Version vor 3.11 installiert ist, müssen Sie Python 3.11 installieren, um die neueste Version von JAX auszuführen.

1. Prüfen Sie, welche Python-Version auf Ihren TPU-VMs ausgeführt wird:

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

   Wenn die Version älter als Python 3.11 ist, installieren Sie 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. Virtuelle Umgebung erstellen:

   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. Installieren Sie JAX in der virtuellen Umgebung:

   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'
   

JAX-Code auf dem Slice ausführen

Um JAX-Code auf einem TPU-Slice auszuführen, müssen Sie den Code auf jedem Host im TPU-Slice ausführen. Der Funktionsaufruf jax.device_count() reagiert erst, wenn er auf jedem Host im Slice aufgerufen wird. Das folgende Beispiel zeigt, wie Sie auf einem TPU-Slice eine JAX-Berechnung ausführen.

Code vorbereiten

Erstellen Sie auf jeder Instanz eine Datei mit dem Namen example.py:

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"

Den Code auf dem Slice ausführen

Führen Sie das Programm example.py auf jeder TPU-VM im Slice aus:

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'

Die Ausgabe sollte in etwa so aussehen:

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