Melatih model menggunakan TPU v5e

Perhatian: Halaman ini berlaku untuk Cloud TPU API, yang tidak lagi dalam pengembangan aktif. Cloud TPU API hanya akan menerima perbaikan bug dan update keamanan. Untuk mendapatkan fitur terbaru dan dukungan untuk versi TPU terbaru, sebaiknya gunakan Compute Engine atau Google Kubernetes Engine untuk mengelola resource TPU. Untuk mengetahui informasi selengkapnya, lihat Resource TPU di Compute Engine dan Tentang TPU di GKE. Cloud TPU API tidak didukung di TPU7x dan versi TPU yang lebih baru.

Dengan jejak 256 chip yang lebih kecil per Pod, TPU v5e dioptimalkan untuk menjadi produk bernilai tinggi untuk pelatihan, fine-tuning, dan inferensi Transformer, text-to-image, dan Jaringan Neural Konvolusional (CNN). Untuk mengetahui informasi selengkapnya tentang penggunaan Cloud TPU v5e untuk inferensi, lihat Inferensi menggunakan v5e.

Untuk mengetahui informasi selengkapnya tentang hardware dan konfigurasi TPU Cloud TPU v5e, lihat TPU v5e.

Mulai

Bagian berikut menjelaskan cara mulai menggunakan TPU v5e.

Kuota permintaan

Anda memerlukan kuota untuk menggunakan TPU v5e untuk pelatihan. Ada berbagai jenis kuota untuk TPU on-demand, TPU yang direservasi, dan Spot VM TPU. Ada kuota terpisah yang diperlukan jika Anda menggunakan TPU v5e untuk inferensi. Untuk mengetahui informasi selengkapnya tentang kuota, lihat Kuota. Untuk meminta kuota TPU v5e, hubungi Cloud Sales.

Membuat akun dan project Google Cloud

Anda memerlukan akun dan project Google Cloud untuk menggunakan Cloud TPU. Untuk mengetahui informasi selengkapnya, lihat Menyiapkan lingkungan Cloud TPU.

Buat Cloud TPU

Praktik terbaiknya adalah menyediakan Cloud TPU v5e sebagai resource yang diantrekan menggunakan perintah queued-resource create. Untuk mengetahui informasi selengkapnya, lihat Mengelola resource dalam antrean.

Anda juga dapat menggunakan Create Node API (gcloud compute tpus tpu-vm create) untuk menyediakan Cloud TPU v5e. Untuk mengetahui informasi selengkapnya, lihat Mengelola resource TPU.

Untuk mengetahui informasi selengkapnya tentang konfigurasi v5e yang tersedia untuk pelatihan, lihat Jenis Cloud TPU v5e untuk pelatihan.

Penyiapan framework

Bagian ini menjelaskan proses penyiapan umum untuk pelatihan model kustom menggunakan JAX dengan TPU v5e. Untuk opsi pelatihan PyTorch, lihat Melatih model menggunakan PyTorch di Cloud TPU.

Untuk mengetahui petunjuk penyiapan inferensi, lihat pengantar inferensi v5e.

Tentukan beberapa variabel lingkungan:

export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west4-a
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=your_queued_resource_id

Penyiapan untuk JAX

Jika Anda memiliki bentuk slice yang lebih besar dari 8 chip, Anda akan memiliki beberapa VM dalam satu slice. Dalam hal ini, Anda harus menggunakan flag --worker=all untuk menjalankan penginstalan di semua VM TPU dalam satu langkah tanpa menggunakan SSH untuk login ke setiap VM secara terpisah:

gcloud compute tpus tpu-vm ssh $TPU_NAME  \
   --project=$PROJECT_ID \
   --zone=$ZONE \
   --worker=all \
   --command='pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'

Deskripsi flag perintah

• TPU_NAME: ID teks TPU yang ditetapkan pengguna yang dibuat saat permintaan resource yang diantrekan dialokasikan.
• PROJECT_ID: Google Cloud Nama Project. Gunakan project yang sudah ada atau buat project baru di Siapkan project Google Cloud
• ZONE: Lihat dokumen region dan zona TPU untuk mengetahui zona yang didukung.
• worker: VM TPU yang memiliki akses ke TPU yang mendasarinya.

Anda dapat menjalankan perintah berikut untuk memeriksa jumlah perangkat (output yang ditampilkan di sini dihasilkan dengan slice v5litepod-16). Kode ini menguji apakah semuanya diinstal dengan benar dengan memeriksa apakah JAX melihat TensorCore Cloud TPU dan dapat menjalankan operasi dasar:

