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 mendemonstrasikan hal ini.

Sebelum memulai

Pastikan Anda memiliki:

Membuat image pelatihan menggunakan Maxtext

MaxText adalah project model bahasa besar (LLM) open source yang dikembangkan oleh Google. Project ini ditulis dalam JAX dan dirancang agar memiliki performa tinggi dan dapat diskalakan, serta berjalan secara efisien di TPU dan GPU Google Cloud.

Untuk membuat 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
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 akan mengirimkan image Kubernetes MaxText ke gcr.io/$PROJECT/${USER}_runner. Anda dapat menggunakan image Docker ini untuk menjalankan pelatihan di TPU menggunakan backend Pathways.

Menjalankan workload batch menggunakan XPK

Sekarang Anda dapat mengirimkan image Docker Maxtext yang telah dibuat sebelumnya menggunakan XPK dengan perintah yang sama seperti yang Anda gunakan sebelumnya.

xpk workload create-pathways \
--workload=WORKLOAD \
--cluster=CLUSTER \
--num-slices=WORKLOAD_NODEPOOL_COUNT \
--tpu-type=TPU_TYPE \
--project=PROJECT \
--zone=ZONE \
--docker-image='gcr.io/PROJECT/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"

Ganti kode berikut:

  • WORKLOAD: nama unik untuk mengidentifikasi workload Anda
  • CLUSTER: nama cluster GKE Anda
  • WORKLOAD_NODEPOOL_COUNT : jumlah maksimum tugas dapat dimulai ulang
  • TPU_TYPE: jenis TPU menentukan versi dan ukuran Cloud TPU yang ingin Anda buat. Untuk mengetahui informasi selengkapnya tentang jenis TPU yang didukung untuk setiap versi TPU, lihat Versi TPU
  • PROJECT : Project ID Anda Google Cloud
  • ZONE: zona tempat Anda berencana menjalankan workload
  • USER : User ID Google Cloud Anda
  • RUN_NAME : nama yang ditetapkan pengguna untuk mengidentifikasi eksekusi alur kerja

Anda akan melihat output seperti berikut:

[XPK] Follow your Pathways workload and other resources here : https://console.cloud.google.com/logs/query;query=resource.type%3D"k8s_container"%0Aresource.labels.project_id%3D"<project-name>"%0Aresource.labels.location%3D"<your-zone>"%0Aresource.labels.cluster_name%3D"<your-cluster-name>"%0Aresource.labels.pod_name:"<your-pod-name>"%0Aseverity>%3DDEFAULT

Gunakan link dalam output dari perintah XPK sebelumnya untuk mengikuti progres workload Anda. Anda dapat memfilter log untuk container JAX dengan memilih jax-tpu di bagian filter Nama Container.

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

Workload akan selesai setelah jumlah langkah yang ditentukan. Jika Anda ingin menghentikannya sebelum waktunya, gunakan perintah berikut:

xpk workload delete --workload=WORKLOAD --cluster=CLUSTER --project=PROJECT --zone=ZONE

Langkah berikutnya