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 transformer, teks ke gambar, 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:

• Inferensi MaxText JetStream pada v6e
• Inferensi PyTorch JetStream di v6e
• Inferensi MaxDiffusion di v6e
• Inferensi vLLM di v6e
• Melakukan inferensi multihost menggunakan Pathways

Sebelum memulai

Sebelum memulai, Anda perlu:

• Buat akun dan project dengan penagihan diaktifkan Google Cloud
• 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

Verifikasi bahwa 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 Seimbang dan Hyperdisk ML. v6e tidak mendukung Persistent Disk (PD). Untuk mengetahui informasi selengkapnya, lihat Opsi penyimpanan untuk data Cloud TPU.

Jika Anda menggunakan GKE dengan XPK, Anda memerlukan izin tambahan di konsol Google Cloud . Untuk mengetahui informasi selengkapnya, lihat Izin yang diperlukan di konsol Google Cloud .

Menyediakan TPU

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. 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. 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. 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 Cloud TPU v6e dengan GKE dan XPK

Jika menggunakan perintah GKE dengan v6e, Anda dapat menggunakan perintah Kubernetes atau XPK untuk menyediakan Cloud TPU dan melatih atau menayangkan model. Lihat Merencanakan Cloud TPU di GKE untuk mempelajari cara merencanakan konfigurasi Cloud TPU di cluster GKE. Bagian berikut menyediakan perintah untuk membuat cluster XPK dengan dukungan NIC tunggal dan dukungan multi-NIC.

Membuat cluster XPK dengan dukungan NIC tunggal

export CLUSTER_NAME=xpk-cluster-name
export ZONE=us-east1-d
export PROJECT_ID=your-project-id
export TPU_TYPE=v6e-256
export NUM_SLICES=2

export NETWORK_NAME=${CLUSTER_NAME}-mtu9k
export NETWORK_FW_NAME=${NETWORK_NAME}-fw

Catatan: Gunakan jaringan kustom untuk performa yang lebih baik dan agar tidak membebani jaringan default.

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

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

export CLUSTER_ARGUMENTS="--network=${NETWORK_NAME} --subnetwork=${NETWORK_NAME}"

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

Deskripsi tanda perintah

Variabel	Deskripsi
CLUSTER_NAME	Nama yang ditetapkan pengguna untuk cluster XPK.
PROJECT_ID	Google Cloud nama project. Gunakan project yang sudah ada atau buat project baru. Untuk mengetahui informasi selengkapnya, lihat Menyiapkan project Google Cloud .
ZONE	Lihat dokumen Region dan zona Cloud TPU untuk mengetahui zona yang didukung.
TPU_TYPE	Lihat Jenis Akselerator.
NUM_SLICES	Jumlah irisan yang ingin Anda buat
CLUSTER_ARGUMENTS	Jaringan dan subnetwork yang akan digunakan. Contoh: --network=${NETWORK_NAME} --subnetwork=${NETWORK_NAME}
NUM_SLICES	Jumlah irisan yang akan dibuat.
NETWORK_NAME	Nama jaringan sekunder yang akan digunakan.
NETWORK_FW_NAME	Nama firewall jaringan sekunder yang akan digunakan.

Membuat cluster XPK dengan dukungan multi-NIC

export CLUSTER_NAME=xpk-cluster-name
export REGION=your-region
export ZONE=us-east1-d
export PROJECT_ID=your-project-id
export TPU_TYPE=v6e-256
export NUM_SLICES=2

export NETWORK_NAME_1=${CLUSTER_NAME}-mtu9k-1-${ZONE}
export SUBNET_NAME_1=${CLUSTER_NAME}-privatesubnet-1-${ZONE}
export NETWORK_FW_NAME_1=${NETWORK_NAME_1}-fw-1-${ZONE}
export FIREWALL_RULE_NAME=${CLUSTER_NAME}-privatefirewall-1-${ZONE}
export ROUTER_NAME=${CLUSTER_NAME}-network-1-${ZONE}
export NAT_CONFIG=${CLUSTER_NAME}-natconfig-1-${ZONE}

Catatan: Gunakan jaringan kustom untuk performa yang lebih baik dan agar tidak membebani jaringan default.

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

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

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

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

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

# Secondary subnet for multi-nic experience.
# Need custom IP routing to be different from the first network's subnet.

export NETWORK_NAME_2=${CLUSTER_NAME}-privatenetwork-2-${ZONE}
export SUBNET_NAME_2=${CLUSTER_NAME}-privatesubnet-2-${ZONE}
export FIREWALL_RULE_NAME=${CLUSTER_NAME}-privatefirewall-2-${ZONE}
export ROUTER_NAME=${CLUSTER_NAME}-network-2-${ZONE}
export NAT_CONFIG=${CLUSTER_NAME}-natconfig-2-${ZONE}

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

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

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

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

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

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}"

