Melatih model menggunakan TPU v6e

Dokumen ini memandu Anda dalam melatih model di Cloud TPU v6e (juga disebut Trillium), yang mencakup penyiapan lingkungan, pengoptimalan performa, dan contoh pelatihan praktis menggunakan JAX dan PyTorch/XLA.

TPU v6e, yang juga disebut Trillium, adalah TPU generasi ke-6 Google. Di semua platform teknis, seperti API dan log, serta di seluruh dokumen ini, Trillium akan disebut sebagai v6e. Dengan 256 chip per Pod, arsitektur TPU v6e memiliki banyak kesamaan dengan v5e. TPU v6e dioptimalkan untuk pelatihan, penyesuaian, dan penayangan model transformer, text-to-image, dan jaringan neural konvolusional (CNN). Untuk mengetahui informasi selengkapnya tentang arsitektur dan konfigurasi sistem TPU v6e, lihat TPU v6e.

Untuk mengetahui informasi tentang menjalankan inferensi di Cloud TPU v6e, lihat tutorial berikut:

Sebelum memulai

Sebelum memulai, Anda perlu:

Buat akun Google Cloud dan project dengan penagihan diaktifkan
Menginstal komponen alfa Google Cloud CLI
Aktifkan Cloud TPU API
Membuat agen layanan Cloud TPU
Buat akun layanan Cloud TPU dan berikan izin

Untuk mengetahui informasi selengkapnya, lihat Menyiapkan lingkungan Cloud TPU.

Memverifikasi kuota dan izin

Pastikan project Anda memiliki kuota berikut:

Kuota TPU v6e preemptible atau on-demand
Kuota alamat IP
Kuota untuk Hyperdisk Balanced dan untuk jenis disk lain yang ingin Anda gunakan

Catatan: v6e mendukung Hyperdisk Balanced dan Hyperdisk ML. v6e tidak mendukung Persistent Disk (PD). Untuk mengetahui informasi selengkapnya, lihat Opsi penyimpanan untuk data Cloud TPU.

Jika Anda menggunakan Google Kubernetes Engine (GKE) dengan XPK (Accelerated Processing Kit), Anda memerlukan izin tambahan di konsol Google Cloud . Untuk mengetahui informasi selengkapnya, lihat Izin yang diperlukan di konsol Google Cloud .

Opsi penyediaan

Anda dapat menyediakan dan mengelola TPU v6e menggunakan metode berikut:

GKE: Anda dapat menggunakan GKE untuk menyediakan dan mengelola TPU sebagai kumpulan akselerator untuk workload machine learning dalam container Anda. Untuk mengetahui informasi selengkapnya, lihat Tentang TPU di GKE.
GKE dan XPK: XPK adalah alat command line yang menyederhanakan pembuatan cluster dan eksekusi workload di GKE. Mode ini dirancang bagi praktisi ML untuk menyediakan TPU dan menjalankan tugas pelatihan tanpa memerlukan keahlian Kubernetes yang mendalam. Untuk mengetahui informasi selengkapnya, lihat repositori GitHub XPK.
Resource dalam antrean Cloud TPU: Resource dalam antrean memungkinkan Anda meminta kapasitas TPU yang disediakan saat tersedia. Sangat ideal untuk tugas batch dan workload fault-tolerant yang dapat menunggu dalam antrean. Anda dapat menentukan jangka waktu untuk permintaan Anda. Untuk mengetahui informasi selengkapnya, lihat Mengelola resource dalam antrean.

Menyediakan TPU v6e dengan GKE dan XPK

Jika menggunakan GKE dengan TPU v6e, Anda dapat menggunakan perintah Kubernetes atau XPK untuk menyediakan TPU dan melatih atau menayangkan model. Untuk mengetahui informasi selengkapnya tentang penggunaan GKE dengan TPU, lihat Tentang TPU di GKE.

Cloud TPU v6e mendukung konfigurasi kartu antarmuka jaringan (NIC) yang memungkinkan Anda menskalakan throughput di beberapa jaringan. Bagian berikut memberikan perintah untuk membuat cluster GKE dengan dukungan NIC tunggal atau dukungan multi-NIC menggunakan XPK. Untuk sebagian besar workload slice tunggal, NIC tunggal memberikan performa yang memadai dengan lebih sedikit konfigurasi. Untuk workload Multislice dan workload yang memerlukan kecepatan penyerapan data yang tinggi, gunakan multi-NIC.

Membuat cluster dengan dukungan NIC tunggal menggunakan XPK

Untuk sebagian besar workload slice tunggal, NIC tunggal memberikan performa yang memadai dengan konfigurasi yang lebih sedikit. Untuk workload Multislice dan workload yang memerlukan kecepatan penyerapan data yang tinggi, gunakan multi-NIC.

Bagian berikut menunjukkan cara membuat cluster GKE dengan dukungan NIC tunggal menggunakan XPK.

Instal XPK dan siapkan variabel lingkungan

