Modell mit TPU7x (Ironwood) trainieren

In diesem Dokument wird beschrieben, wie Sie TPU7x-Ressourcen mit dem Cluster Toolkit bereitstellen und eine Trainingsarbeitslast mit MaxText (JAX) oder TorchTitan auf TPU (PyTorch) bereitstellen.

TPU7x ist die erste Version der Ironwood-Familie,der TPU der siebten Generation 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 in den folgenden Repositories auf GitHub:

Hinweis

Bevor Sie beginnen, sollten Sie prüfen, ob die folgenden Anforderungen erfüllt sind:

Führen Sie dann die folgenden Einrichtungsschritte aus:

  1. Legen Sie die Umgebungsvariablen für Ihr Projekt und Ihren Cluster fest:

    export PROJECT_ID=YOUR_PROJECT_ID
export ZONE=YOUR_ZONE
export CLUSTER_NAME=YOUR_CLUSTER_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.
    • YOUR_CLUSTER_NAME: Der Name des neuen Clusters.
    • YOUR_BUCKET_NAME: Der Name Ihres Cloud Storage-Bucket für die Trainingsausgabe.

  2. Wenn Sie noch keinen Cloud Storage-Bucket haben, erstellen Sie einen:

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

TPU7x-Cluster (Ironwood) bereitstellen

Stellen Sie einen Google Kubernetes Engine-Cluster (GKE) mit TPU7x bereit. Folgen Sie dazu der Anleitung unter GKE-Cluster mit TPU7x bereitstellen. Achten Sie darauf, dass die Topologie des Clusterknotenpools den Anforderungen Ihrer Trainingsarbeitslast entspricht, z. B. 4x4x8 für das MaxText-Beispiel oder 2x2x4 für das TorchTitan-Beispiel.

Trainingsarbeitslast mit MaxText bereitstellen

In den folgenden Abschnitten wird beschrieben, wie Sie ein Container-Image erstellen und einen MaxText-Trainingsjob mit dem Cluster Toolkit einreichen.

MaxText-Docker-Image erstellen und hochladen

Erstellen Sie das MaxText-Docker-Image und übertragen Sie es per Push:

# 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

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

MaxText-Trainingsbefehl definieren

Bereiten Sie den Befehl zum Ausführen des Trainingsskripts im Docker-Container vor.

Mit der MaxText-1B-Modellkonfiguration wird ein Sprachmodell mit etwa 1 Milliarde Parametern trainiert. Verwenden Sie dieses Modell, um kleine Chipgrößen zu testen. 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_training \
        steps=30"

MaxText-Trainingsjob senden

Senden Sie Ihren Trainingsjob mit gcluster job submit. Weitere Informationen zu den Optionen für die Jobeinreichung finden Sie im Cluster Toolkit-Jobleitfaden:

gcluster job submit \
    --name="maxtext-1b-$(date +%H%M)" \
    --cluster=${CLUSTER_NAME} \
    --project=${PROJECT_ID} \
    --location=${ZONE} \
    --compute-type=tpu7x-standard-4t \
    --topology=4x4x8 \
    --image="gcr.io/${PROJECT_ID}/${CLOUD_IMAGE_NAME}" \
    --command="${MAXTEXT_COMMAND}"

Arbeitslastnamen müssen innerhalb des Clusters eindeutig sein. In diesem Beispiel wird $(date +%H%M) an den Namen der Arbeitslast angehängt.

Trainingsarbeitslast mit TorchTitan bereitstellen

TorchTitan verwendet TorchTPU, um verteiltes PyTorch-Training auf Cloud TPU auszuführen. Weitere Informationen finden Sie in der README-Datei zu TorchTitan auf TPU.

Nächste Schritte