Batcharbeitslast mit Pathways ausführen

Für die Zwecke dieses Dokuments werden Batcharbeitslasten als JAX-Arbeitslasten definiert, die bis zum Abschluss ausgeführt werden und im selben GKE-Cluster wie der Pathways-Cluster bereitgestellt werden, insbesondere neben den Pathways-Controllerkomponenten (IFRT-Proxyserver und Pathways-Ressourcenmanager). Nach Abschluss der JAX-Arbeitslast werden die Pathways-Clusterkomponenten beendet. In diesem Leitfaden wird dies anhand einer JAX-Trainingsarbeitslast veranschaulicht.

Hinweis

Sie benötigen Folgendes:

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 so konzipiert, dass es sehr leistungsstark und skalierbar ist und effizient auf Google Cloud-TPUs und ‑GPUs ausgeführt werden kann.

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

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

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.

Batcharbeitslast mit dem Cluster Toolkit ausführen

Senden Sie das vorgefertigte 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=ZONE \
    --num-slices=WORKLOAD_NODEPOOL_COUNT \
    --compute-type=COMPUTE_TYPE \
    --topology=TOPOLOGY \
    --image="gcr.io/PROJECT_ID/USER_runner" \
    --command="python3 -m MaxText.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"

Weitere Informationen zu Optionen zum Einreichen von Jobs finden Sie im Cluster Toolkit Job Submission Guide.

Ersetzen Sie Folgendes:

  • WORKLOAD: Ein eindeutiger Name zur Identifizierung Ihres Workloads. Aufgrund von DNS-Label-Beschränkungen darf er maximal 22 Zeichen lang sein.
  • CLUSTER: der Name 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)
  • PROJECT_ID: Ihre Google Cloud Projekt-ID
  • ZONE: Die Zone, in der Sie Ihre Arbeitslast ausführen möchten
  • USER: Ihre Google Cloud Nutzer-ID
  • BUCKET_NAME: der Cloud Storage-Bucket für Ausgaben
  • RUN_NAME: Ein vom Nutzer zugewiesener Name zur Identifizierung des Workflow-Laufs.

Mit dem Befehl gcluster job logs können Sie den Fortschritt Ihrer Arbeitslast verfolgen:

gcluster job logs WORKLOAD \
    --cluster=CLUSTER \
    --project=PROJECT_ID \
    --location=ZONE \
    --main-only=false

completed step: 1, seconds: 0.484, TFLOP/s/device: 87.349, Tokens/s/device: 2117.382, total_weights: 2945, loss: 10.888
completed step: 2, seconds: 0.407, TFLOP/s/device: 103.699, Tokens/s/device: 2513.735, total_weights: 3253, loss: 9.697
completed step: 3, seconds: 0.248, TFLOP/s/device: 170.300, Tokens/s/device: 4128.167, total_weights: 3154, loss: 9.641
completed step: 4, seconds: 0.216, TFLOP/s/device: 195.122, Tokens/s/device: 4729.880, total_weights: 3119, loss: 9.547
completed step: 5, seconds: 0.272, TFLOP/s/device: 155.298, Tokens/s/device: 3764.512, total_weights: 2837, loss: 10.179
completed step: 6, seconds: 0.472, TFLOP/s/device: 89.489, Tokens/s/device: 2169.266, total_weights: 3069, loss: 9.776

Verwenden Sie den gcluster job cancel-Befehl, um die Arbeitslast vor Abschluss abzubrechen:

gcluster job cancel WORKLOAD --cluster=CLUSTER --project=PROJECT_ID --location=ZONE

Nächste Schritte