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 Cluster Toolkit.

TPU7x ist die erste Version der Ironwood-Familie,der siebten Generation von TPUs von Google Cloud. 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 auf GitHub unter Training Recipes for Ironwood TPU.

Trainingsarbeitslast mit MaxText und Cluster Toolkit bereitstellen

Mit dem Cluster Toolkit können Sie produktionsbereite GKE-Cluster erstellen und Trainingsarbeitslasten ausführen.

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

Hinweis

Führen Sie zuerst die folgenden Schritte aus:

  1. Legen Sie die folgenden Umgebungsvariablen 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, der das Ausgabeverzeichnis für das Modelltraining ist.

  2. Wenn Sie noch 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

TPU7x-Cluster (Ironwood) bereitstellen

Stellen Sie einen GKE-TPU7x-Cluster bereit, indem Sie der Anleitung unter GKE-TPU7x-Cluster bereitstellen folgen. Achten Sie darauf, dass die Topologie des Knotenpools Ihres Clusters den Anforderungen Ihrer Trainingsarbeitslast entspricht, z. B. 4x4x8 für die Beispielarbeitslast „MaxText“ unten.

MaxText-Docker-Image erstellen und hochladen

Erstellen und übertragen Sie das MaxText-Docker-Image:

# 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 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 Language Models mit etwa 1 Milliarde Parametern entwickelt wurde. Verwenden Sie dieses Modell, um mit kleinen Chip-Skalen 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_training \
        steps=30"

Trainingsarbeitslast bereitstellen

Senden Sie Ihren Trainingsjob mit gcluster job submit. Weitere Informationen zu den Optionen zum Einreichen von Jobs finden Sie im Cluster Toolkit Job Guide:

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, um die Eindeutigkeit zu gewährleisten.

Nächste Schritte