Menjalankan pelatihan reinforcement learning di VM TPU menggunakan MaxText

Tutorial ini memberikan panduan langkah demi langkah untuk menjalankan pelatihan reinforcement learning (RL) pada satu v6e-8 instance virtual machine (VM) Tensor Processing Unit (TPU) di Google Cloud dengan menggunakan MaxText, tumpukan pelatihan berbasis JAX berperforma tinggi untuk model bahasa besar (LLM).

Ringkasan workload

Menjalankan workload pasca-pelatihan RL ini di Cloud TPU melibatkan komponen berikut:

• Model: Llama 3.1 8B Instruct (llama3.1-8b-Instruct), LLM open-weights yang dioptimalkan untuk tugas dialog dan mengikuti petunjuk. Dalam tutorial ini, bobot model dikonversi dari format Hugging Face ke format yang kompatibel dengan MaxText (mt-format) dan disiapkan untuk eksekusi terdistribusi di TPU v6e-8 menggunakan JAX.

• Algoritma Reinforcement Learning (RL): Group Relative Policy Optimization (GRPO), yang workload-nya menjalankan pipeline pasca-pelatihan yang skalabel menggunakan pelatih RL bawaan MaxText. GRPO menghilangkan kebutuhan akan model nilai atau kritik terpisah dengan mengevaluasi grup output yang dihasilkan untuk perintah yang sama, menggunakan reward relatif untuk mengoptimalkan kebijakan aktor secara efisien.

• Set data dan template: Grade School Math 8K (GSM8K), set data berkualitas tinggi yang terdiri dari soal matematika sekolah dasar yang beragam secara linguistik. Workload ini menggunakannya untuk membuat rollout, menghitung reward, dan meningkatkan kemampuan penalaran multi-langkah selama fase penyelarasan RL.

• Infrastruktur dan framework: Satu node Cloud TPU v6e-8 (ct6e-standard-8t) menggunakan MaxText (maxtext[tpu-post-train]), tumpukan pelatihan berperforma tinggi yang didukung oleh JAX.

Tujuan

• Menyiapkan instance VM Cloud TPU.
• Menginstal MaxText dan dependensinya.
• Mengonversi model Hugging Face ke format MaxText.
• Menjalankan workload RL Group Relative Policy Optimization (GRPO) di TPU.
• Mengonversi model terlatih kembali ke format Hugging Face untuk penayangan.

Biaya

Dalam dokumen ini, Anda akan menggunakan komponen yang dapat ditagih sebagai berikut Google Cloud:

• Cloud TPU

Untuk membuat perkiraan biaya berdasarkan proyeksi penggunaan Anda, gunakan kalkulator harga.

Pengguna Google Cloud baru mungkin memenuhi syarat untuk mendapatkan uji coba gratis.

Setelah menyelesaikan tugas yang dijelaskan dalam dokumen ini, Anda dapat menghindari penagihan berkelanjutan dengan menghapus resource yang Anda buat. Untuk mengetahui informasi selengkapnya, lihat Pembersihan.

Penting: Untuk menyelesaikan tutorial ini, Anda harus memiliki akses ke TPU v6e-8. Tutorial ini mengasumsikan bahwa Anda sudah memiliki pemesanan untuk TPU ini di Google Cloud project Anda guna memastikan Anda memiliki kapasitas yang tersedia saat Anda membutuhkannya. Untuk mempelajari lebih lanjut opsi Anda untuk memesan kapasitas Cloud TPU untuk tanggal dan waktu mendatang, lihat Tentang pemesanan Cloud TPU.

Sebelum memulai

• Anda memerlukan token akses Hugging Face untuk menggunakan tutorial ini. Anda dapat mendaftar akun gratis di Hugging Face. Setelah memiliki akun, buat token akses:

  1. Di halaman Welcome to Hugging Face, klik avatar akun Anda, lalu pilih Token akses.
  2. Di halaman Token akses, klik Buat token baru.
  3. Pilih jenis token Read dan masukkan nama untuk token Anda.
  4. Token akses Anda akan ditampilkan. Simpan token di tempat yang aman.

• Di situs Hugging Face, terima perjanjian lisensi untuk model yang Anda rencanakan untuk dilatih. Tutorial ini menggunakan model llama3.1-8b-Instruct.

