Menjalankan workload batch dengan Pathways

Untuk tujuan dokumen ini, workload batch ditentukan sebagai workload JAX yang dieksekusi hingga selesai dan di-deploy dalam cluster GKE yang sama dengan cluster Pathways, khususnya bersama komponen pengontrol Pathways (server proxy IFRT dan pengelola resource Pathways). Penyelesaian workload JAX akan menghentikan komponen cluster Pathways. Panduan ini menggunakan workload pelatihan JAX untuk mendemonstrasikannya.

Sebelum memulai

Pastikan Anda memiliki:

Membangun image pelatihan menggunakan Maxtext

MaxText adalah project model bahasa besar (LLM) open source yang dikembangkan oleh Google. Ditulis dalam JAX dan dirancang agar berperforma tinggi dan skalabel, serta berjalan secara efisien di TPU dan GPU Google Cloud.

Untuk membangun image Docker MaxText menggunakan JAX stabil versi terbaru dari repositori GitHub OSS, jalankan perintah berikut:

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.

Perintah ini mengirim image Kubernetes MaxText ke gcr.io/$PROJECT_ID/${USER}_runner. Anda dapat menggunakan image Docker ini untuk menjalankan pelatihan di TPU menggunakan backend Pathways.

Menjalankan workload batch dengan Cluster Toolkit

Kirimkan image Docker MaxText bawaan menggunakan perintah 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"

Untuk mengetahui informasi selengkapnya tentang opsi pengiriman tugas, lihat Panduan Pengiriman Tugas Cluster Toolkit.

Ganti kode berikut:

  • WORKLOAD: nama unik untuk mengidentifikasi workload Anda; karena batas label DNS, nama ini harus terdiri dari 22 karakter atau kurang
  • CLUSTER: nama cluster GKE Anda
  • WORKLOAD_NODEPOOL_COUNT: jumlah node pool slice TPU
  • COMPUTE_TYPE: jenis mesin TPU (misalnya, ct6e-standard-4t). Untuk mengetahui informasi selengkapnya tentang jenis TPU yang didukung untuk setiap versi TPU, lihat versi TPU.
  • TOPOLOGY: topologi penempatan TPU (misalnya, 2x4)
  • PROJECT_ID: Project ID Google Cloud Anda
  • ZONE: zona tempat Anda berencana menjalankan workload
  • USER: ID Google Cloud pengguna Anda
  • BUCKET_NAME: bucket Cloud Storage untuk output
  • RUN_NAME: nama yang ditetapkan pengguna untuk mengidentifikasi eksekusi alur kerja

Pantau progres beban kerja Anda menggunakan perintah gcluster job logs:

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

Untuk membatalkan workload sebelum selesai, gunakan perintah gcluster job cancel:

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

Langkah berikutnya