Diese Seite wurde von der Cloud Translation API übersetzt.

Modell mit TPU7x (Ironwood) trainieren

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

TPU7x ist die erste Version der Ironwood-Familie,der siebten Generation von TPUs von Google Cloud. Die Ironwood-Generation ist für umfangreiches KI-Training und ‑Inferenz konzipiert. Weitere Informationen finden Sie unter TPU7x.

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

TPUs bereitstellen

Sie können TPU7x mit den folgenden Methoden bereitstellen und verwalten:

GKE: Mit der GKE können Sie TPUs als Pool von Beschleunigern für Ihre containerisierten ML-Arbeitslasten bereitstellen und verwalten. Verwenden Sie die Google Cloud CLI, um Ihre GKE-Clusterinstanz manuell zu erstellen, um vorhandene GKE-Produktionsumgebungen präzise anzupassen oder zu erweitern. Weitere Informationen finden Sie unter TPUs in GKE.
GKE und XPK: XPK ist ein Befehlszeilentool, das die Clustererstellung und die Ausführung von Arbeitslasten in der GKE vereinfacht. Es wurde für ML-Nutzer ohne umfassende Kubernetes-Kenntnisse entwickelt, die TPUs bereitstellen und Trainingsjobs ausführen möchten. Mit XPK können Sie schnell GKE-Cluster erstellen und Arbeitslasten für Proof-of-Concept und Tests ausführen. Weitere Informationen finden Sie im XPK-GitHub-Repository.
GKE und TPU Cluster Director: TPU Cluster Director ist über eine Reservierung im Modus „All Capacity“ verfügbar. Damit haben Sie vollen Zugriff auf Ihre gesamte reservierte Kapazität (ohne Holdbacks) und vollständigen Einblick in die TPU-Hardwaretopologie, den Auslastungsstatus und den Integritätsstatus. Weitere Informationen finden Sie unter Übersicht über den Modus „Alle Kapazitäten“.

Trainingsarbeitslast mit MaxText und XPK bereitstellen

Verwenden Sie das Accelerated Processing Kit (XPK), um GKE-Cluster für Proof-of-Concept und Tests zu erstellen. XPK ist ein Befehlszeilentool, das die Bereitstellung, Verwaltung und Ausführung von Arbeitslasten für maschinelles Lernen vereinfacht.

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

Hinweise

Führen Sie zuerst die folgenden Schritte aus:

Sie benötigen ein Google Cloud -Projekt mit aktivierter Abrechnung.
Zugriff auf TPU7x erhalten Weitere Informationen erhalten Sie von Ihrem Account-Management-Team.
Das Konto, das Sie mit XPK verwenden, muss die im XPK-GitHub-Repository aufgeführten Rollen haben.

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 der offiziellen Installationsanleitung. 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 RESERVATION_NAME=YOUR_RESERVATION_NAME
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. In der Vorschau wird nur us-central1-c unterstützt.
- 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_RESERVATION_NAME: Der Name Ihrer Reservierung. Verwenden Sie für freigegebene Reservierungen projects/YOUR_PROJECT_NUMBER/reservations/YOUR_RESERVATION_NAME.
- YOUR_BUCKET_NAME: Der Name Ihres Cloud Storage-Bucket, der das Ausgabeverzeichnis für das Modelltraining ist.

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

Folgen Sie der Anleitung im Abschnitt MTU konfigurieren, um Ihre Netzwerkkonfiguration zu optimieren.
Füllen Sie die Variable ${CLUSTER_ARGUMENTS} aus, die Sie im Befehl xpk cluster create verwenden:
```
export CLUSTER_ARGUMENTS="--network=${NETWORK_NAME} --subnetwork=${SUBNET_NAME}"
```
Erstellen Sie Ihren GKE-Cluster mit TPU7x-Knotenpools mit dem Befehl xpk cluster create:
```
xpk cluster create \
    --project=${PROJECT_ID} \
    --zone=${ZONE} \
    --cluster ${CLUSTER_NAME} \
    --cluster-cpu-machine-type=n1-standard-8 \
    --tpu-type=${ACCELERATOR_TYPE} \
    --reservation=${RESERVATION_NAME} \
    --custom-cluster-arguments="${CLUSTER_ARGUMENTS}"
```
Wenn Sie das Flag --cluster-cpu-machine-type auf n1-standard-8 (oder höher) setzen, hat der Standardknotenpool genügend CPU für System-Pods, z. B. den JobSet-Webhook, wodurch Fehler vermieden werden. Standardmäßig verwendet XPK e2-standard-16. Einige Zonen unterstützen nur bestimmte CPU-Typen. Möglicherweise müssen Sie also zwischen den Typen n1, n2 und e2 wechseln. Andernfalls können Kontingentfehler auftreten.

Fügen Sie einen Wartungsausschluss hinzu, um Upgrades für den Cluster zu verhindern:

gcloud container clusters update ${CLUSTER_NAME} \
    --zone=${ZONE} \
    --add-maintenance-exclusion-name="no-upgrade-next-month" \
    --add-maintenance-exclusion-start="EXCLUSION_START_TIME" \
    --add-maintenance-exclusion-end="EXCLUSION_END_TIME" \
    --add-maintenance-exclusion-scope="no_upgrades"

Ersetzen Sie Folgendes:

EXCLUSION_START_TIME: Die ausgewählte Startzeit für den Wartungsausschluss im Format YYYY-MM-DDTHH:MM:SSZ.
EXCLUSION_END_TIME: Die ausgewählte Endzeit für den Wartungsausschluss im Format YYYY-MM-DDTHH:MM:SSZ.

MaxText-Docker-Image erstellen oder hochladen

Sie können entweder ein Docker-Image lokal mit Skripts erstellen, die von MaxText bereitgestellt werden, oder ein vorgefertigtes Image verwenden.

Build 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

# Custom Jax and LibTPU wheels
pip download libtpu==0.0.28.dev20251104+nightly -f "https://storage.googleapis.com/jax-releases/libtpu_releases.html"
pip download --pre jax==0.8.1.dev20251104 jaxlib==0.8.1.dev20251104 --index https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/

# Build the Docker image
bash docker_build_dependency_image.sh MODE=custom_wheels

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 xpk-Arbeitslastbefehl verwenden.

Bild hochladen (optional)

Nachdem Sie das Docker-Image lokal gemäß der Anleitung im vorherigen Abschnitt erstellt haben, können Sie es 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}

Nach der erfolgreichen Ausführung dieses Befehls sollte das MaxText-Bild in gcr.io mit dem Namen gcr.io/PROJECT_ID/CLOUD_IMAGE_NAME angezeigt werden.

MaxText-Trainingsbefehl definieren

Bereiten Sie den Befehl vor, mit dem Ihr Trainingsskript im Docker-Container ausgeführt wird.

Das MaxText 1B-Modell ist eine Konfiguration im MaxText-Framework, die für das Training eines Sprachmodells mit etwa 1 Milliarde Parametern entwickelt wurde. Mit diesem Modell können Sie mit kleinen Chip-Skalen 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]

Die Namen von Arbeitslasten müssen innerhalb des Clusters eindeutig sein. In diesem Beispiel wird $(date +%H%M) an den Namen der Arbeitslast angehängt, um Eindeutigkeit zu gewährleisten.