TPU-Slice mit mehreren Hosts erstellen

TPU-Slice mit mehreren Hosts erstellen

Hier erfahren Sie, wie Sie ein TPU-Slice mit mehreren Hosts mithilfe einer verwalteten Instanzgruppe (Managed Instance Group, MIG) erstellen, eine Verbindung zum Slice herstellen und eine Berechnung ausführen. In dieser Kurzanleitung wird die Option für die On-Demand-Nutzung verwendet. Führen Sie die Befehle in dieser Kurzanleitung in Ihrem lokalen Terminal oder in Cloud Shell aus.

Hinweis

1. Installieren Sie die Google Cloud CLI.

   Hinweis:Wenn Sie die gcloud CLI zuvor installiert haben, prüfen Sie, ob Sie die aktuelle Version haben, indem Sie gcloud components update ausführen.

2. Konfigurieren Sie die gcloud CLI für die Verwendung Ihrer föderierten Identität.

   Weitere Informationen finden Sie unter Mit Ihrer föderierten Identität in der gcloud CLI anmelden.

3. Führen Sie den folgenden Befehl aus, um die gcloud CLI zu initialisieren:

   gcloud init

4. Google Cloud Projekt erstellen oder auswählen

   Erforderliche Rollen zum Auswählen oder Erstellen eines Projekts

   • Projekt auswählen: Für die Auswahl eines Projekts ist keine bestimmte IAM-Rolle erforderlich. Sie können jedes Projekt auswählen, für das Ihnen eine Rolle zugewiesen wurde.
   • Projekt erstellen: Zum Erstellen eines Projekts benötigen Sie die Rolle „Project Creator“ (roles/resourcemanager.projectCreator), die die Berechtigung resourcemanager.projects.create enthält. Informationen zum Zuweisen von Rollen
   Hinweis: Wenn Sie die Ressourcen, die Sie in diesem Verfahren erstellen, nicht behalten möchten, erstellen Sie ein Projekt, anstatt ein vorhandenes Projekt auszuwählen. Wenn Sie fertig sind, können Sie das Projekt löschen und dadurch alle mit dem Projekt verknüpften Ressourcen entfernen.

   • So erstellen Sie ein Google Cloud Projekt:

     gcloud projects create PROJECT_ID

     Ersetzen Sie PROJECT_ID durch einen Namen für das Google Cloud Projekt, das Sie erstellen.

   • Wählen Sie das von Ihnen erstellte Google Cloud Projekt aus:

     gcloud config set project PROJECT_ID

     Ersetzen Sie PROJECT_ID durch Ihren Google Cloud Projektnamen.

5. Wenn Sie für diese Anleitung ein vorhandenes Projekt verwenden, prüfen Sie, ob Sie die erforderlichen Berechtigungen haben. Wenn Sie ein neues Projekt erstellt haben, haben Sie bereits die erforderlichen Berechtigungen.

6. Prüfen Sie, ob die Abrechnung für Ihr Google Cloud Projekt aktiviert ist.

7. Aktivieren Sie die Compute Engine API, falls sie noch nicht aktiviert ist:

   Rollen, die zum Aktivieren von APIs erforderlich sind

   Zum Aktivieren von APIs benötigen Sie die Berechtigung serviceusage.services.enable. Wenn Sie das Projekt erstellt haben, haben Sie diese Berechtigung wahrscheinlich bereits über die Rolle „Inhaber“ (roles/owner). Andernfalls können Sie diese Berechtigung über die Rolle „Service Usage-Administrator“ (roles/serviceusage.serviceUsageAdmin) erhalten. Informationen zum Zuweisen von Rollen

   gcloud services enable compute.googleapis.com

Erforderliche Rollen

Bitten Sie Ihren Administrator, Ihnen die folgenden IAM-Rollen für Ihr Projekt zuzuweisen, um die Berechtigungen zu erhalten, die Sie zum Erstellen einer MIG benötigen, die einen TPU-Slice mit mehreren Hosts bildet, und um eine SSH-Verbindung zu jeder VM in der MIG herzustellen und Befehle auszuführen:

• Compute Instance Admin (v1) (roles/compute.instanceAdmin.v1)
• Service Account User (roles/iam.serviceAccountUser)
• Service Usage-Administrator ()roles/serviceusage.serviceUsageAdmin

Weitere Informationen zum Zuweisen von Rollen finden Sie unter Zugriff auf Projekte, Ordner und Organisationen verwalten.

Sie können die erforderlichen Berechtigungen auch über benutzerdefinierte Rollen oder andere vordefinierte Rollen erhalten.

Instanzvorlage erstellen

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

Hinweis :Wenn Sie einen anderen Maschinentyp verwenden, müssen Sie ein unterstütztes Betriebssystem-Image verwenden. Eine Liste der von Google verwalteten Betriebssystem-Images für jede unterstützte TPU-Version finden Sie unter TPU-Betriebssystem-Images.

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 physische Eigenschaften Ihrer Compute-Instanzen. In TPU-Slices definiert die Beschleunigertopologie die physische Anordnung von TPU-Chips. Die Angabe einer Beschleunigertopologie ist für verbundene TPU-Slices mit mehreren Hosts erforderlich.

Verwenden Sie zum Erstellen einer Arbeitslastrichtlinie für ein Multi-Host-TPU-Slice den Befehl gcloud compute resource-policies create workload-policy mit dem Flag --accelerator-topology. Mit dem folgenden Befehl wird eine Arbeitslastrichtlinie mit einer 2x4-Topologie erstellt:

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.

Hinweis :Möglicherweise tritt ein vorübergehender Fehler bei der Ressourcenverfügbarkeit auf, während die MIG auf Kapazität wartet. Wenn Sie längere Wartezeiten haben, können Sie es mit einer anderen Region versuchen. Eine Liste der unterstützten Standorte finden Sie unter TPU-Regionen und ‑Zonen. Wenn Sie eine andere Region auswählen, müssen Sie zuerst eine Arbeitslastrichtlinie in dieser Region erstellen. Wiederholen Sie dazu den Schritt Arbeitslastrichtlinie erstellen. Achten Sie darauf, die Region in allen nachfolgenden Befehlen zu aktualisieren.

1. Verwenden Sie den Befehl gcloud compute instance-groups managed create, um eine MIG zu erstellen, die ein 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 dabei PROJECT_ID durch die ID Ihres Projekts in Google Cloud .

2. Prüfen Sie mit den folgenden Befehlen, ob die verwalteten Instanzen ausgeführt werden:

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

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

   • Verwenden Sie den Befehl gcloud compute instance-groups managed list-instances, um den Status der einzelnen Instanzen aufzurufen:

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

   Achtung :Alle Instanzen müssen den Status RUNNING haben, bevor Sie fortfahren. Der Versuch, über SSH eine Verbindung zu Instanzen herzustellen, während sie noch bereitgestellt werden, führt zu einem Fehler.

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 aktuelle 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 jax.device_count()-Funktionsaufruf reagiert nicht mehr, bis er auf jedem Host im Slice aufgerufen wurde. 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.]

Bereinigen

Löschen Sie das Projekt von Google Cloud zusammen mit den Ressourcen, damit Ihrem Konto von Google Cloud die auf dieser Seite verwendeten Ressourcen nicht in Rechnung gestellt werden.

Wenn Sie Ihr Projekt behalten möchten, können Sie auch nur die MIG und alle VMs in der Gruppe mit dem Befehl gcloud compute instance-groups managed delete löschen:

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

Nächste Schritte

• TPU-Ressourcen in Compute Engine
• TPU-Slice mit mehreren Hosts mit anderen Nutzungsoptionen erstellen
• TPU-Instanzen in verwalteten Instanzgruppen