Instal XPK. Ikuti petunjuk di repositori GitHub XPK.
Siapkan variabel lingkungan untuk cluster Anda:
```
export CLUSTER_NAME=XPK_CLUSTER_NAME
export ZONE=us-east1-d
export PROJECT_ID=PROJECT_ID
export ACCELERATOR_TYPE=ACCELERATOR_TYPE
export NUM_SLICES=1
```
Tetapkan variabel lingkungan berikut:
- CLUSTER_NAME: Nama untuk cluster Anda.
- ZONE: Zona tempat cluster TPU akan dibuat. Untuk mengetahui informasi selengkapnya tentang zona yang didukung, lihat Region dan zona.
- PROJECT_ID: Project ID Google Cloud Anda.
- ACCELERATOR_TYPE: Jenis TPU, yang juga disebut jenis akselerator, menentukan versi dan ukuran Cloud TPU yang ingin Anda buat. Misalnya, v6e-256. Untuk mengetahui informasi selengkapnya tentang jenis akselerator yang didukung untuk setiap versi TPU, lihat versi TPU.
- NUM_SLICES: Jumlah slice TPU untuk cluster Anda. Setiap slice memiliki jumlah chip yang ditentukan dalam ACCELERATOR_TYPE. Untuk cluster slice tunggal, tetapkan NUM_SLICES ke 1. Untuk cluster Multislice, tentukan jumlah slice berdasarkan persyaratan skalabilitas beban kerja Anda. Jumlah total chip dalam cluster adalah jumlah chip dalam ACCELERATOR_TYPE dikalikan dengan NUM_SLICES.

Membuat cluster

Pilih salah satu opsi berikut untuk membuat cluster Anda. Sebaiknya gunakan jaringan kustom dengan MTU 8.896 untuk performa yang optimal. Untuk mengetahui informasi selengkapnya, lihat Mengonfigurasi MTU.

Jaringan kustom

Untuk membuat jaringan kustom dengan MTU 8.896 dan menggunakannya untuk cluster Anda, ikuti langkah-langkah berikut:

Tetapkan variabel lingkungan untuk nama jaringan dan firewall:
```
export NETWORK_NAME=NETWORK_NAME
export NETWORK_FW_NAME=FIREWALL_NAME
```
Ganti kode berikut:
- NETWORK_NAME: Nama untuk jaringan.
- FIREWALL_NAME: Nama untuk aturan firewall jaringan.

Buat jaringan kustom dengan MTU 8.896:

gcloud compute networks create ${NETWORK_NAME} \
    --mtu=8896 \
    --project=${PROJECT_ID} \
    --subnet-mode=auto \
    --bgp-routing-mode=regional

Buat aturan firewall yang mengizinkan traffic TCP, ICMP, dan UDP di jaringan Anda:

gcloud compute firewall-rules create ${NETWORK_FW_NAME} \
    --network=${NETWORK_NAME} \
    --allow tcp,icmp,udp \
    --project=${PROJECT_ID}

Tetapkan variabel lingkungan untuk argumen cluster XPK guna menggunakan jaringan yang Anda buat:
```
export CLUSTER_ARGUMENTS="--network=${NETWORK_NAME} --subnetwork=${NETWORK_NAME}"
```

Buat cluster XPK. Perintah berikut menyediakan kapasitas sesuai permintaan:

xpk cluster create --cluster=${CLUSTER_NAME} \
    --cluster-cpu-machine-type=e2-standard-8 \
    --num-slices=${NUM_SLICES} \
    --tpu-type=${ACCELERATOR_TYPE} \
    --zone=${ZONE} \
    --project=${PROJECT_ID} \
    --on-demand \
    --custom-cluster-arguments="${CLUSTER_ARGUMENTS}"

Untuk menggunakan kapasitas yang dicadangkan, ganti --on-demand dengan --reservation=RESERVATION_NAME. Untuk menggunakan VM Spot TPU, ganti --on-demand dengan --spot.

Jaringan default

Jika tidak memerlukan jaringan MTU tinggi, Anda dapat membuat cluster yang menggunakan jaringan VPC default. Perintah berikut menyediakan kapasitas sesuai permintaan:

xpk cluster create --cluster=${CLUSTER_NAME} \
    --cluster-cpu-machine-type=e2-standard-8 \
    --num-slices=${NUM_SLICES} \
    --tpu-type=${ACCELERATOR_TYPE} \
    --zone=${ZONE} \
    --project=${PROJECT_ID} \
    --on-demand

Untuk menggunakan kapasitas yang dicadangkan, ganti --on-demand dengan --reservation=RESERVATION_NAME. Untuk menggunakan VM Spot TPU, ganti --on-demand dengan --spot.

Membuat cluster dengan dukungan multi-NIC menggunakan XPK

Untuk beban kerja Multislice atau beban kerja lain yang memerlukan bandwidth jaringan tinggi, seperti untuk penyerapan data, Anda dapat menggunakan multi-NIC untuk meningkatkan performa. Saat Anda menggunakan multi-NIC, setiap VM TPU dialokasikan antarmuka jaringan tambahan, yang masing-masing terhubung ke jaringan VPC yang unik, sehingga meningkatkan throughput jaringan secara keseluruhan. Untuk sebagian besar beban kerja slice tunggal, NIC tunggal memberikan performa yang memadai dengan konfigurasi yang lebih sedikit.

Bagian berikut menunjukkan cara membuat cluster GKE dengan dukungan multi-NIC menggunakan XPK.

Instal XPK dan siapkan variabel lingkungan