gcloud compute tpus tpu-vm ssh $TPU_NAME \
   --project=$PROJECT_ID \
   --zone=$ZONE \
   --worker=all \
   --command='python3 -c "import jax; print(jax.device_count()); print(jax.local_device_count())"'

Outputnya akan seperti ini:

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
SSH: Attempting to connect to worker 2...
SSH: Attempting to connect to worker 3...
16
4
16
4
16
4
16
4

jax.device_count() menunjukkan jumlah total chip dalam slice tertentu. jax.local_device_count() menunjukkan jumlah chip yang dapat diakses oleh satu VM di slice ini.

# Check the number of chips in the given slice by summing the count of chips
# from all VMs through the
# jax.local_device_count() API call.
gcloud compute tpus tpu-vm ssh $TPU_NAME \
   --project=$PROJECT_ID \
   --zone=$ZONE \
   --worker=all \
   --command='python3 -c "import jax; xs=jax.numpy.ones(jax.local_device_count()); print(jax.pmap(lambda x: jax.lax.psum(x, \"i\"), axis_name=\"i\")(xs))"'

Outputnya akan seperti ini:

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
SSH: Attempting to connect to worker 2...
SSH: Attempting to connect to worker 3...
[16. 16. 16. 16.]
[16. 16. 16. 16.]
[16. 16. 16. 16.]
[16. 16. 16. 16.]

Coba Tutorial JAX dalam dokumen ini untuk mulai menggunakan pelatihan v5e menggunakan JAX.

Pembersihan

Hapus TPU dan resource yang diantrekan di akhir sesi Anda. Untuk menghapus resource yang diantrekan, hapus slice, lalu hapus resource yang diantrekan dalam 2 langkah:

gcloud compute tpus tpu-vm delete $TPU_NAME \
   --project=$PROJECT_ID \
   --zone=$ZONE \
   --quiet

gcloud compute tpus queued-resources delete $QUEUED_RESOURCE_ID \
   --project=$PROJECT_ID \
   --zone=${ZONE} \
   --quiet

Dua langkah ini juga dapat digunakan untuk menghapus permintaan resource dalam antrean yang berada dalam status FAILED.

Contoh JAX/FLAX

Bagian berikut menjelaskan contoh cara melatih model JAX dan FLAX di TPU v5e.

Melatih ImageNet di v5e

Tutorial ini menjelaskan cara melatih ImageNet di v5e menggunakan data input palsu. Jika Anda ingin menggunakan data nyata, lihat file README di GitHub.

Siapkan

1. Buat variabel lingkungan:

   export PROJECT_ID=your-project-id
   export TPU_NAME=your-tpu-name
   export ZONE=us-west4-a
   export ACCELERATOR_TYPE=v5litepod-8
   export RUNTIME_VERSION=v2-alpha-tpuv5-lite
   export SERVICE_ACCOUNT=your-service-account
   export QUEUED_RESOURCE_ID=your-queued-resource-id

   Deskripsi variabel lingkungan

   • PROJECT_ID: Project ID Google Cloud Anda. Gunakan project yang sudah 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.
   • SERVICE_ACCOUNT: Alamat email untuk akun layanan Anda. Anda dapat menemukannya dengan membuka halaman Akun Layanan di konsol Google Cloud .

     Contoh: tpu-service-account@PROJECT_ID.iam.

   • QUEUED_RESOURCE_ID: ID teks yang ditetapkan pengguna untuk permintaan resource yang diantrekan.

2. Membuat resource TPU:

   gcloud compute tpus queued-resources create $QUEUED_RESOURCE_ID \
      --node-id=$TPU_NAME \
      --project=$PROJECT_ID \
      --zone=$ZONE \
      --accelerator-type=$ACCELERATOR_TYPE \
      --runtime-version=$RUNTIME_VERSION \
      --service-account=$SERVICE_ACCOUNT
   

   Catatan: Untuk menggunakan pemesanan, tambahkan flag --reserved. Untuk menggunakan VM Spot TPU, tambahkan flag --spot.

   Anda akan dapat melakukan SSH ke VM TPU setelah resource yang diantrekan berada dalam status ACTIVE:

   gcloud compute tpus queued-resources describe $QUEUED_RESOURCE_ID  \
      --project=$PROJECT_ID \
      --zone=$ZONE
   

   Jika QueuedResource berada dalam status ACTIVE, output-nya akan mirip dengan berikut ini:

    state: ACTIVE
   

