Menjalankan workload batch dengan Pathways

Panduan ini menggunakan workload pelatihan JAX yang berjalan hingga selesai di cluster yang sama dengan komponen pengontrol Pathways untuk menunjukkan cara menjalankan workload batch.

Anda dapat menjalankan workload batch di cluster Google Kubernetes Engine (GKE) atau di cluster Cluster Director dengan Slurm:

  • Di GKE, penyelesaian workload JAX akan menghentikan komponen cluster Pathways.
  • Di Cluster Director, komponen pengontrol Pathways berjalan sebagai tugas Slurm terpisah, sehingga terus berjalan setelah workload JAX selesai, hingga Anda menghentikannya.

Halaman ini berisi petunjuk untuk kedua platform. Di setiap bagian, pilih tab GKE atau Cluster Director.

Sebelum memulai

Sebelum memulai

GKE

Selesaikan prasyarat berikut:

Cluster Director

Selesaikan prasyarat berikut:

Membangun image pelatihan menggunakan MaxText

MaxText adalah project model bahasa besar (LLM) open source yang dikembangkan oleh Google. Model ini ditulis dalam JAX dan dirancang untuk berjalan secara efisien di TPU dan GPU Google Cloud .

GKE

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

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

Ganti kode berikut:

  • PROJECT_ID:Google Cloud project ID Anda.
  • USER: ID pengguna Google Cloud Anda.

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.

Cluster Director

Anda tidak perlu membangun image pelatihan: setup_env.sh menginstal MaxText di lingkungan Pathways di /home/pathways, dan run_maxtext.sh menjalankannya.

Menjalankan workload batch

GKE

  1. Jika biner gcluster tidak ada dalam variabel lingkungan PATH, tambahkan direktori yang berisi gcluster ke PATH:

    export PATH="$PATH:CLUSTER_TOOLKIT_PATH"

    Ganti CLUSTER_TOOLKIT_PATH dengan jalur absolut ke direktori yang berisi biner gcluster, seperti $HOME/cluster-toolkit.

  2. Kirim 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=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"

    Ganti kode berikut:

    • BUCKET_NAME: bucket Cloud Storage untuk output.
    • WORKLOAD: nama unik untuk mengidentifikasi workload Anda (22 karakter atau kurang karena batas label DNS).
    • CLUSTER: nama cluster GKE Anda.
    • PROJECT_ID:Google Cloud project ID Anda.
    • LOCATION: region atau zona 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.
    • USER: ID pengguna Google Cloud Anda.
    • RUN_NAME: nama yang ditetapkan pengguna untuk mengidentifikasi operasi alur kerja.

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

  3. Ikuti progres workload Anda menggunakan perintah gcluster job logs:

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

    Outputnya mirip dengan hal berikut ini:

    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. Opsional: Untuk membatalkan workload sebelum selesai, gunakan perintah gcluster job cancel:

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

Cluster Director

Jalankan perintah berikut pada node login cluster dari ~/pathways. Skrip akan saling mencari dan menulis log-nya di direktori tersebut.

  1. Mulai layanan Pathways:

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

    Ganti kode berikut:

    • TOPOLOGY: topologi TPU dari setiap slice. Versi ini harus cocok dengan acceleratorTopology nodeset TPU Anda—misalnya, 2x4.
    • MIN_SLICES: jumlah minimum slice yang berjalan yang diperlukan untuk memulai dan melanjutkan pelatihan.
    • MAX_SLICES: jumlah maksimum slice yang akan digunakan, biasanya jumlah slice dalam nodeset.
    • BUCKET_NAME: bucket Cloud Storage untuk data sementara dan checkpoint Pathways.

    Jika partisi TPU Anda bukan partisi default cluster, tetapkan juga PARTITION ke namanya.

    Layanan siap saat log-nya, pathways-service-JOB_ID.out (dengan JOB_ID adalah ID pekerjaan Slurm), menampilkan output yang mirip dengan berikut ini:

    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. Kirim workload pelatihan MaxText. Anda dapat mengirimkannya tepat setelah tugas layanan: tugas ini menunggu hingga satu jam agar layanan siap.

    sbatch run_maxtext.sh

    Tugas MaxText menunggu hingga layanan Pathways memiliki setidaknya MIN_SLICES slice yang berjalan sebelum memulai pelatihan. Workload ini melatih Llama 3.1 8B pada data sintetis. Setelan slice dan bucket berasal dari layanan, jadi jangan tetapkan di tugas ini. Anda dapat menetapkan variabel opsional berikut:

    • STEPS: jumlah langkah pelatihan. Defaultnya adalah 30.
    • CHECKPOINT_PERIOD: jumlah langkah di antara checkpoint. Defaultnya adalah 5.
    • BASE_OUTPUT_DIR: jalur Cloud Storage tempat MaxText menulis checkpoint. Defaultnya adalah gs://BUCKET_NAME/maxtext-runs.
    • REDUCED_LAYERS=1: melatih varian 8 lapisan, bukan 32 lapisan penuh, untuk slice 2x2.

    Ikuti progres workload Anda (ganti JOB_ID dengan ID pekerjaan Slurm dari tugas MaxText). MaxText menulis log-nya ke error standar:

    tail -f maxtext-llama31-JOB_ID.err

    Setiap langkah pelatihan akan mencetak baris yang mirip dengan berikut ini:

    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

    Anda dapat mengirimkan run_maxtext.sh lagi nanti tanpa memulai ulang layanan.

  3. Setelah selesai, hentikan workload, lalu layanan Pathways:

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

    Membatalkan layanan juga akan membatalkan tugas worker Pathways. VM TPU adalah node statis, sehingga terus berjalan dan tersedia untuk tugas Slurm lainnya.

Jika Anda mengalami masalah, lihat Memecahkan masalah Pathways di Cloud.

Langkah berikutnya