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

Melatih model menggunakan TPU v6e

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

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, fine-tuning, dan inferensi 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 cara menjalankan inferensi di Cloud TPU v6e, lihat tutorial berikut:

Sebelum memulai

Sebelum memulai, Anda harus:

Menyiapkan project Google Cloud untuk TPU
Menginstal Google Cloud CLI
Mengaktifkan Compute Engine API
Membuat akun layanan Compute Engine dan memberikan izin

Untuk mengetahui informasi selengkapnya, lihat Menyiapkan lingkungan Cloud TPU.

Memverifikasi kuota dan izin

Pastikan project Anda memiliki kuota berikut:

Kuota akselerator TPU Compute Engine
Kuota VM Compute Engine
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:

Compute Engine (Google Cloud CLI): Anda dapat menggunakan alat Google Cloud CLI untuk mengelola VM TPU secara langsung sebagai resource Compute Engine. Anda dapat menyediakan VM TPU individual atau menggunakan Grup Instance Terkelola (MIG) untuk mengelola grup VM TPU. Untuk mengetahui informasi selengkapnya, lihat Ringkasan pembuatan TPU.
GKE: Anda dapat menggunakan GKE untuk menyediakan dan mengelola TPU sebagai kumpulan akselerator untuk workload machine learning dalam container. 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. Alat 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.

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 menyajikan model. Untuk mengetahui informasi selengkapnya tentang cara menggunakan 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 menyediakan 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 konfigurasi yang lebih sedikit. 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 single-slice, 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.

Menginstal XPK dan menyiapkan 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. Contoh, 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 workload 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 agar menggunakan jaringan yang Anda buat:
```
export CLUSTER_ARGUMENTS="--network=${NETWORK_NAME} --subnetwork=${NETWORK_NAME}"
```

Buat cluster XPK. Perintah berikut menyediakan kapasitas on-demand:

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 dipesan, 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 on-demand:

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 dipesan, 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 workload Multislice atau workload 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 akan dialokasikan antarmuka jaringan tambahan, yang masing-masing terhubung ke jaringan VPC yang unik, sehingga meningkatkan throughput jaringan secara keseluruhan. Untuk sebagian besar workload 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.

Menginstal XPK dan menyiapkan 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. Contoh, 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 workload Anda. Jumlah total chip dalam cluster adalah jumlah chip dalam ACCELERATOR_TYPE dikalikan dengan NUM_SLICES.

Membuat resource jaringan utama

Buat jaringan utama dengan unit transmisi maksimum (MTU) 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 akan 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}

Konfigurasikan NAT untuk jaringan utama. Perintah berikut memungkinkan traffic dari cluster Anda 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 agar dapat 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 agar menggunakan dua jaringan yang Anda buat untuk dukungan multi-NIC.