3. Instal JAX dan jaxlib versi terbaru:

   gcloud compute tpus tpu-vm ssh $TPU_NAME \
      --project=$PROJECT_ID \
      --zone=$ZONE \
      --worker=all \
      --command='pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
   

   Catatan: Jika Anda melihat error yang menyarankan Anda menggunakan jax>=0.4.8, Anda dapat mengabaikan pesan tersebut.

4. Buat clone model ImageNet dan instal persyaratan yang sesuai:

   gcloud compute tpus tpu-vm ssh $TPU_NAME \
      --project=$PROJECT_ID \
      --zone=$ZONE \
      --worker=all \
      --command="git clone https://github.com/coolkp/flax.git && cd flax && git checkout pmap-orbax-conversion && git pull"
   

5. Untuk membuat data palsu, model memerlukan informasi tentang dimensi set data. Informasi ini dapat dikumpulkan dari metadata set data ImageNet:

   gcloud compute tpus tpu-vm ssh $TPU_NAME \
      --project=$PROJECT_ID \
      --zone=$ZONE \
      --worker=all \
      --command="cd flax/examples/imagenet && pip install -r requirements-cloud-tpu.txt"
   

Melatih model

Setelah semua langkah sebelumnya selesai, Anda dapat melatih model.

gcloud compute tpus tpu-vm ssh $TPU_NAME \
   --project=$PROJECT_ID \
   --zone=$ZONE \
   --worker=all \
   --command="cd flax/examples/imagenet && bash ../../tests/download_dataset_metadata.sh && JAX_PLATFORMS=tpu python imagenet_fake_data_benchmark.py"

Menghapus TPU dan resource yang diantrekan

Hapus TPU dan resource yang diantrekan di akhir sesi Anda.

gcloud compute tpus tpu-vm delete $TPU_NAME \
   --project=$PROJECT_ID \
   --zone=$ZONE \
   --quiet

gcloud compute tpus queued-resources delete $QUEUED_RESOURCE_ID \
   --project=$PROJECT_ID \
   --zone=$ZONE \
   --quiet

Model FLAX Hugging Face

Model Hugging Face yang diimplementasikan di FLAX dapat langsung digunakan di Cloud TPU v5e. Bagian ini memberikan petunjuk untuk menjalankan model populer.

Melatih ViT di Imagenette

Tutorial ini menunjukkan cara melatih model Vision Transformer (ViT) dari HuggingFace menggunakan set data Imagenette Fast AI di Cloud TPU v5e.

Model ViT adalah model pertama yang berhasil melatih encoder Transformer di ImageNet dengan hasil yang sangat baik dibandingkan dengan jaringan konvolusional. Untuk mengetahui informasi selengkapnya, lihat Ringkasan ViT.

Siapkan

1. Buat variabel lingkungan:

   export PROJECT_ID=your-project-id
   export TPU_NAME=your-tpu-name
   export ZONE=us-west4-a
   export ACCELERATOR_TYPE=v5litepod-16
   export RUNTIME_VERSION=v2-alpha-tpuv5-lite
   export SERVICE_ACCOUNT=your-service-account
   export QUEUED_RESOURCE_ID=your-queued-resource-id

   Deskripsi variabel lingkungan

   • PROJECT_ID: Project ID Google Cloud Anda. Gunakan project yang sudah 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.
   • SERVICE_ACCOUNT: Alamat email untuk akun layanan Anda. Anda dapat menemukannya dengan membuka halaman Akun Layanan di konsol Google Cloud .

     Contoh: tpu-service-account@PROJECT_ID.iam.

   • QUEUED_RESOURCE_ID: ID teks yang ditetapkan pengguna untuk permintaan resource yang diantrekan.

2. Membuat resource TPU:

   gcloud compute tpus queued-resources create $QUEUED_RESOURCE_ID \
      --node-id=$TPU_NAME \
      --project=$PROJECT_ID \
      --zone=$ZONE \
      --accelerator-type=$ACCELERATOR_TYPE \
      --runtime-version=$RUNTIME_VERSION \
      --service-account=$SERVICE_ACCOUNT
   

   Catatan: Untuk menggunakan pemesanan, tambahkan flag --reserved. Untuk menggunakan VM Spot TPU, tambahkan flag --spot.

   Anda akan dapat melakukan SSH ke VM TPU setelah resource yang diantrekan berada dalam status ACTIVE:

   gcloud compute tpus queued-resources describe $QUEUED_RESOURCE_ID  \
      --project=$PROJECT_ID \
      --zone=$ZONE
   

   Jika resource yang diantrekan berada dalam status ACTIVE, output-nya akan mirip dengan berikut ini:

    state: ACTIVE
   