python3 xpk.py cluster create \
   --cluster=${CLUSTER_NAME} \
   --cluster-cpu-machine-type=e2-standard-8 \
   --num-slices=${NUM_SLICES} \
   --tpu-type=${TPU_TYPE} \
   --zone=${ZONE}  \
   --project=${PROJECT_ID} \
   --on-demand \
   --custom-cluster-arguments="${CLUSTER_ARGUMENTS}" \
   --custom-nodepool-arguments="${NODE_POOL_ARGUMENTS}" \
   --create-vertex-tensorboard

Deskripsi tanda perintah

Variabel	Deskripsi
CLUSTER_NAME	Nama yang ditetapkan pengguna untuk cluster XPK.
PROJECT_ID	Google Cloud nama project. Gunakan project yang sudah ada atau buat project baru. Untuk mengetahui informasi selengkapnya, lihat Menyiapkan project Google Cloud .
ZONE	Lihat dokumen Region dan zona Cloud TPU untuk mengetahui zona yang didukung.
TPU_TYPE	Lihat Jenis Akselerator.
NUM_SLICES	Jumlah irisan yang ingin Anda buat
CLUSTER_ARGUMENTS	Jaringan dan subnetwork yang akan digunakan. Contoh: --enable-dataplane-v2 --enable-ip-alias --enable-multi-networking --network=${NETWORK_NAME_1} --subnetwork=${SUBNET_NAME_1}
NODE_POOL_ARGUMENTS	Jaringan node tambahan yang akan digunakan. Contoh: --additional-node-network network=${NETWORK_NAME_2},subnetwork=${SUBNET_NAME_2}
NUM_SLICES	Jumlah irisan yang akan dibuat (hanya diperlukan untuk Multislice).
NETWORK_NAME	Nama jaringan sekunder yang akan digunakan.
NETWORK_FW_NAME	Nama firewall jaringan sekunder yang akan digunakan.

Menyiapkan JAX atau PyTorch

Resource berikut menunjukkan cara menyiapkan JAX atau PyTorch di Cloud 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: Jalankan 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 .

Mengoptimalkan performa jaringan

Bagian ini menjelaskan cara mengoptimalkan performa jaringan dengan mengonfigurasi unit transmisi maksimum (MTU), menggunakan multi-NIC untuk lingkungan Multislice, dan meningkatkan setelan TCP.

Mengonfigurasi MTU

Untuk performa jaringan terbaik, gunakan jaringan dengan MTU (unit transmisi maksimum) 8.896.

Secara default, Virtual Private Cloud (VPC) hanya menyediakan MTU sebesar 1.460 byte, yang memberikan performa jaringan yang kurang optimal. Anda dapat menyetel MTU jaringan VPC ke nilai antara 1.300 byte dan 8.896 byte (inklusif). Ukuran MTU kustom umum adalah 1.500 byte (Ethernet standar) atau 8.896 byte (maksimum yang mungkin). Untuk mengetahui informasi selengkapnya, lihat Ukuran MTU jaringan VPC yang valid.

Untuk mengetahui informasi selengkapnya tentang cara mengubah setelan MTU untuk jaringan default atau yang sudah ada, lihat Mengubah setelan MTU jaringan VPC.

Contoh berikut membuat jaringan dengan MTU 8.896 dan aturan firewall yang sesuai yang mengizinkan traffic TCP, ICMP, dan UDP dalam jaringan.

export RESOURCE_NAME=your-resource-name
export NETWORK_NAME=${RESOURCE_NAME}-privatenetwork
export NETWORK_FW_NAME=${RESOURCE_NAME}-privatefirewall
gcloud compute networks create ${NETWORK_NAME} --mtu=8896 --project=${PROJECT_ID} \
    --subnet-mode=auto --bgp-routing-mode=regional
gcloud compute firewall-rules create ${NETWORK_FW_NAME} --network=${NETWORK_NAME} \
    --allow tcp,icmp,udp --project=${PROJECT_ID}

Ganti your-resource-name dengan nama dasar untuk jaringan dan firewall.

Menggunakan opsi multi-NIC untuk Multislice

Jika Anda menggunakan lingkungan Multislice, tetapkan variabel lingkungan berikut, yang diperlukan untuk subnet sekunder:

export NETWORK_NAME_2=${RESOURCE_NAME}
export SUBNET_NAME_2=${RESOURCE_NAME}
export FIREWALL_RULE_NAME=${RESOURCE_NAME}
export ROUTER_NAME=${RESOURCE_NAME}-network-2
export NAT_CONFIG=${RESOURCE_NAME}-natconfig-2
export REGION=your-region

Gunakan perintah berikut untuk membuat perutean IP kustom untuk jaringan dan subnet.

Catatan: Kedua subnetwork harus memiliki rentang IP yang saling eksklusif. Hal ini mencakup rentang CIDR IP utama dan sekunder di subnetwork.

