Google menggunakan teknologi AI untuk menerjemahkan konten ke dalam bahasa pilihan Anda. Terjemahan AI mungkin mengandung kesalahan.

Menjalankan supervised fine-tuning di VM TPU menggunakan MaxText

Tutorial ini memberikan panduan langkah demi langkah untuk menjalankan supervised fine-tuning (SFT) pada satu instance virtual machine (VM) v6e-8Tensor Processing Unit (TPU) di Google Cloud menggunakan MaxText, stack pelatihan berbasis JAX berperforma tinggi untuk model bahasa besar (LLM).

Ringkasan workload

Menjalankan workload SFT ini di VM Cloud TPU melibatkan komponen berikut:

Model: Gemma 3 4B (gemma3-4b), model bahasa ringan dengan bobot terbuka. Dalam tutorial ini, bobot model dikonversi dari format Hugging Face ke format yang kompatibel dengan MaxText (mt-format) dan disiapkan untuk dieksekusi di satu VM Cloud TPU v6e-8 menggunakan JAX.
Metode pelatihan: SFT menggunakan MaxText. Workload menjalankan pipeline fine-tuning pada satu instance Cloud TPU v6e-8 menggunakan JAX. Model ini mengoptimalkan parameter model selama 1.000 langkah menggunakan pelatih SFT bawaan MaxText (maxtext.trainers.post_train.sft.train_sft) dan set data penyesuaian instruksi yang berfokus pada dialog.
Set data dan template: UltraChat 200k (HuggingFaceH4/ultrachat_200k), set data dialog multi-turn berskala besar dan berkualitas tinggi (pemisahan train_sft dengan kolom messages). Data ini digunakan untuk melakukan fine-tuning model dan meningkatkan kemampuan percakapan dan mengikuti petunjuknya.
Infrastruktur dan framework: satu instance VM Cloud TPU v6e-8 (ct6e-standard-8t). Framework-nya adalah MaxText (maxtext[tpu-post-train]) yang didukung oleh JAX.

Tujuan

Siapkan instance VM Cloud TPU.
Instal MaxText dan dependensinya.
Mengonversi model Hugging Face ke format MaxText.
Menjalankan workload pelatihan SFT di TPU.
Mengonversi model yang telah disesuaikan kembali ke format Hugging Face untuk inferensi.

Biaya

Dalam dokumen ini, Anda akan menggunakan komponen Google Cloudyang dapat ditagih berikut:

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.

Sebelum memulai

Anda memerlukan token akses Hugging Face untuk menggunakan tutorial ini. Anda dapat mendaftar untuk mendapatkan akun gratis di Hugging Face. Setelah Anda 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 Create new token.
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, setujui perjanjian lisensi untuk model yang ingin Anda latih. Tutorial ini menggunakan model gemma3-4b.

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

TPU Admin (roles/tpu.admin)
Service Account User (roles/iam.serviceAccountUser)
Compute Editor (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 NAME="YOUR_TPU_NAME"
export NETWORK="default"

Ganti kode berikut:

YOUR_PROJECT_ID: Project ID Google Cloud Anda
YOUR_ZONE: zona yang ingin Anda gunakan
YOUR_RESERVATION_NAME: reservasi kapasitas Anda
YOUR_TPU_NAME: nama instance VM Cloud TPU Anda

Lakukan autentikasi dengan Google Cloud dengan menjalankan perintah berikut:

gcloud auth login

Membuat VM Cloud TPU

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

gcloud compute instances create "${NAME}" \
    --zone="${ZONE}" \
    --project="${PROJECT}" \
    --network="${NETWORK}" \
    --tags="${NAME}" \
    --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 "${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.2" --resolution=lowest

Instal dependensi lain yang diperlukan 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.

Tentukan variabel lingkungan Anda, seperti token akses Hugging Face, nama model yang ingin Anda gunakan, dan direktori tempat Anda ingin menyimpan model dalam format MaxText.

export HF_TOKEN="YOUR_HF_TOKEN"
export MODEL_NAME='gemma3-4b'
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 SFT.

Konfigurasikan parameter pelatihan workload SFT.

# -- 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 STEPS=1000
export PER_DEVICE_BATCH_SIZE=1

# -- Dataset configuration --
export DATASET_NAME="HuggingFaceH4/ultrachat_200k"
export TRAIN_SPLIT="train_sft"
export TRAIN_DATA_COLUMNS="['messages']"

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

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

python3 -m maxtext.trainers.post_train.sft.train_sft \
    run_name="${RUN_NAME?}" \
    base_output_directory="${BASE_OUTPUT_DIRECTORY?}" \
    model_name="${MODEL_NAME?}" \
    load_parameters_path="${MAXTEXT_CKPT_PATH?}" \
    per_device_batch_size="${PER_DEVICE_BATCH_SIZE?}" \
    steps="${STEPS?}" \
    hf_path="${DATASET_NAME?}" \
    train_split="${TRAIN_SPLIT?}" \
    train_data_columns="${TRAIN_DATA_COLUMNS?}" \
    profiler=xplane

Mengonversi kembali model terlatih ke format Hugging Face

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

Tetapkan jalur untuk ekspor dan parameter terlatih.

export HF_EXPORT=/dev/shm/$MODEL_NAME/hf-trained/
export POST_TRAIN_PATH=$BASE_OUTPUT_DIRECTORY/$RUN_NAME/checkpoints/$STEPS/model_params

Jalankan konversi kembali ke format Hugging Face.

python3 -m maxtext.checkpoint_conversion.to_huggingface \
    model_name="${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 telah disesuaikan dan disimpan di /dev/shm/gemma3-4b/hf-trained siap digunakan. Karena Anda akan kehilangan akses ke konten folder /dev/shm saat VM dimulai ulang, Anda harus memindahkan model yang telah disesuaikan ke penyimpanan persisten atau menguploadnya ke Hugging Face Hub.

Pembersihan

Agar tidak menimbulkan biaya tambahan, hapus resource yang dibuat selama tutorial ini.

Menghapus instance VM TPU

Keluar dari instance VM TPU Cloud Anda, lalu hapus instance tersebut.

# 1. Delete TPU instance
echo "Deleting TPU instance: ${NAME}..."
gcloud compute instances delete "${NAME}" \
    --zone="${ZONE}" \
    --project="${PROJECT}" \
    --quiet || true

# 2. Delete IAP firewall rule
echo "Deleting firewall rule: ${FIREWALL_RULE_NAME}..."
gcloud compute firewall-rules delete "${FIREWALL_RULE_NAME}" \
    --project="${PROJECT}" \
    --quiet || true

Langkah berikutnya

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