3. Instal JAX dan library-nya:

   gcloud compute tpus tpu-vm ssh $TPU_NAME \
      --project=$PROJECT_ID \
      --zone=$ZONE \
      --worker=all \
      --command='pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
   

4. Download repositori Hugging Face dan instal persyaratan:

   gcloud compute tpus tpu-vm ssh $TPU_NAME \
      --project=$PROJECT_ID \
      --zone=$ZONE \
      --worker=all \
      --command='git clone https://github.com/huggingface/transformers.git && cd transformers && pip install . && pip install -r examples/flax/_tests_requirements.txt && pip install --upgrade huggingface-hub urllib3 zipp && pip install tensorflow==2.19 && sed -i "s/torchvision==0.12.0+cpu/torchvision==0.22.1/" examples/flax/vision/requirements.txt && pip install -r examples/flax/vision/requirements.txt && pip install tf-keras'
   

5. Download set data Imagenette:

   gcloud compute tpus tpu-vm ssh $TPU_NAME \
      --project=$PROJECT_ID \
      --zone=$ZONE \
      --worker=all \
      --command='cd transformers && wget https://s3.amazonaws.com/fast-ai-imageclas/imagenette2.tgz && tar -xvzf imagenette2.tgz'
   

Melatih model

Latih model dengan buffer yang telah dipetakan sebelumnya sebesar 4 GB.

gcloud compute tpus tpu-vm ssh $TPU_NAME \
   --project=$PROJECT_ID \
   --zone=$ZONE \
   --worker=all \
   --command='cd transformers && JAX_PLATFORMS=tpu python3 examples/flax/vision/run_image_classification.py --train_dir "imagenette2/train" --validation_dir "imagenette2/val" --output_dir "./vit-imagenette" --learning_rate 1e-3 --preprocessing_num_workers 32 --per_device_train_batch_size 8 --per_device_eval_batch_size 8 --model_name_or_path google/vit-base-patch16-224-in21k --num_train_epochs 3'

Menghapus TPU dan resource yang diantrekan

Hapus TPU dan resource yang diantrekan di akhir sesi Anda.

gcloud compute tpus tpu-vm delete $TPU_NAME \
   --project=$PROJECT_ID \
   --zone=$ZONE \
   --quiet

gcloud compute tpus queued-resources delete $QUEUED_RESOURCE_ID \
   --project=$PROJECT_ID \
   --zone=$ZONE \
   --quiet

Hasil tolok ukur ViT

Skrip pelatihan dijalankan di v5litepod-4, v5litepod-16, dan v5litepod-64. Tabel berikut menunjukkan throughput dengan berbagai jenis akselerator.

Jenis akselerator	v5litepod-4	v5litepod-16	v5litepod-64
Epoch	3	3	3
Ukuran batch global	32	128	512
Throughput (contoh/dtk)	263,40	429,34	470,71

Melatih Difusi pada Pokémon

Tutorial ini menunjukkan cara melatih model Stable Diffusion dari HuggingFace menggunakan set data Pokémon di Cloud TPU v5e.

Model Stable Diffusion adalah model text-to-image laten yang menghasilkan gambar fotorealistik dari input teks apa pun. Untuk informasi selengkapnya, lihat referensi berikut:

• Ringkasan Stable Diffusion
• Sumber kode Stable Diffusion

Siapkan

1. Tetapkan variabel lingkungan untuk nama bucket penyimpanan Anda:

   export GCS_BUCKET_NAME=your_bucket_name
   

2. Siapkan bucket penyimpanan untuk output model Anda:

   Google Cloud CLI storage buckets create gs://$GCS_BUCKET_NAME \
        --project=$PROJECT_ID \
        --location=$ZONE
   

