Batcharbeitslast mit Pathways ausführen

In diesem Leitfaden wird anhand einer JAX-Trainingsarbeitslast, die im selben Cluster wie die Pathways-Controller-Komponenten ausgeführt wird, gezeigt, wie eine Batcharbeitslast ausgeführt wird.

Sie können Batcharbeitslasten in einem Google Kubernetes Engine-Cluster (GKE) oder in einem Cluster Director-Cluster mit Slurm ausführen:

  • In GKE werden die Pathways-Clusterkomponenten nach Abschluss der JAX-Arbeitslast beendet.
  • Auf dem Cluster Director werden die Pathways-Controller-Komponenten als separater Slurm-Job ausgeführt. Sie werden also nach Abschluss der JAX-Arbeitslast weiter ausgeführt, bis Sie sie beenden.

Diese Seite enthält Anleitungen für beide Plattformen. Wählen Sie in jedem Abschnitt den Tab GKE oder Cluster Director aus.

Hinweis

Hinweis

GKE

Erfüllen Sie die folgenden Voraussetzungen:

Cluster Director

Erfüllen Sie die folgenden Voraussetzungen:

Trainings-Image mit MaxText erstellen

MaxText ist ein von Google entwickeltes Open-Source-Projekt für Large Language Models (LLMs). Es ist in JAX geschrieben und für die effiziente Ausführung auf TPUs und GPUs von Google Cloud konzipiert.

GKE

Wenn Sie ein MaxText-Docker-Image mit der aktuellen Version von stabilem JAX aus dem Open-Source-GitHub-Repository erstellen möchten, führen Sie den folgenden Befehl aus:

git clone https://github.com/AI-Hypercomputer/maxtext
cd maxtext
gcloud config set project PROJECT_ID
bash src/dependencies/scripts/docker_build_dependency_image.sh MODE=stable
gcloud auth configure-docker
# This script needs bash version >= 4.2 to execute.
bash src/dependencies/scripts/docker_upload_runner.sh \
    CLOUD_IMAGE_NAME=USER_runner

Ersetzen Sie Folgendes:

  • PROJECT_ID: IhreGoogle Cloud -Projekt-ID.
  • USER: IhreGoogle Cloud -Nutzer-ID.

Mit diesem Befehl wird das MaxText-Kubernetes-Image an gcr.io/PROJECT_ID/USER_runner übertragen. Sie können dieses Docker-Image verwenden, um das Training auf TPUs mit dem Pathways-Backend auszuführen.

Cluster Director

Sie müssen kein Trainings-Image erstellen: setup_env.sh hat MaxText in der Pathways-Umgebung in /home/pathways installiert und run_maxtext.sh führt es aus.

Batcharbeitslast ausführen

