Google uses AI technology to translate content into your preferred language. AI translations can contain errors.

Modell mit TPU7x (Ironwood) trainieren

In diesem Dokument wird beschrieben, wie Sie TPU7x-Ressourcen bereitstellen. Außerdem finden Sie ein Beispiel für die Bereitstellung einer Trainingsarbeitslast mit MaxText und XPK.

TPU7x ist die erste Version der Ironwood-Familie, Google Cloud's siebten Generation TPU. Die Ironwood-Generation wurde für umfangreiches KI-Training und Inferenz entwickelt. Weitere Informationen finden Sie unter TPU7x.

Weitere für TPU7x optimierte Beispiele finden Sie unter Training Recipes for Ironwood TPU auf GitHub.

Trainingsarbeitslast mit MaxText und XPK bereitstellen

Mit dem Accelerated Processing Kit (XPK) können Sie GKE-Cluster für Proof-of-Concept- und Testzwecke erstellen.

In den folgenden Abschnitten wird beschrieben, wie Sie eine Trainingsarbeitslast mit MaxText und XPK bereitstellen.

Hinweis

Führen Sie vor dem Start die folgenden Schritte aus:

Sie benötigen ein Google Cloud Projekt mit aktivierter Abrechnung.
Greifen Sie auf TPU7x zu. Weitere Informationen erhalten Sie von Ihrem Account-Management-Team.
Achten Sie darauf, dass das Konto, das Sie mit XPK verwenden, die im XPK GitHub Repository aufgeführten Rollen hat.

XPK und Abhängigkeiten installieren

Installieren Sie XPK. Folgen Sie der Anleitung im XPK-GitHub-Repository.
Installieren Sie Docker gemäß der Anleitung Ihres Administrators oder folgen Sie den offiziellen Installationsanleitungen. Führen Sie nach der Installation die folgenden Befehle aus, um Docker zu konfigurieren und die Installation zu testen:
```
gcloud auth configure-docker
sudo usermod -aG docker $USER # relaunch the terminal and activate venv after running this command
docker run hello-world # Test Docker
```
Legen Sie die folgenden Umgebungsvariablen fest:
```
export PROJECT_ID=YOUR_PROJECT_ID
export ZONE=YOUR_ZONE
export CLUSTER_NAME=YOUR_CLUSTER_NAME
export ACCELERATOR_TYPE=YOUR_ACCELERATOR_TYPE
export BASE_OUTPUT_DIR="gs://YOUR_BUCKET_NAME"
```
Ersetzen Sie Folgendes:
- YOUR_PROJECT_ID: Ihre Google Cloud Projekt-ID
- YOUR_ZONE: Die Zone, in der der Cluster erstellt werden soll
- YOUR_CLUSTER_NAME: Der Name des neuen Clusters
- YOUR_ACCELERATOR_TYPE: Die TPU-Version und -Topologie Beispiel: tpu7x-4x4x8 Eine Liste der unterstützten Topologien finden Sie unter Unterstützte Konfigurationen.
- YOUR_BUCKET_NAME: Der Name Ihres Cloud Storage-Bucket, der das Ausgabeverzeichnis für das Modelltraining sein wird

Wenn Sie keinen Cloud Storage-Bucket haben, erstellen Sie einen mit dem folgenden Befehl:

gcloud storage buckets create ${BASE_OUTPUT_DIR} \
    --project=${PROJECT_ID} \
    --location=US \
    --default-storage-class=STANDARD \
    --uniform-bucket-level-access

Cluster mit einer einzelnen NIC und einem einzelnen Slice erstellen

Wählen Sie eine der folgenden Optionen aus, um Ihren Cluster zu erstellen. Für eine optimale Leistung wird die Verwendung eines benutzerdefinierten Netzwerks mit einer MTU von 8.896 empfohlen.

Benutzerdefinierte Netzwerke

So erstellen Sie ein benutzerdefiniertes Netzwerk mit einer MTU von 8.896 und verwenden es für Ihren Cluster:

Legen Sie Umgebungsvariablen für die Netzwerk- und Firewallnamen fest:
```
export NETWORK_NAME=NETWORK_NAME
export NETWORK_FW_NAME=FIREWALL_NAME
```
Ersetzen Sie Folgendes:
- NETWORK_NAME: Ein Name für das Netzwerk
- FIREWALL_NAME: Ein Name für die Netzwerk-Firewallregel

Erstellen Sie ein benutzerdefiniertes Netzwerk mit einer MTU von 8.896:

gcloud compute networks create ${NETWORK_NAME} \
    --mtu=8896 \
    --project=${PROJECT_ID} \
    --subnet-mode=auto \
    --bgp-routing-mode=regional

Erstellen Sie eine Firewallregel, die TCP-, ICMP- und UDP-Traffic in Ihrem Netzwerk zulässt:

gcloud compute firewall-rules create ${NETWORK_FW_NAME} \
    --network=${NETWORK_NAME} \
    --allow tcp,icmp,udp \
    --project=${PROJECT_ID}

Legen Sie eine Umgebungsvariable für die XPK-Clusterargumente fest, um das erstellte Netzwerk zu verwenden:
```
export CLUSTER_ARGUMENTS="--network=${NETWORK_NAME} --subnetwork=${NETWORK_NAME}"
```

Erstellen Sie den XPK-Cluster. Mit dem folgenden Befehl wird On-Demand-Kapazität bereitgestellt:

xpk cluster create --cluster=${CLUSTER_NAME} \
    --cluster-cpu-machine-type=n1-standard-8 \
    --num-slices=${NUM_SLICES} \
    --tpu-type=${ACCELERATOR_TYPE} \
    --zone=${ZONE} \
    --project=${PROJECT_ID} \
    --on-demand \
    --custom-cluster-arguments="${CLUSTER_ARGUMENTS}"

Wenn Sie reservierte Kapazität verwenden möchten, ersetzen Sie --on-demand durch --reservation=RESERVATION_NAME. Wenn Sie TPU-Spot-VMs verwenden möchten, ersetzen Sie --on-demand durch --spot.

Standardnetzwerk

Wenn Sie kein Netzwerk mit hoher MTU benötigen, können Sie einen Cluster erstellen, der das Standard-VPC-Netzwerk verwendet. Mit dem folgenden Befehl wird On-Demand-Kapazität bereitgestellt:

xpk cluster create --cluster=${CLUSTER_NAME} \
    --cluster-cpu-machine-type=n1-standard-8 \
    --num-slices=${NUM_SLICES} \
    --tpu-type=${ACCELERATOR_TYPE} \
    --zone=${ZONE} \
    --project=${PROJECT_ID} \
    --on-demand

MaxText-Docker-Image erstellen oder hochladen

Sie können entweder ein Docker-Image lokal mit von MaxText bereitgestellten Skripts erstellen oder ein vorgefertigtes Image verwenden.

Lokal erstellen

Mit den folgenden Befehlen wird Ihr lokales Verzeichnis in den Container kopiert:

# Make sure you're running on a virtual environment with python3.12. If nothing is printed, you have the correct version.
[[ "$(python3 -c 'import sys; print(f"{sys.version_info.major}.{sys.version_info.minor}")' 2>/dev/null)" == "3.12" ]] || { >&2 echo "Error: Python version must be 3.12."; false; }

# Clone MaxText
git clone https://github.com/AI-Hypercomputer/maxtext.git
cd maxtext
git checkout maxtext-tutorial-v1.0.0

# Build the Docker image
bash docker_build_dependency_image.sh MODE=stable JAX_VERSION=0.8.2

Nachdem die Befehle erfolgreich ausgeführt wurden, sollte lokal ein Image mit dem Namen maxtext_base_image erstellt worden sein. Sie können Ihr lokales Image direkt im Befehl für die XPK-Arbeitslast verwenden.

Image hochladen (optional)

Nachdem Sie das Docker-Image lokal gemäß der Anleitung im vorherigen Abschnitt erstellt haben, können Sie das MaxText-Docker-Image mit dem folgenden Befehl in die Registry hochladen:

export CLOUD_IMAGE_NAME="${USER}-maxtext-runner"
bash docker_upload_runner.sh CLOUD_IMAGE_NAME=${CLOUD_IMAGE_NAME}

Nachdem dieser Befehl erfolgreich ausgeführt wurde, sollte das MaxText-Image in gcr.io mit dem Namen gcr.io/PROJECT_ID/CLOUD_IMAGE_NAME angezeigt werden.

MaxText-Trainingsbefehl definieren

Bereiten Sie den Befehl vor, um Ihr Trainingsskript im Docker-Container auszuführen.

Das MaxText 1B-Modell ist eine Konfiguration im MaxText-Framework, die für das Training eines Language Models mit etwa 1 Milliarde Parametern entwickelt wurde. Verwenden Sie dieses Modell, um mit kleinen Chip-Skalierungen zu experimentieren. Die Leistung ist nicht optimiert.

export MAXTEXT_COMMAND="JAX_PLATFORMS=tpu,cpu \
    ENABLE_PJRT_COMPATIBILITY=true \
    python3 src/MaxText/train.py src/MaxText/configs/base.yml \
        base_output_directory=${BASE_OUTPUT_DIR} \
        dataset_type=synthetic \
        per_device_batch_size=2 \
        enable_checkpointing=false \
        gcs_metrics=true \
        run_name=maxtext_xpk \
        steps=30"

Trainingsarbeitslast bereitstellen

Führen Sie den Befehl xpk workload create aus, um Ihren Trainingsjob bereitzustellen. Sie müssen entweder das Flag --base-docker-image angeben, um das MaxText-Basis-Image zu verwenden, oder das Flag --docker-image und das gewünschte Image. Sie können das Flag --enable-debug-logs einfügen, um das Debugging-Logging zu aktivieren.

xpk workload create \
    --cluster ${CLUSTER_NAME} \
    --base-docker-image maxtext_base_image \
    --workload maxtext-1b-$(date +%H%M) \
    --tpu-type=${ACCELERATOR_TYPE} \
    --zone ${ZONE} \
    --project ${PROJECT_ID} \
    --command "${MAXTEXT_COMMAND}"
    # [--enable-debug-logs]

Arbeitslastnamen müssen im Cluster eindeutig sein. In diesem Beispiel wird $(date +%H%M) an den Arbeitslastnamen angehängt, um die Eindeutigkeit zu gewährleisten.

Nächste Schritte

Verwenden Sie die Google Cloud ML Diagnostics-Plattform, um Ihre Arbeitslasten zu optimieren und zu diagnostizieren
Trainingsarbeitslast mit einem für TPU7x optimierten Rezept ausführen
TPU7x-Microbenchmark ausführen