1. Buat jaringan sekunder.

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

2. 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}
   

3. Buat aturan firewall untuk mengizinkan traffic dalam subnetwork baru.

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

4. Buat Cloud Router untuk jaringan sekunder.

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

5. Buat konfigurasi NAT untuk Cloud Router.

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

Setelah membuat slice multi-jaringan, Anda dapat memvalidasi bahwa kedua kartu antarmuka jaringan (NIC) sedang digunakan dengan menyiapkan cluster XPK dan menambahkan flag --command ifconfig ke perintah pembuatan workload XPK.

1. Gunakan perintah workload create berikut untuk menampilkan output perintah ifconfig di log konsol Google Cloud dan periksa apakah eth0 dan eth1 telah menetapkan MTU ke 8.896.

   Catatan: Dalam perintah workload create, tentukan flag --docker-image untuk menggunakan image Docker yang ada atau flag --base-docker-image untuk membuat image container baru tempat menjalankan workload XPK Anda. Flag --base-docker-image menggunakan image dasar yang ditentukan dan menambahkan konten direktori saat ini.
   Catatan: Sebelum menjalankan perintah ini, pastikan Anda telah membuat cluster XPK.

   python3 xpk.py workload create \
       --cluster CLUSTER_NAME \
       {--base-docker-image maxtext_base_image | --docker-image your-cloud-image-name} \
       --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

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

Meningkatkan kualitas setelan TCP

Jika Anda menyediakan Cloud TPU 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"'

Mengoptimalkan performa alokasi memori

Library tcmalloc digunakan secara default di VM Cloud TPU untuk meningkatkan performa model dengan alokasi memori yang cukup besar dan sering. Ini dikonfigurasi melalui variabel lingkungan LD_PRELOAD.

Namun, untuk beberapa workload (misalnya, DLRM dengan alokasi tabel penyematan yang sangat besar), tcmalloc dapat menyebabkan perlambatan. Dalam kasus seperti itu, Anda dapat kembali ke fungsi malloc standar dengan membatalkan setelan variabel LD_PRELOAD di sesi shell sebelum menjalankan skrip pelatihan:

unset LD_PRELOAD

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 lokasi dan harga 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 setiap hari:
  • 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:

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

Membangun image dasar

Instal MaxText atau MaxDiffusion dan bangun image Docker:

1. 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
   

2. Konfigurasi Docker agar menggunakan Google Cloud CLI:

   gcloud auth configure-docker
   

3. 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
   

4. Tetapkan project ID Anda dalam konfigurasi gcloud CLI aktif:

   gcloud config set project ${PROJECT_ID}
   

5. 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

1. 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

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

   Jangan jalankan skrip model terlebih dahulu.

   Catatan: Pembuatan titik pemeriksaan otomatis tidak diaktifkan. Jika Anda melihat error seperti Could not find a known environment for initializing distributed JAX, gunakan JAX_FORCE_TPU_INIT=true.

   MaxText

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

   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.

   Catatan: Untuk menjalankan model ini, Anda harus menyediakan TPU yang berukuran v6e-256 atau lebih besar.

   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).

   Catatan: Untuk menjalankan model ini, Anda harus menyediakan TPU yang berukuran v6e-256 atau lebih besar.

   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.

   Catatan: Untuk menjalankan model ini, Anda harus menyediakan TPU yang berukuran v6e-256 atau lebih besar.

   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 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
   

   Catatan: Secara opsional, tetapkan dcn_data_parallelism untuk Run multislice.

3. 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

   Variabel	Deskripsi
   CLUSTER_NAME	Nama cluster XPK 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.

4. 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.

   python3 xpk.py 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 workload Anda menggunakan tanda --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 ini. 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.

1. 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.

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

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

3. Buat cluster XPK menggunakan flag --create-vertex-tensorboard, seperti yang didokumentasikan di Membuat Vertex AI TensorBoard. Anda juga dapat menjalankan perintah ini di cluster yang sudah ada.

4. 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.

Catatan: Jika Anda menjalankan beban kerja non-MaxText, gunakan contoh MaxText ini untuk membuat perubahan upload data yang sesuai.

Log mencakup 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.

Menghapus cluster XPK

Gunakan perintah xpk cluster delete untuk menghapus cluster Anda:

python3 xpk.py 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.

1. 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

   Variabel	Deskripsi
   PROJECT_ID	ID project Google Cloud Anda. Gunakan project yang ada atau buat project baru.
   TPU_NAME	Nama TPU.
   ZONE	Zona tempat VM TPU akan dibuat. 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.

2. 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 dari transformer Hugging Face 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 Paralel Data yang Sepenuhnya Di-shard (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.

1. 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.

   Penting: Tetapkan use_cache ke false untuk pelatihan.

   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
   

2. 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 .

3. 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.

1. Login ke Hugging Face di Cloud TPU menggunakan perintah berikut:

   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'

2. 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.