GKE

  1. Wenn sich die Binärdatei gcluster nicht in Ihrer Umgebungsvariablen PATH befindet, fügen Sie das Verzeichnis mit gcluster zu PATH hinzu:

    export PATH="$PATH:CLUSTER_TOOLKIT_PATH"

    Ersetzen Sie CLUSTER_TOOLKIT_PATH durch den absoluten Pfad zum Verzeichnis, das die Binärdatei gcluster enthält, z. B. $HOME/cluster-toolkit.

  2. Senden Sie das vordefinierte MaxText-Docker-Image mit dem Befehl gcluster job submit:

    gcluster job submit \
    --pathways \
    --pathways-gcs-location="gs://BUCKET_NAME/pathways-artifacts" \
    --name=WORKLOAD \
    --cluster=CLUSTER \
    --project=PROJECT_ID \
    --location=LOCATION \
    --num-slices=WORKLOAD_NODEPOOL_COUNT \
    --compute-type=COMPUTE_TYPE \
    --topology=TOPOLOGY \
    --image="gcr.io/PROJECT_ID/USER_runner" \
    --command="python3 -m maxtext.trainers.pre_train.train /deps/src/maxtext/configs/base.yml base_output_directory=gs://BUCKET_NAME per_device_batch_size=1 enable_checkpointing=false remat_policy=full global_parameter_scale=1 steps=20 max_target_length=2048 use_iota_embed=true reuse_example_batch=1 dataset_type=synthetic attention=flash gcs_metrics=True enable_single_controller=True run_name=RUN_NAME-pathways-job"

    Ersetzen Sie Folgendes:

    • BUCKET_NAME: der Cloud Storage-Bucket für Ausgaben.
    • WORKLOAD: Ein eindeutiger Name zur Identifizierung Ihrer Arbeitslast (maximal 22 Zeichen aufgrund von DNS-Label-Limits).
    • CLUSTER: der Name Ihres GKE-Cluster.
    • PROJECT_ID: IhreGoogle Cloud -Projekt-ID.
    • LOCATION: die Region oder Zone Ihres GKE-Cluster.
    • WORKLOAD_NODEPOOL_COUNT: Die Anzahl der TPU-Slice-Knotenpools.
    • COMPUTE_TYPE: der TPU-Maschinentyp, z. B. ct6e-standard-4t. Weitere Informationen zu den unterstützten TPU-Typen für die einzelnen TPU-Versionen finden Sie unter TPU-Versionen.
    • TOPOLOGY: die TPU-Platzierungstopologie, z. B. 2x4.
    • USER: IhreGoogle Cloud -Nutzer-ID.
    • RUN_NAME: ein vom Nutzer zugewiesener Name zur Identifizierung der Workflowausführung.

    Weitere Informationen zu den Optionen für die Jobübermittlung finden Sie im Leitfaden zur Jobübermittlung mit dem Cluster Toolkit.

  3. Verfolgen Sie den Fortschritt Ihrer Arbeitslast mit dem Befehl gcluster job logs:

    gcluster job logs WORKLOAD \
    --cluster=CLUSTER \
    --project=PROJECT_ID \
    --location=LOCATION \
    --main-only \
    --follow

    Die Ausgabe sieht etwa so aus:

    completed step: 0, seconds: 1.259, TFLOP/s/device: 10.792, Tokens/s/device: 1626.636, total_weights: 65536, loss: 10.876, lm_loss: 10.876, perplexity: 52906.090
completed step: 1, seconds: 3.366, TFLOP/s/device: 4.036, Tokens/s/device: 608.366, total_weights: 65536, loss: 10.876, lm_loss: 10.876, perplexity: 52906.090
completed step: 2, seconds: 0.264, TFLOP/s/device: 51.415, Tokens/s/device: 7749.885, total_weights: 65536, loss: 10.563, lm_loss: 10.563, perplexity: 38665.316
completed step: 3, seconds: 1.025, TFLOP/s/device: 13.257, Tokens/s/device: 1998.187, total_weights: 65536, loss: 9.983, lm_loss: 9.983, perplexity: 21646.125
completed step: 4, seconds: 1.194, TFLOP/s/device: 11.376, Tokens/s/device: 1714.750, total_weights: 65536, loss: 9.462, lm_loss: 9.462, perplexity: 12864.218
completed step: 5, seconds: 1.185, TFLOP/s/device: 11.465, Tokens/s/device: 1728.203, total_weights: 65536, loss: 8.961, lm_loss: 8.961, perplexity: 7791.373

  4. Optional: Wenn Sie die Arbeitslast abbrechen möchten, bevor sie abgeschlossen ist, verwenden Sie den Befehl gcluster job cancel:

    gcluster job cancel WORKLOAD \
    --cluster=CLUSTER \
    --project=PROJECT_ID \
    --location=LOCATION

Cluster Director

Führen Sie die folgenden Befehle auf dem Login-Knoten des Clusters von ~/pathways aus. Die Skripts suchen nacheinander und schreiben ihre Logs in dieses Verzeichnis.

  1. Starten Sie den Pathways-Dienst:

    cd ~/pathways