Instal XPK. Ikuti petunjuk di repositori GitHub XPK.
Siapkan variabel lingkungan untuk cluster dan jaringan utama Anda:
```
export CLUSTER_NAME=XPK_CLUSTER_NAME
export REGION=REGION
export ZONE=us-east1-d
export PROJECT_ID=PROJECT_ID
export ACCELERATOR_TYPE=ACCELERATOR_TYPE
export NUM_SLICES=2

export NETWORK_NAME_1=${CLUSTER_NAME}-mtu9k-1-${ZONE}
export SUBNET_NAME_1=${CLUSTER_NAME}-privatesubnet-1-${ZONE}
export FIREWALL_RULE_NAME_1=${CLUSTER_NAME}-privatefirewall-1-${ZONE}
export ROUTER_NAME_1=${CLUSTER_NAME}-network-1-${ZONE}
export NAT_CONFIG_1=${CLUSTER_NAME}-natconfig-1-${ZONE}

export NETWORK_NAME_2=${CLUSTER_NAME}-mtu9k-2-${ZONE}
export SUBNET_NAME_2=${CLUSTER_NAME}-privatesubnet-2-${ZONE}
export FIREWALL_RULE_NAME_2=${CLUSTER_NAME}-privatefirewall-2-${ZONE}
export ROUTER_NAME_2=${CLUSTER_NAME}-network-2-${ZONE}
export NAT_CONFIG_2=${CLUSTER_NAME}-natconfig-2-${ZONE}
```
Tetapkan variabel lingkungan berikut:
- CLUSTER_NAME: Nama untuk cluster Anda.
- REGION: Region tempat cluster TPU Anda akan dibuat.
- ZONE: Zona tempat cluster TPU akan dibuat. Untuk mengetahui informasi selengkapnya tentang zona yang didukung, lihat Region dan zona.
- PROJECT_ID: Project ID Google Cloud Anda.
- ACCELERATOR_TYPE: Jenis akselerator menentukan versi dan ukuran Cloud TPU yang ingin Anda buat. Misalnya, v6e-256. Untuk mengetahui informasi selengkapnya tentang jenis akselerator yang didukung untuk setiap versi TPU, lihat Versi TPU.
- NUM_SLICES: Jumlah slice TPU untuk cluster Anda. Setiap slice memiliki jumlah chip yang ditentukan dalam ACCELERATOR_TYPE. Untuk cluster slice tunggal, tetapkan NUM_SLICES ke 1. Untuk cluster Multislice, tentukan jumlah slice berdasarkan persyaratan skalabilitas beban kerja Anda. Jumlah total chip dalam cluster adalah jumlah chip dalam ACCELERATOR_TYPE dikalikan dengan NUM_SLICES.

Buat resource jaringan utama

Buat jaringan utama dengan unit transmisi maksimum (MTU) sebesar 8.896:
```
gcloud compute networks create ${NETWORK_NAME_1} \
    --mtu=8896 \
    --bgp-routing-mode=regional \
    --subnet-mode=custom \
    --project=${PROJECT_ID}
```
Menggunakan jaringan kustom dengan MTU 8.896 memberikan performa yang lebih baik. Untuk mengetahui informasi selengkapnya, lihat Mengonfigurasi MTU.

Buat subnetwork utama:

gcloud compute networks subnets create ${SUBNET_NAME_1} \
    --network=${NETWORK_NAME_1} \
    --range=10.11.0.0/18 \
    --region=${REGION} \
    --project=${PROJECT_ID}

Buat aturan firewall untuk jaringan utama yang mengizinkan traffic tcp, icmp, dan udp di jaringan utama:

gcloud compute firewall-rules create ${FIREWALL_RULE_NAME_1} \
    --network=${NETWORK_NAME_1} \
    --allow tcp,icmp,udp \
    --project=${PROJECT_ID}

Buat Cloud Router untuk jaringan utama:

gcloud compute routers create ${ROUTER_NAME_1} \
    --project=${PROJECT_ID} \
    --network=${NETWORK_NAME_1} \
    --region=${REGION}

Konfigurasi NAT untuk jaringan utama. Perintah berikut mengizinkan traffic dari cluster Anda untuk menjangkau internet:

gcloud compute routers nats create ${NAT_CONFIG_1} \
    --router=${ROUTER_NAME_1} \
    --region=${REGION} \
    --auto-allocate-nat-external-ips \
    --nat-all-subnet-ip-ranges \
    --project=${PROJECT_ID} \
    --enable-logging

Membuat resource jaringan sekunder

Buat jaringan sekunder:

gcloud compute networks create ${NETWORK_NAME_2} --mtu=8896 \
    --bgp-routing-mode=regional \
    --subnet-mode=custom \
    --project=${PROJECT_ID}

Buat subnetwork untuk jaringan sekunder:

gcloud compute networks subnets create ${SUBNET_NAME_2} \
    --network=${NETWORK_NAME_2} \
    --range=10.10.0.0/18 \
    --region=${REGION} \
    --project=${PROJECT_ID}

Buat aturan firewall untuk mengizinkan traffic dalam jaringan baru:

gcloud compute firewall-rules create ${FIREWALL_RULE_NAME_2} \
    --network=${NETWORK_NAME_2} \
    --allow tcp,icmp,udp \
    --source-ranges 10.10.0.0/18 \
    --project=${PROJECT_ID}

Buat Cloud Router untuk jaringan sekunder:

gcloud compute routers create ${ROUTER_NAME_2} \
    --project=${PROJECT_ID} \
    --network=${NETWORK_NAME_2} \
    --region=${REGION}

Buat konfigurasi NAT untuk Cloud Router:

gcloud compute routers nats create ${NAT_CONFIG_2} \
    --router=${ROUTER_NAME_2} \
    --region=${REGION} \
    --auto-allocate-nat-external-ips \
    --nat-all-subnet-ip-ranges \
    --project=${PROJECT_ID} \
    --enable-logging

Membuat cluster

Tetapkan variabel lingkungan untuk argumen cluster dan node pool guna menggunakan jaringan dan subnetwork yang Anda buat:

export CLUSTER_ARGUMENTS="--enable-dataplane-v2 --enable-ip-alias --enable-multi-networking --network=${NETWORK_NAME_1} --subnetwork=${SUBNET_NAME_1}"
export NODE_POOL_ARGUMENTS="--additional-node-network network=${NETWORK_NAME_2},subnetwork=${SUBNET_NAME_2}"

Argumen ini mengonfigurasi cluster untuk menggunakan dua jaringan yang Anda buat untuk dukungan multi-NIC.

Buat cluster. Perintah berikut menyediakan kapasitas sesuai permintaan:
```
xpk cluster create \
    --cluster=${CLUSTER_NAME} \
    --cluster-cpu-machine-type=e2-standard-8 \
    --num-slices=${NUM_SLICES} \
    --tpu-type=${ACCELERATOR_TYPE} \
    --zone=${ZONE}  \
    --project=${PROJECT_ID} \
    --on-demand \
    --custom-cluster-arguments="${CLUSTER_ARGUMENTS}" \
    --custom-nodepool-arguments="${NODE_POOL_ARGUMENTS}" \
    --create-vertex-tensorboard
```
Untuk menggunakan kapasitas yang dicadangkan, ganti --on-demand dengan --reservation=RESERVATION_NAME. Untuk menggunakan VM Spot TPU, ganti --on-demand dengan --spot.

Catatan: Flag --cluster-cpu-machine-type menentukan jenis mesin untuk VM host CPU di cluster. Meskipun e2-standard-8 adalah default yang direkomendasikan untuk banyak workload, Anda dapat memilih jenis mesin yang berbeda jika workload Anda memiliki persyaratan CPU atau memori tertentu untuk tugas seperti pra-pemrosesan data.

Memvalidasi penyiapan multi-NIC

Setelah membuat cluster dengan dukungan multi-NIC, Anda dapat memvalidasi bahwa kedua NIC sedang digunakan dengan membuat XPK workload dan menambahkan flag --command ifconfig.

Gunakan perintah berikut untuk menampilkan output perintah ifconfig di log konsolGoogle Cloud . Anda harus menentukan flag --base-docker-image maxtext_base_image untuk menggunakan image dasar MaxText, seperti pada contoh berikut, atau menentukan flag --docker-image dan image yang ingin Anda gunakan.

xpk workload create \
    --cluster ${CLUSTER_NAME} \
    --base-docker-image maxtext_base_image \
    --workload=${USER}-xpk-${ACCELERATOR_TYPE}-${NUM_SLICES} \
    --tpu-type=${ACCELERATOR_TYPE} \
    --num-slices=${NUM_SLICES}  \
    --on-demand \
    --zone=${ZONE} \
    --project=${PROJECT_ID} \
    --command "ifconfig"

Jika Anda ingin mengaktifkan log debug atau menggunakan Vertex AI TensorBoard, tambahkan argumen opsional berikut ke perintah:

--enable-debug-logs \
--use-vertex-tensorboard

Pastikan eth0 dan eth1 memiliki MTU yang ditetapkan ke 8.896 dengan memeriksa output workload XPK di Google Cloud log konsol.

Menyiapkan JAX atau PyTorch

Referensi berikut menunjukkan cara menyiapkan JAX atau PyTorch di TPU Anda, bergantung pada metode penyediaan dan pengelolaan yang Anda gunakan:

GKE Autopilot: Menyiapkan aplikasi TPU
GKE Standard: Menyiapkan workload Anda
GKE dan XPK: README XPK
Cloud TPU host tunggal menggunakan JAX: Menjalankan penghitungan di VM Cloud TPU menggunakan JAX
Cloud TPU multi-host menggunakan JAX: Menjalankan kode JAX pada slice TPU
Cloud TPU host tunggal menggunakan PyTorch: Menjalankan perhitungan di VM Cloud TPU menggunakan PyTorch
Cloud TPU multi-host menggunakan PyTorch: Menjalankan kode PyTorch pada slice TPU

Untuk menyiapkan dan menjalankan XPK dengan MaxText, lihat Menjalankan MaxText dalam Skala Besar dengan XPK .

Meningkatkan kualitas setelan TCP

Jika Anda menyediakan TPU v6e menggunakan resource dalam antrean, Anda dapat menjalankan perintah berikut untuk meningkatkan performa jaringan dengan meningkatkan batas buffer penerimaan TCP.

gcloud alpha compute tpus queued-resources ssh "${QUEUED_RESOURCE_ID}" \
    --project "${PROJECT_ID}" \
    --zone "${ZONE}" \
    --node=all \
    --worker=all \
    --command='
    sudo sh -c "echo \"4096 41943040 314572800\" > /proc/sys/net/ipv4/tcp_rmem"'

Menggunakan SkyPilot

Anda dapat menggunakan Cloud TPU v6e dengan SkyPilot. SkyPilot adalah framework open source yang menyederhanakan proses menjalankan, mengelola, dan menskalakan beban kerja AI. Anda dapat menambahkan informasi harga dan lokasi terkait v6e ke SkyPilot. Untuk mengetahui informasi selengkapnya, lihat contoh SkyPilot TPU v6e.

Contoh pelatihan

Bagian berikut memberikan contoh untuk melatih model MaxText, MaxDiffusion, dan PyTorch di Cloud TPU v6e.

Contoh ini telah diuji dengan versi software berikut:

Python 3.10 atau yang lebih baru
Versi software harian:
- JAX harian 0.4.32.dev20240912
- LibTPU Harian 0.1.dev20240912+nightly
Versi software stabil:
- JAX + JAX Lib v0.4.37

Melatih MaxText dan MaxDiffusion di Cloud TPU v6e

Bagian berikut mencakup siklus proses pelatihan model MaxText dan MaxDiffusion.