Untuk mendapatkan izin yang Anda perlukan untuk menyelesaikan tutorial ini, minta administrator untuk memberi Anda peran IAM berikut di project Anda:

• Admin TPU (roles/tpu.admin)
• Pengguna Akun Layanan (roles/iam.serviceAccountUser)
• Editor Compute (roles/compute.editor)

Untuk mengetahui informasi selengkapnya tentang pemberian peran, lihat Mengelola akses ke project, folder, dan organisasi.

Anda mungkin juga bisa mendapatkan izin yang diperlukan melalui peran khusus atau peran bawaan lainnya.

Menyiapkan lingkungan

Siapkan variabel lingkungan Anda dengan menjalankan skrip berikut:

export PROJECT="YOUR_PROJECT_ID"
export ZONE="YOUR_ZONE"
export RESERVATION="YOUR_RESERVATION_NAME"
export TPU_NAME="YOUR_TPU_NAME"

Ganti kode berikut:

• YOUR_PROJECT_ID: ID project Anda Google Cloud
• YOUR_ZONE: zona yang ingin Anda gunakan
• YOUR_RESERVATION_NAME: pemesanan kapasitas Anda
• YOUR_TPU_NAME: nama instance VM Cloud TPU Anda

Lakukan autentikasi dengan Google Cloud menjalankan perintah berikut:

gcloud auth login

Membuat VM Cloud TPU

Buat instance VM Cloud TPU dengan 8 chip TPU v6e, yang terikat ke pemesanan kapasitas Anda.

gcloud compute instances create "${TPU_NAME}" \
    --zone="${ZONE}" \
    --project="${PROJECT}" \
    --machine-type=ct6e-standard-8t \
    --image-project=ubuntu-os-accelerator-images \
    --image-family=ubuntu-accel-2204-amd64-tpu-v5e-v5p-v6e \
    --boot-disk-size=200GB \
    --maintenance-policy=TERMINATE \
    --instance-termination-action=DELETE \
    --provisioning-model=RESERVATION_BOUND \
    --reservation-affinity=specific \
    --reservation="${RESERVATION}"

Setelah instance VM dibuat, hubungkan ke instance tersebut menggunakan SSH.

gcloud compute ssh "${TPU_NAME}" --zone "${ZONE}" --project "${PROJECT}"

Selesaikan langkah-langkah berikut dalam instance VM TPU Anda.

Menginstal MaxText

Update paket sistem dalam instance VM TPU.

sudo apt update && sudo apt upgrade -y --fix-missing

Instal Python 3.12, yang diperlukan MaxText, dan paket lingkungan virtualnya.

sudo apt install -y build-essential cmake ninja-build

Gunakan uv untuk mempercepat penginstalan paket Python.

curl -LsSf https://astral.sh/uv/install.sh | sh
source "${HOME}/.local/bin/env"

Buat lingkungan virtual bernama maxtext_venv dan aktifkan.

uv python install 3.12
uv venv --python 3.12 --seed maxtext_venv
source maxtext_venv/bin/activate

Instal MaxText dan dependensi yang diperlukan untuk tugas pasca-pelatihan.

UV_TORCH_BACKEND=cpu uv pip install "maxtext[tpu-post-train]==0.2.3" --resolution=lowest

Instal dependensi yang diperlukan yang tersisa dengan menjalankan perintah berikut:

install_tpu_post_train_extra_deps

Mengonversi model ke format MaxText

Untuk melatih model dalam format MaxText, Anda harus mengonversinya dari format Hugging Face ke format MaxText.

Berikan nilai berikut:

• Token akses Hugging Face Anda
• Nama model yang ingin Anda gunakan
• Direktori tempat Anda ingin menyimpan model dalam format MaxText
• Opsi untuk pemuatan dan penyimpanan

export HF_TOKEN="YOUR_HF_TOKEN"
export MODEL_NAME='llama3.1-8b-Instruct'
export MODEL_CHECKPOINT_DIRECTORY=/dev/shm/$MODEL_NAME/mt-format/
export USE_PATHWAYS=0 # Set to 1 for Pathways, 0 for McJAX
export LAZY_LOAD_TENSORS=False # True to use lazy load, False to use eager load.

Ganti YOUR_HF_TOKEN dengan token akses Hugging Face yang Anda buat sebelumnya.

Untuk mengonversi model dari format Hugging Face ke format MaxText, jalankan skrip berikut. Konversi ini memerlukan waktu sekitar lima menit untuk diselesaikan.