SLICE_SHAPE=TOPOLOGY \
    MIN_SLICES=MIN_SLICES \
    MAX_SLICES=MAX_SLICES \
    GCS_BUCKET=gs://BUCKET_NAME \
    sbatch start_pathways.sh

    Ersetzen Sie Folgendes:

    • TOPOLOGY: die TPU-Topologie der einzelnen Slices. Sie muss mit der acceleratorTopology Ihres TPU-Knotensatzes übereinstimmen, z. B. 2x4.
    • MIN_SLICES: Die Mindestanzahl der laufenden Slices, die zum Starten und Fortsetzen des Trainings erforderlich sind.
    • MAX_SLICES: die maximale Anzahl von Slices, die verwendet werden sollen, normalerweise die Anzahl der Slices im Knotensatz.
    • BUCKET_NAME: der Cloud Storage-Bucket für Pathways-Scratch-Daten und -Prüfpunkte.

    Wenn Ihre TPU-Partition nicht die Standardpartition des Clusters ist, legen Sie für PARTITION auch den Namen fest.

    Der Dienst ist bereit, wenn das Log pathways-service-JOB_ID.out (wobei JOB_ID die Slurm-Job-ID ist) eine Ausgabe wie die folgende enthält:

    Pathways Runtime Service is READY!
 Discovery file    : /home/USER/pathways/cluster.env
 IFRT Proxy Target : grpc://RM_IP:29000
 Running TPU Slices : 2 / 2 (min=2)

  2. Senden Sie die MaxText-Trainingsarbeitslast. Sie können ihn direkt nach dem Dienstjob einreichen: Er wartet bis zu einer Stunde, bis der Dienst bereit ist.

    sbatch run_maxtext.sh

    Der MaxText-Job wartet, bis der Pathways-Dienst mindestens MIN_SLICES ausgeführte Slices hat, bevor er mit dem Training beginnt. Die Arbeitslast trainiert Llama 3.1 8B mit synthetischen Daten. Die Slice-Einstellungen und der Bucket stammen vom Dienst. Legen Sie sie daher nicht für diesen Job fest. Sie können die folgenden optionalen Variablen festlegen:

    • STEPS: die Anzahl der Trainingsschritte. Der Standardwert ist 30.
    • CHECKPOINT_PERIOD: die Anzahl der Schritte zwischen Prüfpunkten. Der Standardwert ist 5.
    • BASE_OUTPUT_DIR: der Cloud Storage-Pfad, in den MaxText Prüfpunkte schreibt. Der Standardwert ist gs://BUCKET_NAME/maxtext-runs.
    • REDUCED_LAYERS=1: Trainiert eine Variante mit 8 Schichten anstelle der vollständigen 32 Schichten für 2x2 Slices.

    Verfolgen Sie den Fortschritt Ihrer Arbeitslast (ersetzen Sie JOB_ID durch die Slurm-Job-ID des MaxText-Jobs). MaxText schreibt sein Log in den Standardfehlerkanal:

    tail -f maxtext-llama31-JOB_ID.err

    In jedem Trainingsschritt wird eine Zeile wie die folgende ausgegeben:

    completed step: 24, seconds: 0.326, live slice count: 2, TFLOP/s/device: 88.199, Tokens/s/device: 6289.27, total_weights: 32768, loss: 6.238

    Sie können run_maxtext.sh später noch einmal einreichen, ohne den Dienst neu starten zu müssen.

  3. Wenn Sie fertig sind, beenden Sie die Arbeitslast und dann den Pathways-Dienst:

    scancel -u ${USER} -n maxtext-llama31
scancel -u ${USER} -n pathways-service

    Wenn Sie den Dienst abbrechen, werden auch die Pathways-Worker-Jobs abgebrochen. Die TPU-VMs sind statische Knoten, sodass sie weiter ausgeführt werden und für andere Slurm-Jobs verfügbar sind.

Wenn Probleme auftreten, lesen Sie den Abschnitt Fehlerbehebung für Pathways in der Cloud.

Nächste Schritte