Secara umum, langkah-langkah tingkat tingginya adalah:

Buat image dasar workload.
Jalankan workload Anda menggunakan XPK.
1. Buat perintah pelatihan untuk workload.
2. Deploy workload.
Ikuti workload dan lihat metrik.
Hapus workload XPK jika tidak diperlukan.
Hapus cluster jika tidak diperlukan lagi.

Membangun image dasar

Instal MaxText atau MaxDiffusion dan bangun image Docker:

Clone repositori yang ingin Anda gunakan dan ubah ke direktori untuk repositori:

MaxText:

git clone https://github.com/google/maxtext.git && cd maxtext

MaxDiffusion:

git clone https://github.com/google/maxdiffusion.git && cd maxdiffusion && git checkout 4a8155ec0129512812b31930f0a91c6d5a141103

Konfigurasi Docker agar menggunakan Google Cloud CLI:
```
gcloud auth configure-docker
```

Bangun image Docker menggunakan perintah berikut atau menggunakan gambar AI JAX. Untuk mengetahui informasi selengkapnya tentang gambar AI JAX, lihat Gambar AI JAX.

MaxText:

bash docker_build_dependency_image.sh MODE=stable JAX_VERSION=0.4.35

MaxDiffusion:

bash .github/workflows/build_and_upload_images.sh CLOUD_IMAGE_NAME=maxdiffusion_jax_stable_stack MODE=jax_ai_image PROJECT=${PROJECT_ID} LOCAL_IMAGE_NAME=maxdiffusion_jax_stable_stack BASEIMAGE=us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest

Tetapkan project ID Anda dalam konfigurasi gcloud CLI aktif:
```
gcloud config set project ${PROJECT_ID}
```
Jika Anda meluncurkan beban kerja dari mesin yang tidak memiliki image yang dibuat secara lokal, upload image tersebut.

Catatan: Pastikan Anda menggunakan bash versi 4.2 atau yang lebih baru.
1. Tetapkan variabel lingkungan CLOUD_IMAGE_NAME:
```
export CLOUD_IMAGE_NAME=${USER}_runner
```
2. Upload gambar:
```
bash docker_upload_runner.sh ${CLOUD_IMAGE_NAME}
```

Menjalankan workload menggunakan XPK

Tetapkan variabel lingkungan berikut jika Anda tidak menggunakan nilai default yang ditetapkan oleh MaxText atau MaxDiffusion:

export BASE_OUTPUT_DIR=gs://YOUR_BUCKET
export PER_DEVICE_BATCH_SIZE=2
export NUM_STEPS=30
export MAX_TARGET_LENGTH=8192

Buat skrip model Anda. Skrip ini akan disalin sebagai perintah pelatihan pada langkah selanjutnya.

Jangan jalankan skrip model terlebih dahulu.

MaxText

MaxText adalah LLM open source berperforma tinggi dan sangat skalabel yang ditulis dalam Python dan JAX murni serta menargetkan Google Cloud TPU dan GPU untuk pelatihan dan inferensi.

JAX_PLATFORMS=tpu,cpu \
ENABLE_PJRT_COMPATIBILITY=true \
TPU_SLICE_BUILDER_DUMP_CHIP_FORCE=true \
TPU_SLICE_BUILDER_DUMP_ICI=true && \
python3 -m MaxText.train MaxText/configs/base.yml \
     base_output_directory=${BASE_OUTPUT_DIR} \
     dataset_type=synthetic \
     per_device_batch_size=${PER_DEVICE_BATCH_SIZE} \
     enable_checkpointing=false \
     gcs_metrics=true \
     profiler=xplane \
     skip_first_n_steps_for_profiler=5 \
     steps=${NUM_STEPS}  # attention='dot_product'"

Gemma2

Gemma adalah serangkaian LLM dengan bobot terbuka yang dikembangkan oleh Google DeepMind, berdasarkan riset dan teknologi Gemini.

python3 -m MaxText.train MaxText/configs/base.yml \
    model_name=gemma2-27b \
    run_name=gemma2-27b-run \
    base_output_directory=${BASE_OUTPUT_DIR} \
    max_target_length=${MAX_TARGET_LENGTH} \
    per_device_batch_size=${PER_DEVICE_BATCH_SIZE} \
    steps=${NUM_STEPS} \
    enable_checkpointing=false \
    use_iota_embed=true \
    gcs_metrics=true \
    dataset_type=synthetic \
    profiler=xplane \
    attention=flash

Mixtral 8x7b

Mixtral adalah model AI canggih yang dikembangkan oleh Mistral AI, menggunakan arsitektur sparse mixture-of-experts (MoE).

python3 -m MaxText.train MaxText/configs/base.yml \
    base_output_directory=${BASE_OUTPUT_DIR} \
    per_device_batch_size=${PER_DEVICE_BATCH_SIZE} \
    model_name=mixtral-8x7b \
    steps=${NUM_STEPS} \
    max_target_length=${MAX_TARGET_LENGTH} \
    tokenizer_path=assets/tokenizer.mistral-v1 \
    attention=flash \
    dtype=bfloat16 \
    dataset_type=synthetic \
    profiler=xplane

Llama3-8b

Llama adalah serangkaian LLM dengan bobot terbuka yang dikembangkan oleh Meta.

Untuk mengetahui contoh cara menjalankan Llama3 di PyTorch, lihat model torch_xla di repositori torchprime.

MaxDiffusion