python3 -m maxtext.checkpoint_conversion.to_maxtext \
    model_name="${MODEL_NAME?}" \
    hf_access_token="${HF_TOKEN?}" \
    base_output_directory="${MODEL_CHECKPOINT_DIRECTORY?}" \
    scan_layers=True \
    use_multimodal=False \
    hardware=cpu \
    skip_jax_distributed_system=true \
    checkpoint_storage_use_zarr3=$((1 - USE_PATHWAYS)) \
    checkpoint_storage_use_ocdbt=$((1 - USE_PATHWAYS)) \
    --lazy_load_tensors="${LAZY_LOAD_TENSORS?}"

Memulai workload pelatihan

Setelah proses konversi selesai, Anda dapat memulai workload RL.

1. Konfigurasi parameter pelatihan workload RL.

   # -- MaxText configuration --
   export BASE_OUTPUT_DIRECTORY=/dev/shm/$MODEL_NAME/post-train/
   RUN_NAME=$(date +%Y-%m-%d-%H-%M-%S)
   export RUN_NAME
   export CHIPS_PER_VM=8
   export NUM_BATCHES=50
   export MAXTEXT_CKPT_PATH=$MODEL_CHECKPOINT_DIRECTORY/0/items
   export TPU_ACCELERATOR_TYPE=v6e-8
   export TPU_WORKER_ID=0
   TPU_NAME=$(hostname)
   export TPU_NAME
   export TPU_SKIP_MDS_QUERY=1
   export TPU_WORKER_HOSTNAMES=localhost
   export TPU_TOPOLOGY=2x4
   export TPU_CHIPS_PER_HOST_BOUNDS=2,4,1
   export TPU_HOST_BOUNDS=1,1,1

2. Mulai tugas pelatihan. Proses ini memerlukan waktu sekitar 10 menit di instance VM v6e-8.

   python3 -m maxtext.trainers.post_train.rl.train_rl \
       model_name="${MODEL_NAME?}" \
       load_parameters_path="${MAXTEXT_CKPT_PATH?}" \
       run_name="${RUN_NAME?}" \
       base_output_directory="${BASE_OUTPUT_DIRECTORY?}" \
       chips_per_vm="${CHIPS_PER_VM?}" \
       num_batches="${NUM_BATCHES?}" \
       num_test_batches=10 \
       rollout_data_parallelism=1 \
       rollout_tensor_parallelism=-1

Mengonversi model terlatih kembali ke format Hugging Face

Setelah workload pelatihan selesai, konversi model kembali ke format Hugging Face.

1. Tetapkan jalur untuk ekspor dan parameter terlatih.

   export HF_EXPORT=/dev/shm/$MODEL_NAME/hf-trained/
   export HF_MODEL_NAME=llama3.1-8b
   export POST_TRAIN_PATH=$BASE_OUTPUT_DIRECTORY/$RUN_NAME/checkpoints/actor/$NUM_BATCHES/model_params

2. Jalankan konversi kembali ke format Hugging Face.

   python3 -m maxtext.checkpoint_conversion.to_huggingface \
       model_name="${HF_MODEL_NAME?}" \
       load_parameters_path="${POST_TRAIN_PATH?}" \
       base_output_directory="${HF_EXPORT?}" \
       scan_layers=True \
       use_multimodal=False \
       weight_dtype=bfloat16

Setelah konversi selesai, model yang disesuaikan yang disimpan di /dev/shm/$MODEL_NAME/hf-trained siap digunakan. Karena Anda akan kehilangan akses ke konten folder /dev/shm saat VM di-reboot, Anda harus memindahkan model yang disesuaikan ke penyimpanan persisten atau menguploadnya ke Hugging Face Hub.

Pembersihan

Untuk menghindari biaya tambahan, hapus resource yang dibuat selama tutorial ini.

Menghapus instance VM TPU

Hapus instance VM Cloud TPU Anda.

gcloud compute instances delete "${TPU_NAME}" --zone="${ZONE}" --project="${PROJECT}" --quiet

Langkah berikutnya

• Untuk mengetahui informasi selengkapnya tentang Cloud TPU, lihat Pengantar Cloud TPU.
• Untuk mengetahui detail arsitektur dan konfigurasi untuk TPU v6e-8, lihat TPU v6e.