Buat cluster. Perintah berikut menyediakan kapasitas on-demand:
```
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 dipesan, 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 prapemrosesan data.

Memvalidasi penyiapan multi-NIC

Setelah membuat cluster dengan dukungan multi-NIC, Anda dapat memvalidasi bahwa kedua NIC digunakan dengan membuat workload XPK 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 dalam 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 log konsol Google Cloud .

Menyiapkan JAX atau PyTorch

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

GKE Autopilot: Menyiapkan aplikasi TPU Anda
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 di slice TPU
Cloud TPU menggunakan PyTorch: PyTorch di Cloud TPU

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

Menggunakan SkyPilot

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

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 Nightly 0.4.32.dev20240912
- LibTPU Nightly 0.1.dev20240912+nightly
Versi software yang stabil:
- JAX + JAX Lib v0.4.37

Melatih MaxText dan MaxDiffusion di Cloud TPU v6e

Bagian berikut membahas siklus proses pelatihan model MaxText dan MaxDiffusion.

Secara umum, langkah-langkah tingkat tinggi adalah:

Membangun image dasar workload.
Jalankan workload Anda menggunakan XPK.
1. Membuat perintah pelatihan untuk workload.
2. Men-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
git checkout maxtext-v0.2.0

MaxDiffusion:

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

Konfigurasikan Docker untuk 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 dependencies/scripts/docker_build_dependency_image.sh MODE=stable

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 di konfigurasi gcloud CLI yang aktif:
```
gcloud config set project ${PROJECT_ID}
```
Jika Anda meluncurkan workload dari mesin yang tidak memiliki image yang di-build 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 dependencies/scripts/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 berikutnya.

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 TPU dan GPU Google Cloud 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.trainers.pre_train.train src/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 rangkaian LLM berbobot terbuka yang dikembangkan oleh Google DeepMind, berdasarkan riset dan teknologi Gemini.

python3 -m maxtext.trainers.pre_train.train src/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, yang menggunakan arsitektur sparse mixture-of-experts (MoE).

python3 -m maxtext.trainers.pre_train.train src/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

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 fotorealistik dari input teks apa pun.

Anda perlu 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 dijalankan sebagai perintah pelatihan.
Jalankan model menggunakan skrip yang Anda buat di 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 flag opsional berikut:
- Anda dapat mengaktifkan logging debug dengan menyertakan flag --enable-debug-logs. Untuk mengetahui informasi selengkapnya, lihat Men-debug JAX di MaxText.
- Anda dapat membuat Vertex AI Experiment untuk mengupload data ke Vertex AI TensorBoard dengan menyertakan flag --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}"
```
Output-nya menyertakan link untuk mengikuti workload Anda. Buka link dan klik tab Logs untuk melacak workload Anda secara real time.

Men-debug JAX di MaxText

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

Daftar workload XPK
XPK inspector
Aktifkan logging panjang di log workload Anda menggunakan flag --enable-debug-logs saat Anda membuat workload 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 tidak menjadi masalah untuk workload 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 dalam Membuat Vertex AI TensorBoard. Anda juga dapat menjalankan perintah ini pada cluster yang ada.
Buat eksperimen Vertex AI saat menjalankan workload XPK menggunakan flag --use-vertex-tensorboard dan flag --experiment-name opsional. Untuk mengetahui daftar lengkap langkah-langkahnya, lihat Membuat Eksperimen Vertex AI untuk mengupload data ke Vertex AI TensorBoard.

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

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 region yang sesuai dari menu drop-down.

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

Menghapus workload XPK Anda

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 dijalankan lagi, atau jika ada 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 tolok ukur 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 menggunakan PyTorch di Cloud TPU

Anda dapat menggunakan TorchTPU secara langsung untuk melatih model atau menggunakan kembali library referensi TorchTitan. Untuk mulai menggunakan workload v6e, lihat resep TorchTitan Trillium. Untuk mengetahui informasi selengkapnya tentang ekosistem PyTorch di Cloud TPU, lihat PyTorch di Cloud TPU.

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

Menginstal XPK dan menyiapkan variabel lingkungan

Membuat cluster

Jaringan kustom

Jaringan default

Membuat cluster dengan dukungan multi-NIC menggunakan XPK

Menginstal XPK dan menyiapkan variabel lingkungan

Membuat resource jaringan utama

Membuat resource jaringan sekunder

Membuat cluster

Memvalidasi penyiapan multi-NIC

Menyiapkan JAX atau PyTorch

Menggunakan SkyPilot

Contoh pelatihan

Melatih MaxText dan MaxDiffusion di Cloud TPU v6e

Membangun image dasar

Menjalankan workload menggunakan XPK

MaxText

Gemma2

Mixtral 8x7b

MaxDiffusion

Deskripsi variabel lingkungan

Men-debug JAX di MaxText

Memantau JAX di MaxText menggunakan Vertex AI

Menghapus workload XPK Anda

Hapus cluster Anda

Hasil tolok ukur MaxDiffusion

Melatih model menggunakan PyTorch di Cloud TPU