MaxDiffusion adalah kumpulan implementasi referensi dari berbagai model difusi laten yang ditulis dalam Python dan JAX murni yang berjalan di perangkat XLA, termasuk Cloud TPU dan GPU. Stable Diffusion adalah model text-to-image laten yang menghasilkan gambar fotorealistis dari input teks apa pun.

Anda harus menginstal cabang Git tertentu untuk menjalankan MaxDiffusion seperti yang ditunjukkan dalam skrip pelatihan berikut.

git clone https://github.com/google/maxdiffusion.git
&& cd maxdiffusion
&& git checkout 4a8155ec0129512812b31930f0a91c6d5a141103
&& pip install -r requirements.txt && pip install .
&& pip install huggingface_hub==0.30.2 && OUT_DIR=${BASE_OUTPUT_DIR}
&& python src/maxdiffusion/train_sdxl.py \
    src/maxdiffusion/configs/base_xl.yml \
    revision=refs/pr/95 \
    activations_dtype=bfloat16 \
    weights_dtype=bfloat16 \
    resolution=1024 \
    per_device_batch_size=1 \
    output_dir=${OUT_DIR} \
    jax_cache_dir=${OUT_DIR}/cache_dir/ \
    max_train_steps=200 \
    attention=flash \
    run_name=sdxl-ddp-v6e

Ekspor variabel berikut:
```
export CLUSTER_NAME=CLUSTER_NAME
export ACCELERATOR_TYPE=ACCELERATOR_TYPE
export NUM_SLICES=NUM_SLICES
export YOUR_MODEL_SCRIPT=YOUR_MODEL_SCRIPT
```
Deskripsi variabel lingkungan
- CLUSTER_NAME: Nama cluster Anda.
- ACCELERATOR_TYPE: Jenis akselerator menentukan versi dan ukuran Cloud TPU yang ingin Anda buat. Untuk mengetahui informasi selengkapnya tentang jenis akselerator yang didukung untuk setiap versi TPU, lihat Versi TPU.
- NUM_SLICES: Jumlah slice TPU.
- YOUR_MODEL_SCRIPT: Skrip model yang akan dieksekusi sebagai perintah pelatihan.
Jalankan model menggunakan skrip yang Anda buat pada langkah sebelumnya. Anda harus menentukan flag --base-docker-image untuk menggunakan image dasar MaxText atau menentukan flag --docker-image dan image yang ingin Anda gunakan.

Anda dapat memilih untuk menambahkan tanda opsional berikut:
- Anda dapat mengaktifkan logging debug dengan menyertakan tanda --enable-debug-logs. Untuk mengetahui informasi selengkapnya, lihat Men-debug JAX di MaxText.
- Anda dapat membuat Eksperimen Vertex AI untuk mengupload data ke Vertex AI TensorBoard dengan menyertakan tanda --use-vertex-tensorboard. Untuk mengetahui informasi selengkapnya, lihat Memantau JAX di MaxText menggunakan Vertex AI.
```
xpk workload create \
  --cluster ${CLUSTER_NAME} \
  {--base-docker-image maxtext_base_image | --docker-image gcr.io/${PROJECT_ID}/${CLOUD_IMAGE_NAME}:latest} \
  --workload=${USER}-xpk-${ACCELERATOR_TYPE}-${NUM_SLICES} \
  --tpu-type=${ACCELERATOR_TYPE} \
  --num-slices=${NUM_SLICES}  \
  --on-demand \
  --zone=${ZONE} \
  --project=${PROJECT_ID} \
  --command="${YOUR_MODEL_SCRIPT}"
```
Outputnya mencakup link untuk memantau beban kerja Anda. Buka link dan klik tab Log untuk melacak beban kerja Anda secara real time.

Men-debug JAX di MaxText

Gunakan perintah XPK tambahan untuk mendiagnosis alasan cluster atau workload tidak berjalan:

Daftar workload XPK
Pemeriksa XPK
Aktifkan logging panjang di log beban kerja Anda menggunakan tanda --enable-debug-logs saat Anda membuat beban kerja XPK

Memantau JAX di MaxText menggunakan Vertex AI

Untuk menggunakan TensorBoard, akun pengguna Google Cloud Anda harus memiliki peran aiplatform.user. Jalankan perintah berikut untuk memberikan peran ini:

gcloud projects add-iam-policy-binding your-project-id \
   --member='user:your-email' \
   --role='roles/aiplatform.user'

Melihat data skalar dan profil melalui TensorBoard terkelola Vertex AI.

Tingkatkan permintaan pengelolaan resource (CRUD) untuk zona yang Anda gunakan dari 600 menjadi 5.000. Hal ini mungkin bukan masalah untuk beban kerja kecil yang menggunakan kurang dari 16 VM.

Instal dependensi seperti cloud-accelerator-diagnostics untuk Vertex AI:

# xpk dependencies will install cloud-accelerator-diagnostics for Vertex AI
cd ~/xpk
pip install .

Buat cluster menggunakan flag --create-vertex-tensorboard, seperti yang didokumentasikan di Membuat Vertex AI TensorBoard. Anda juga dapat menjalankan perintah ini di cluster yang sudah ada.
Buat eksperimen Vertex AI saat menjalankan beban kerja XPK menggunakan tanda --use-vertex-tensorboard dan tanda --experiment-name opsional. Untuk mengetahui daftar lengkap langkah-langkahnya, lihat Membuat Vertex AI Experiment untuk mengupload data ke Vertex AI TensorBoard.

Log tersebut menyertakan link ke Vertex AI TensorBoard, mirip dengan berikut:

View your TensorBoard at https://us-central1.tensorboard.googleusercontent.com/experiment/project_id+locations+us-central1+tensorboards+hash+experiments+name