3. Buat variabel lingkungan:

   export PROJECT_ID=your-project-id
   export TPU_NAME=your-tpu-name
   export ZONE=us-west1-c
   export ACCELERATOR_TYPE=v5litepod-16
   export RUNTIME_VERSION=v2-alpha-tpuv5-lite
   export SERVICE_ACCOUNT=your-service-account
   export QUEUED_RESOURCE_ID=your-queued-resource-id

   Deskripsi variabel lingkungan

   • PROJECT_ID: Project ID Google Cloud Anda. Gunakan project yang sudah 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.
   • SERVICE_ACCOUNT: Alamat email untuk akun layanan Anda. Anda dapat menemukannya dengan membuka halaman Akun Layanan di konsol Google Cloud .

     Contoh: tpu-service-account@PROJECT_ID.iam.

   • QUEUED_RESOURCE_ID: ID teks yang ditetapkan pengguna untuk permintaan resource yang diantrekan.

4. Membuat resource TPU:

   gcloud compute tpus queued-resources create $QUEUED_RESOURCE_ID \
      --node-id=$TPU_NAME \
      --project=$PROJECT_ID \
      --zone=$ZONE \
      --accelerator-type=$ACCELERATOR_TYPE \
      --runtime-version=$RUNTIME_VERSION \
      --service-account=$SERVICE_ACCOUNT
   

   Catatan: Untuk menggunakan pemesanan, tambahkan flag --reserved. Untuk menggunakan VM Spot TPU, tambahkan flag --spot.

   Anda akan dapat melakukan SSH ke VM TPU setelah resource yang diantrekan berada dalam status ACTIVE:

   gcloud compute tpus queued-resources describe $QUEUED_RESOURCE_ID  \
      --project=$PROJECT_ID \
      --zone=$ZONE
   

   Jika resource yang diantrekan berada dalam status ACTIVE, output-nya akan mirip dengan berikut ini:

    state: ACTIVE
   

5. Instal JAX dan library-nya.

   gcloud compute tpus tpu-vm ssh $TPU_NAME \
      --project=$PROJECT_ID \
      --zone=$ZONE \
      --worker=all \
      --command='pip install "jax[tpu]==0.4.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
   

6. Download repositori HuggingFace dan instal persyaratan.

   gcloud compute tpus tpu-vm ssh $TPU_NAME \
        --project=$PROJECT_ID \
        --zone=$ZONE \
        --worker=all \
        --command='git clone https://github.com/RissyRan/diffusers.git && cd diffusers && pip install . && pip install -U -r examples/text_to_image/requirements_flax.txt && pip install tensorflow==2.17.1 clu && pip install tensorboard==2.17.1'
   

Melatih model

Latih model dengan buffer yang telah dipetakan sebelumnya sebesar 4 GB.

gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE --project=$PROJECT_ID --worker=all --command="
    git clone https://github.com/google/maxdiffusion
    cd maxdiffusion
    pip3 install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
    pip3 install -r requirements.txt
    pip3 install .
    pip3 install gcsfs
    export LIBTPU_INIT_ARGS=''
    python -m src.maxdiffusion.train src/maxdiffusion/configs/base_2_base.yml run_name=my_run \
    jax_cache_dir=gs://${GCS_BUCKET_NAME} activations_dtype=bfloat16 weights_dtype=bfloat16 \
    per_device_batch_size=1 precision=DEFAULT dataset_save_location=gs://${GCS_BUCKET_NAME} \
    output_dir=gs://${GCS_BUCKET_NAME}/ attention=flash"

Pembersihan

Hapus TPU, resource yang diantrekan, dan bucket Cloud Storage Anda di akhir sesi.

1. Hapus TPU Anda:

   gcloud compute tpus tpu-vm delete $TPU_NAME \
       --project=$PROJECT_ID \
       --zone=$ZONE \
       --quiet
   

2. Hapus resource yang diantrekan:

   gcloud compute tpus queued-resources delete $QUEUED_RESOURCE_ID \
       --project=$PROJECT_ID \
       --zone=$ZONE \
       --quiet
   

3. Hapus bucket Cloud Storage:

   gcloud storage rm -r gs://$GCS_BUCKET_NAME
   

Hasil tolok ukur untuk difusi

Skrip pelatihan dijalankan di v5litepod-4, v5litepod-16, dan v5litepod-64. Tabel berikut menunjukkan throughput.

Jenis akselerator	v5litepod-4	v5litepod-16	v5litepod-64
Langkah Pelatihan	1500	1500	1500
Ukuran batch global	32	64	128
Throughput (contoh/dtk)	36,53	43,71	49,36

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 VM TPU, lihat Memulai TorchTPU di VM TPU. Untuk mengetahui informasi selengkapnya tentang ekosistem PyTorch di Cloud TPU, lihat PyTorch di Cloud TPU.