Anda juga dapat menemukan link Vertex AI TensorBoard di konsol Google Cloud . Buka Vertex AI Experiments di konsol Google Cloud . Pilih wilayah yang sesuai dari drop-down.

Direktori TensorBoard juga ditulis ke bucket Cloud Storage yang Anda tentukan dengan ${BASE_OUTPUT_DIR}.

Menghapus workload XPK

Gunakan perintah xpk workload delete untuk menghapus satu atau beberapa workload berdasarkan awalan tugas atau status tugas. Perintah ini mungkin berguna jika Anda mengirim workload XPK yang tidak perlu lagi dijalankan, atau jika Anda memiliki tugas yang macet dalam antrean.

Hapus cluster Anda

Gunakan perintah xpk cluster delete untuk menghapus cluster Anda:

xpk cluster delete --cluster ${CLUSTER_NAME} \
    --zone=${ZONE} --project=${PROJECT_ID}

Hasil benchmark MaxDiffusion

Kami menjalankan skrip pelatihan untuk MaxDiffusion di v6e-4, v6e-16, dan dua v6e-16. Tabel berikut menunjukkan throughput yang diukur.

	v6e-4	v6e-16	Dua v6e-16
Langkah-langkah pelatihan	0,069	0,073	0,13
Ukuran batch global	8	32	64
Throughput (contoh/dtk)	115,9	438,4	492,3

Melatih model Llama menggunakan PyTorch/XLA di Cloud TPU v6e

Bagian ini menjelaskan cara melatih model Llama menggunakan PyTorch/XLA di Cloud TPU v6e menggunakan set data WikiText.

Mendapatkan akses ke Hugging Face dan model Llama 3

Anda memerlukan token akses pengguna Hugging Face untuk contoh ini. Untuk mengetahui informasi tentang cara membuat token akses pengguna, lihat dokumentasi Hugging Face tentang token akses pengguna.

Anda juga memerlukan izin untuk mengakses model Llama-3-8B di Hugging Face. Untuk mendapatkan akses, buka model Meta-Llama-3-8B di HuggingFace dan minta akses.

Buat VM Cloud TPU

Buat Cloud TPU v6e dengan 8 chip untuk contoh ini.

Siapkan variabel lingkungan:
```
export PROJECT_ID=your-project-id
export TPU_NAME=your-tpu-name
export ZONE=us-east1-d
export ACCELERATOR_TYPE=v6e-8
export RUNTIME_VERSION=v2-alpha-tpuv6e
```
Deskripsi variabel lingkungan
- PROJECT_ID: Project ID Google Cloud Anda. Gunakan project yang ada atau buat project baru.
- TPU_NAME: Nama TPU.
- ZONE: Zona tempat membuat VM TPU. Untuk mengetahui informasi selengkapnya tentang zona yang didukung, lihat Region dan zona TPU.
- ACCELERATOR_TYPE: Jenis akselerator menentukan versi dan ukuran Cloud TPU yang ingin Anda buat. Untuk mengetahui informasi selengkapnya tentang jenis akselerator yang didukung untuk setiap versi TPU, lihat versi TPU.
- RUNTIME_VERSION: Versi software Cloud TPU.

Buat VM Cloud TPU:

gcloud alpha compute tpus tpu-vm create ${TPU_NAME} --version=${RUNTIME_VERSION} \
   --accelerator-type=${ACCELERATOR_TYPE} \
   --zone=${ZONE} \
   --project=${PROJECT_ID}

Penginstalan

Instal fork pytorch-tpu/transformers Hugging Face transformers dan dependensi. Contoh ini diuji dengan versi dependensi berikut:

torch: kompatibel dengan 2.5.0
torch_xla[tpu]: kompatibel dengan 2.5.0
jax: 0.4.33
jaxlib: 0.4.33

gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone ${ZONE} \
   --worker=all \
   --command='git clone -b flash_attention https://github.com/pytorch-tpu/transformers.git
   cd transformers
   sudo pip3 install -e .
   pip3 install datasets
   pip3 install evaluate
   pip3 install scikit-learn
   pip3 install accelerate
   pip install torch~=2.6.0 torch_xla[tpu]~=2.6.0 -f https://storage.googleapis.com/libtpu-releases/index.html -f https://storage.googleapis.com/libtpu-wheels/index.html
   pip install jax==0.4.38 jaxlib==0.4.38 -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/'

Menyiapkan file konfigurasi model

Perintah pelatihan di bagian berikutnya, Jalankan model, menggunakan dua file konfigurasi JSON untuk menentukan parameter model dan konfigurasi Fully Sharded Data Parallel (FSDP). Sharding FSDP memungkinkan Anda menggunakan ukuran batch yang lebih besar saat melatih dengan melakukan sharding bobot model di beberapa TPU. Saat melatih dengan model yang lebih kecil, mungkin cukup menggunakan paralelisme data dan mereplikasi bobot di setiap perangkat. Untuk mengetahui informasi selengkapnya tentang cara memecah tensor di seluruh perangkat di PyTorch/XLA, lihat Panduan pengguna SPMD PyTorch/XLA.

Buat file konfigurasi parameter model. Berikut adalah konfigurasi parameter model untuk Llama-3-8B. Untuk model lainnya, temukan file konfigurasi di Hugging Face. Misalnya, lihat konfigurasi Llama-2-7B.

cat > llama-config.json << EOF
{
  "architectures": [
    "LlamaForCausalLM"
  ],
  "attention_bias": false,
  "attention_dropout": 0.0,
  "bos_token_id": 128000,
  "eos_token_id": 128001,
  "hidden_act": "silu",
  "hidden_size": 4096,
  "initializer_range": 0.02,
  "intermediate_size": 14336,
  "max_position_embeddings": 8192,
  "model_type": "llama",
  "num_attention_heads": 32,
  "num_hidden_layers": 32,
  "num_key_value_heads": 8,
  "pretraining_tp": 1,
  "rms_norm_eps": 1e-05,
  "rope_scaling": null,
  "rope_theta": 500000.0,
  "tie_word_embeddings": false,
  "torch_dtype": "bfloat16",
  "transformers_version": "4.40.0.dev0",
  "use_cache": false,
  "vocab_size": 128256
}
EOF

Buat file konfigurasi FSDP:

cat > fsdp-config.json << EOF
{
  "fsdp_transformer_layer_cls_to_wrap": [
    "LlamaDecoderLayer"
  ],
  "xla": true,
  "xla_fsdp_v2": true,
  "xla_fsdp_grad_ckpt": true
}
EOF

Untuk mengetahui informasi selengkapnya tentang FSDP, lihat Fully Sharded Data Parallel menggunakan SPMD .

Upload file konfigurasi ke VM TPU menggunakan perintah berikut:

gcloud alpha compute tpus tpu-vm scp llama-config.json fsdp-config.json ${TPU_NAME}:. \
   --worker=all \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

Menjalankan model

Dengan menggunakan file konfigurasi yang Anda buat di bagian sebelumnya, jalankan skrip run_clm.py untuk melatih model Llama-3-8B pada set data WikiText. Skrip pelatihan memerlukan waktu sekitar 10 menit untuk dijalankan di Cloud TPU v6e-8.

gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone ${ZONE} \
   --worker=all \
   --command='
   pip3 install "huggingface_hub[cli]"
   huggingface-cli login --token HUGGING_FACE_TOKEN'

Jalankan pelatihan model:

gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone ${ZONE} \
   --worker=all \
   --command='
   export PJRT_DEVICE=TPU
   export XLA_USE_SPMD=1
   export ENABLE_PJRT_COMPATIBILITY=true
   # Optional variables for debugging:
   export XLA_IR_DEBUG=1
   export XLA_HLO_DEBUG=1
   export PROFILE_EPOCH=0
   export PROFILE_STEP=3
   export PROFILE_DURATION_MS=100000
   # Set PROFILE_LOGDIR to a local VM path or gs://my-bucket/profile_path
   export PROFILE_LOGDIR=PROFILE_PATH
   python3 transformers/examples/pytorch/language-modeling/run_clm.py \
     --dataset_name wikitext \
     --dataset_config_name wikitext-2-raw-v1 \
     --per_device_train_batch_size 16 \
     --do_train \
     --output_dir /home/$USER/tmp/test-clm \
     --overwrite_output_dir \
     --config_name /home/$USER/llama-config.json \
     --cache_dir /home/$USER/cache \
     --tokenizer_name meta-llama/Meta-Llama-3-8B \
     --block_size 8192 \
     --optim adafactor \
     --save_strategy no \
     --logging_strategy no \
     --fsdp "full_shard" \
     --fsdp_config /home/$USER/fsdp-config.json \
     --torch_dtype bfloat16 \
     --dataloader_drop_last yes \
     --flash_attention \
     --max_steps 20'

Memecahkan masalah PyTorch/XLA

Jika Anda menetapkan variabel opsional untuk proses debug di bagian sebelumnya, profil untuk model akan disimpan di lokasi yang ditentukan oleh variabel PROFILE_LOGDIR. Anda dapat mengekstrak file xplane.pb yang disimpan di lokasi ini dan menggunakan tensorboard untuk melihat profil di browser menggunakan petunjuk TensorBoard.

Jika PyTorch/XLA tidak berperforma seperti yang diharapkan, lihat Panduan pemecahan masalah, yang berisi saran untuk men-debug, membuat profil, dan mengoptimalkan model Anda.

Melatih model menggunakan TPU v6e

Sebelum memulai

Memverifikasi kuota dan izin

Opsi penyediaan

Menyediakan TPU v6e dengan GKE dan XPK

Membuat cluster dengan dukungan NIC tunggal menggunakan XPK

Instal XPK dan siapkan variabel lingkungan

Membuat cluster

Jaringan kustom

Jaringan default

Membuat cluster dengan dukungan multi-NIC menggunakan XPK

Instal XPK dan siapkan variabel lingkungan

Buat resource jaringan utama

Membuat resource jaringan sekunder

Membuat cluster

Memvalidasi penyiapan multi-NIC

Menyiapkan JAX atau PyTorch

Meningkatkan kualitas setelan TCP

Menggunakan SkyPilot

Contoh pelatihan

Melatih MaxText dan MaxDiffusion di Cloud TPU v6e

Membangun image dasar

Menjalankan workload menggunakan XPK

MaxText

Gemma2

Mixtral 8x7b

Llama3-8b

MaxDiffusion

Deskripsi variabel lingkungan

Men-debug JAX di MaxText

Memantau JAX di MaxText menggunakan Vertex AI

Menghapus workload XPK

Hapus cluster Anda

Hasil benchmark MaxDiffusion

Melatih model Llama menggunakan PyTorch/XLA di Cloud TPU v6e

Mendapatkan akses ke Hugging Face dan model Llama 3

Buat VM Cloud TPU

Deskripsi variabel lingkungan

Penginstalan

Menyiapkan file konfigurasi model

Menjalankan model

Memecahkan masalah PyTorch/XLA