Pelatihan Cloud TPU v5p

Cloud TPU v5p adalah Cloud TPU generasi kelima dari Google Cloud dan merupakan penerus TPU v4. v5p dioptimalkan untuk pelatihan berskala besar dan menjadi platform terdepan untuk pengembangan LLM, model difusi, dan AI generatif dasar. Pada level tinggi, v5p memberikan performa hingga 2x lebih banyak dibandingkan v4, sekaligus mengemas 2x lebih banyak TPU ke dalam Pod (slice terbesar 6k dibandingkan 3k di v4), sehingga menghasilkan performa hingga 4x lipat pada level Pod. Layanan ini juga berjalan pada frekuensi clock yang lebih tinggi (1,75 Ghz vs. 1,05 Ghz), menambahkan SparseCore untuk embedding skala besar, dan tiga kali lipat kapasitas Memori Bandwidth Tinggi (HBM).

Konsep Cloud TPU v5p

Jika Anda baru menggunakan Cloud TPU, lihat beranda dokumentasi TPU.

Konsep Cloud TPU (misalnya, slice, host, dan TensorCore) dan arsitektur sistem Cloud TPU untuk semua versi Cloud TPU dijelaskan di halaman Arsitektur Sistem Cloud TPU.

Setiap versi Cloud TPU memerlukan jenis akselerator khusus untuk pelatihan atau inferensi. Jenis akselerator ini dijelaskan dalam konfigurasi v5p.

Mengelola resource TPU

Semua perintah yang dapat Anda gunakan untuk mengelola VM TPU dijelaskan dalam Mengelola TPU atau Panduan pengguna resource yang diantrekan untuk mengelola resource yang diantrekan.

Penyiapan Framework

Bagian ini menjelaskan proses penyiapan umum untuk pelatihan model menggunakan JAX atau PyTorch dengan TPU v5p.

Penyiapan untuk JAX

Jika memiliki bentuk slice yang lebih besar dari 4 chip, Anda akan memiliki beberapa VM dalam satu slice. Dalam hal ini, Anda harus menggunakan flag --worker=all untuk menjalankan penginstalan pada semua VM TPU menggunakan satu perintah:

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

Anda dapat menjalankan perintah berikut untuk memeriksa jumlah perangkat (output yang ditampilkan di sini dihasilkan dengan slice v5p-32). Kode ini menguji bahwa semuanya telah diinstal dengan benar dengan memastikan bahwa 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 mirip dengan berikut 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 irisan yang ditentukan. jax.local_device_count() menunjukkan jumlah chip yang dapat diakses oleh satu VM dalam irisan 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 mirip dengan berikut 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.]

Gunakan --node=all untuk menjalankan perintah di semua pekerja Multislice.

gcloud compute tpus queued-resources ssh ${QUEUED_RESOURCE_ID} \
--project ${PROJECT_ID} --zone ${ZONE} --node=all --worker=all \
--command='python3 -c "import jax; print(jax.device_count()); print(jax.local_device_count())"'

Coba tutorial JAX dalam dokumen ini untuk memulai pelatihan v5p menggunakan JAX.

Penyiapan untuk PyTorch

Runtime PJRT adalah satu-satunya runtime yang didukung untuk v5p, dan PyTorch 2.1+ menggunakan PJRT sebagai runtime default untuk semua versi TPU. Bagian ini menjelaskan cara mulai menggunakan PJRT pada Pod v5p dengan PyTorch/XLA 2.2.0 untuk semua pekerja.

Menginstal dependensi

gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
--project ${PROJECT_ID} \
--zone ${ZONE} \
--worker=all \
--command='
sudo apt-get update
sudo apt-get install libopenblas-dev -y
pip3 install numpy
pip install torch~=2.2.0 torch_xla[tpu]~=2.2.0 -f https://storage.googleapis.com/libtpu-releases/index.html
'

Gunakan skrip Python dengan PJRT untuk melakukan validasi penginstalan guna menampilkan perangkat TPU yang tersedia (output yang ditampilkan di sini dibuat dengan irisan v5p-32).

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project ${PROJECT_ID} --zone ${ZONE} --worker=all \
--command='
PJRT_DEVICE=TPU python3 -c "import torch_xla.core.xla_model as xm; print(xm.get_xla_supported_devices(\"TPU\"))"
'

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...
['xla:0', 'xla:1', 'xla:2', 'xla:3']
['xla:0', 'xla:1', 'xla:2', 'xla:3']
['xla:0', 'xla:1', 'xla:2', 'xla:3']
['xla:0', 'xla:1', 'xla:2', 'xla:3']

Gunakan --node=all untuk menjalankan perintah di semua pekerja Multislice.

gcloud compute tpus queued-resources ssh ${QUEUED_RESOURCE_ID} \
--project ${PROJECT_ID} --zone ${ZONE} --node=all --worker=all \
--command='
PJRT_DEVICE=TPU python3 -c "import torch_xla.core.xla_model as xm; print(xm.get_xla_supported_devices(\"TPU\"))"
'

Coba tutorial PyTorch di dokumen ini untuk memulai pelatihan v5p menggunakan PyTorch.

Memantau dan membuat profil

Cloud TPU v5p mendukung pemantauan dan pembuatan profil menggunakan metode yang sama seperti Cloud TPU generasi sebelumnya. Anda dapat membaca Membuat profil model Anda dengan alat Cloud TPU untuk mempelajari lebih lanjut pembuatan profil dan Memantau VM Cloud TPU untuk mempelajari pemantauan lebih lanjut.

Tutorial pelatihan

Bagian ini berfokus pada tutorial pelatihan untuk satu slice. Penyesuaian tutorial ini ke pelatihan Multislice dapat dicapai dengan menambahkan flag --node=all ke perintah SSH. Untuk mengetahui detail dan praktik terbaik, lihat Pengantar multislice.

Catatan: Sebelum menjalankan tutorial ini, Anda perlu menyiapkan lingkungan Cloud TPU. Setelah menyiapkan lingkungan TPU, Anda siap menjalankan tutorial berikut.

• Difusi 2.1
• MaxText
• ResNet di v5p host tunggal
• ResNet pada v5p multi-host

Tutorial JAX

Difusi Kereta 2.1

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

Model Difusi Stabil adalah model teks ke gambar laten yang menghasilkan gambar realistik foto dari input teks apa pun. Untuk informasi selengkapnya, lihat referensi berikut:

• Ringkasan Difusi Stabil
• Sumber kode Diffusion Stabil

Catatan: Sebelum menjalankan tutorial ini, Anda perlu menyiapkan lingkungan Cloud TPU.

Penyiapan

1. Buat variabel lingkungan:

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5p-32
   export ZONE=us-east5-a
   export RUNTIME_VERSION=v2-alpha-tpuv5
   export SERVICE_ACCOUNT=your_service_account
   export TPU_NAME=your_tpu_name
   export QUEUED_RESOURCE_ID=queued_resource_id
   export QUOTA_TYPE=quota_type
   export VALID_UNTIL_DURATION=1d

   Deskripsi flag perintah

   Variabel	Deskripsi
   PROJECT_ID	Nama Project Google Cloud
   ACCELERATOR_TYPE	Lihat halaman versi TPU untuk versi TPU Anda.
   ZONA	Lihat dokumen Region dan zona TPU untuk zona yang didukung.
   RUNTIME_VERSION	Untuk v5p, gunakan v2-alpha-tpuv5 untuk RUNTIME_VERSION.
   SERVICE_ACCOUNT	Alamat ini adalah alamat akun layanan yang dapat Anda temukan di Konsol Google Cloud -> IAM -> Akun Layanan. Misalnya: [e-mail]
   TPU_NAME	ID teks TPU yang ditetapkan pengguna yang dibuat saat permintaan resource dalam antrean dialokasikan.
   QUEUED_RESOURCE_ID	ID teks yang ditetapkan pengguna dari permintaan aset dalam antrean. Lihat dokumen Resource yang Diantrekan untuk informasi tentang resource yang diantrekan.
   QUOTA_TYPE	Dapat berupa reserved atau best-effort. Jika tidak satu pun yang ditentukan, QUOTA_TYPE akan ditetapkan secara default ke on-demand. Lihat quotas untuk mengetahui informasi tentang berbagai jenis kuota yang didukung oleh Cloud TPU.
   VALID_UNTIL_DURATION	Durasi saat permintaan valid. Lihat Resource yang diantrekan untuk mengetahui informasi tentang durasi valid yang berbeda.

2. Buat resource TPU:

   gcloud alpha 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} \
   --valid-until-duration ${VALID_UNTIL_DURATION} \
   --service-account ${SERVICE_ACCOUNT} \
   --${QUOTA_TYPE}

   Catatan: Flag QUOTA_TYPE dapat berupa reserved atau best-effort. Jika tidak satu pun dari opsi tersebut yang ditentukan, defaultnya adalah on-demand. Lihat quotas untuk mengetahui informasi tentang berbagai jenis kuota yang didukung oleh Cloud TPU.

   Anda dapat melakukan SSH ke VM TPU setelah resource yang diantrekan berada dalam status ACTIVE. Periksa status resource yang diantrekan dengan menjalankan perintah berikut:

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

   Ketika resource yang diantrekan berada dalam status ACTIVE, outputnya akan serupa dengan berikut ini:

   state: ACTIVE
   

3. Instal JAX dan dependensinya.

   # compatible with v5p: only jax version 0.4.19 and later \
   # jax 0.4.19 requires py 3.10 \
   
   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} --zone=${ZONE} --worker=all \
   --command='pip install "jax[tpu]==0.4.20" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'

4. Download repositori HuggingFace dan persyaratan penginstalan.

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

5. Melatih model

   Latih model dengan buffer yang telah dipetakan sebelumnya, berukuran 4 GB.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='export PATH=$PATH:$HOME/.local/bin && cd diffusers/examples/text_to_image && JAX_PLATFORMS=tpu,cpu python3 train_text_to_image_flax.py --pretrained_model_name_or_path=stabilityai/stable-diffusion-2-1 --dataset_name=lambdalabs/pokemon-blip-captions --resolution=256 --center_crop --random_flip --train_batch_size=1 --mixed_precision=bf16 --max_train_steps=150 --learning_rate=1e-05 --max_grad_norm=1 --output_dir=sd-pokemon-model --from_pt'

Pembersihan

Hapus TPU dan permintaan resource yang diantrekan di akhir sesi Anda atau untuk menghapus permintaan resource dalam antrean yang berada dalam status "FAILED". Untuk menghapus resource yang diantrekan, hapus slice tersebut, lalu permintaan 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

Atau, gunakan --force untuk menghapus slice dan permintaan resource yang diantrekan dalam satu langkah:

# With --force
gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID}
--project ${PROJECT_ID} --zone ${ZONE} --quiet --force

Hasil benchmark

Skrip pelatihan Difusi Stabil berjalan pada v5p-8, v5p-32, dan v5p-128. Tabel berikut menunjukkan throughput.

	V5p-8	v5p-32	v5p-128
Langkah Kereta	150	150	150
Ukuran tumpukan global	32	64	64
Throughput (contoh/dtk)	12,10	18,08	19,10

MaxText

Tutorial ini menunjukkan cara melatih model MaxText menggunakan set data sintetis di Cloud TPU.

MaxText adalah LLM open source berperforma tinggi, skalabel, dan teruji dengan baik yang ditulis dalam Python/JAX murni yang menargetkan Cloud TPU. MaxText memberdayakan peneliti dan developer dengan alat yang mudah diakses dan mudah disesuaikan untuk memajukan riset dan pengembangan Natural Language Processing (NLP).

Sebelum menjalankan tutorial ini, Anda perlu menyiapkan lingkungan Cloud TPU.

1. Menyiapkan variabel lingkungan

   export PROJECT_ID=your_project_ID
   export TPU_NAME=your_tpu_name # user defined TPU name
   export ACCELERATOR_TYPE=v5p-256
   export ZONE=us-east5-a
   export RUNTIME_VERSION=v2-alpha-tpuv5
   export RUN_NAME=your_experiment_run_name # user defined name for this run
   export GCS_BUCKET_NAME=your_bucket_name # Output cloud folder. Should start with gs://
   export MAXTEXT_OUTPUT_PATH=${GCS_BUCKET_NAME}/your_experiment_output_path
   export NUM_SLICES=1 # Update the value to a number >1 for Multislice.

   Deskripsi flag perintah

   Variabel	Deskripsi
   PROJECT_ID	Nama Project Google Cloud
   TPU_NAME	Nama yang ditetapkan pengguna untuk TPU Anda.
   ACCELERATOR_TYPE	Lihat halaman versi TPU untuk versi TPU Anda.
   ZONA	Lihat dokumen Region dan zona TPU untuk zona yang didukung.
   RUNTIME_VERSION	Untuk v5p, gunakan v2-alpha-tpuv5 untuk versi runtime.
   RUN_NAME	Nama operasi eksperimen yang disediakan pengguna.

   Penyiapan opsional yang direkomendasikan untuk Multislice:

   export NETWORK_NAME=your_network_name
   export FIREWALL_RULE_NAME=your_firewall_rule_name

   Jika Anda menjalankan workload Multislice dan menginginkan performa jaringan yang optimal, sebaiknya buat jaringan khusus dengan Maximum Transmission Unit (MTU) sebesar 8.896 byte dan konfigurasi aturan firewall yang sesuai. Meskipun bersifat opsional, langkah ini dapat meningkatkan performa secara signifikan, terutama saat meningkatkan jumlah slice melalui jaringan pusat data (DCN). Perlu diperhatikan bahwa pembuatan jaringan memerlukan izin compute.networks.create dalam project. Contoh berikut menunjukkan cara membuat aturan jaringan dan firewall khusus.

   Buat jaringan khusus:

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

   Buat aturan firewall:

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

2. Meng-clone repositori MaxText

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

3. Melatih model

   Bagian berikut menjelaskan dua opsi untuk melatih MaxText.

   Opsi 1

   Jika Anda ingin skrip untuk mengelola seluruh alur kerja, mulai dari menyediakan Cloud TPU dan menginstal dependensi, hingga menjalankan model dan menghilangkan resource, Anda dapat menggunakan multihost_job.py.

   cd maxtext && python3 multihost_job.py --PROJECT=${PROJECT_ID} --ZONE=${ZONE} \
   --NUM_SLICES=${NUM_SLICES} --TPU_TYPE=${ACCELERATOR_TYPE} \
   --VERSION=${RUNTIME_VERSION} --RUN_NAME=${RUN_NAME} #user defined run name \
   --BUCKET_NAME=${GCS_BUCKET_NAME} \ #used to store logs and configs
   --COMMAND="bash setup.sh && bash MaxText/configs/experimental/64b.sh RUN_NAME=${RUN_NAME} OUTPUT_PATH=${MAXTEXT_OUTPUT_PATH} PLATFORM=gce"

   Catatan: Tambahkan --CQR_EXTRA_ARGS="--network=${NETWORK_NAME}" jika Anda menggunakan jaringan kustom.

   Setelah memulai skrip, Anda akan melihat pesan yang mirip dengan yang berikut di log. Lokasi log dirujuk dalam pesan output. Klik link pertama untuk mengakses log semua pekerja setelah penyediaan TPU selesai.

   ------------------------------------
   
   multihost_job finished running, TPUs are starting up to run your job remotely.
   
   Logs for your job are displayed here:
   https://console.cloud.google.com/logs/query;query=resource.type%3D%22gce_instance%22%20AND%0Alog_id%2528%22_log%22%2529;?project=PROJECT_ID
   
   To see the output of a single host, you may edit the slice and worker
   number in the `log_file_path` property here:
   
   https://console.cloud.google.com/logs/query;query=resource.type%3D%22gce_instance%22%20AND%0Alog_id%2528%22RUN_NAME_log%22%2529%20AND%0Alabels.%22agent.googleapis.com%2Flog_file_path%22%3D%20%22%2FRUN_NAME%2Fmain_command_log_slice_0_worker_0%22;?project=PROJECT_ID
   
   When your job is finished, the main command log is in your Cloud Storage
   bucket:
   https://console.cloud.google.com/storage/browser/YOUR_BUCKET_NAME/RUN_NAME?project=PROJECT_ID
   
   View the status of the created TPUs using:
   gcloud compute tpus queued-resources list --filter=RUN_NAME
   --zone=ZONE --project=PROJECT_ID
   

Opsi 2

Untuk menjalankan skrip pelatihan beberapa kali di Cloud TPU yang disediakan, gunakan skrip multihost_runner.py untuk menggunakan resource.

1. Menyiapkan variabel untuk membuat TPU.

   export SERVICE_ACCOUNT=your_service_account
   export TPU_NAME=your_tpu_name
   export QUEUED_RESOURCE_ID=your_queued_resource_id
   export VALID_DURATION=1d
   export QUOTA_TYPE=quota_type

   Catatan: Flag berikut hanya diperlukan jika Anda menggunakan Multislice:

   --node-count ${NODE_COUNT} \
   --node-prefix ${NODE_PREFIX} # optional, the default is QUEUED_RESOURCE_ID

   Catatan: Jika Anda menentukan --node-count dan --node-prefix, jangan juga menentukan --node-id saat membuat resource TPU pada langkah berikut; error akan dibuat jika ketiga flag ditentukan.

2. Membuat resource TPU.

   gcloud alpha 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} \
   --valid-until-duration ${VALID_DURATION} \
   --service-account ${SERVICE_ACCOUNT} \
   --${QUOTA_TYPE}

   Catatan: Flag QUOTA_TYPE dapat berupa reserved atau best-effort. Jika tidak satu pun dari opsi tersebut yang ditentukan, defaultnya adalah on-demand. Lihat quotas untuk mengetahui informasi tentang berbagai jenis kuota yang didukung oleh Cloud TPU.

   Anda akan dapat terhubung ke VM TPU menggunakan SSH setelah QueuedResource dalam status ACTIVE:

   Gunakan perintah describe untuk membuat kueri status resource dalam antrean.

   gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}
   --project ${PROJECT_ID} --zone ${ZONE}

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

    state: ACTIVE
   

3. Menghubungkan ke TPU menggunakan SSH

   gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
     --project ${PROJECT_ID} \
     --zone ${ZONE}

4. Menginstal dependensi

   Catatan: Setelah menggunakan ssh untuk terhubung ke VM TPU, Anda perlu menetapkan variabel ${TPU_NAME} dan ${MAXTEXT_OUTPUT_PATH} untuk menggunakannya di VM TPU:

   export TPU_NAME=your_tpu_name
   export MAXTEXT_OUTPUT_PATH=output-path

   cd maxtext && python3 multihost_runner.py --TPU_PREFIX=${TPU_NAME} \
   --COMMAND='bash setup.sh'

5. Jalankan model dengan berbagai skrip konfigurasi, seperti 32b.sh, 64b.sh. Jika menjalankan skrip dari VM TPU, Anda perlu menambahkan flag --INTERNAL_IP=true.

   python3 multihost_runner.py --TPU_PREFIX=${TPU_NAME} \
   --COMMAND="bash MaxText/configs/experimental/64b.sh RUN_NAME=${RUN_NAME}
   OUTPUT_PATH=${MAXTEXT_OUTPUT_PATH} PLATFORM=gce"

   Catatan: Tetapkan --TPU_PREFIX=${NODE_PREFIX} untuk Multislice, bukan --TPU_PREFIX=${TPU_NAME}.

Pembersihan

Hapus TPU dan resource dalam antrean.

Hasil benchmark

Skrip pelatihan MaxText dijalankan dari 32 B hingga 1160 B dengan presisi bf16. Hasil operasi ini ditampilkan dalam tabel berikut.

Jumlah parameter	Jenis Akselerator	TFLOP/chip/dtk	Pemakaian model flop (MFU)
32 M	v5p-128	3,28E+02	71,47%
64 M	v5p-128	3,23E+02	70,31%
128 M	v5p-256	3,15E+02	68,68%
128 M	V5p-512	3,15E+02	68,53%
256 M	v5p-1024	3,16E+02	68,82%
512 M	v5p-1024	2,94E+02	63,99%
1.024 M	v5p-2048	2,49E+02	64,05%
1.024 M	v5p-4096	2,97E+02	64,80%
1.160 M	V5p-7680	2,95E+02	64,27%
1.160 M	v5p-12288	3,04E+02	66,23%

Model parameter 256B telah diuji pada v5p-512 dan v5p-1024 menggunakan bobot bf16 dan int8. Tabel berikut menampilkan hasil pengujian ini.

	V5p-512	V5p-512	v5p-1024	v5p-1024
Ukuran tumpukan global (token)	5,24E+05	5,24E+05	1,05E+06	1,05E+06
Presisi	bf16	int8	bf16	int8
TFLOP/chip/dtk	307	408	308	414
Pemakaian model flop (MFU)	66,98%	88,85%	67,09%	90,23%

Tutorial TensorFlow

Melatih ResNet pada v5p host tunggal

Tutorial ini menjelaskan cara melatih ImageNet di TPU v5p-8 menggunakan set data palsu. Jika Anda ingin menggunakan set data lain, lihat Menyiapkan set data.

Catatan: Sebelum menjalankan tutorial ini, Anda perlu menyiapkan lingkungan Cloud TPU.

Penyiapan

1. Buat variabel lingkungan:

   export PROJECT_ID=your-project-ID
   export ACCELERATOR_TYPE=v5p-8
   export ZONE=us-east1-c
   export RUNTIME_VERSION=tpu-vm-tf-2.16.1-pjrt
   export TPU_NAME=your-tpu-name
   export QUEUED_RESOURCE_ID=your-queued-resource-id
   export QUOTA_TYPE=quota-type

   Untuk tutorial ini, gunakan v5p-8 sebagai ACCELERATOR_TYPE.

2. Buat resource TPU:

   gcloud alpha 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} \
     --${QUOTA_TYPE}

   Catatan: Flag QUOTA_TYPE dapat berupa reserved atau best-effort. Jika tidak satu pun dari opsi tersebut yang ditentukan, defaultnya adalah on-demand. Lihat quotas untuk mengetahui informasi tentang berbagai jenis kuota yang didukung oleh Cloud TPU.

   Anda akan dapat terhubung ke VM TPU menggunakan SSH setelah resource yang diantrekan berada dalam status ACTIVE. Untuk memeriksa status resource yang diantrekan, gunakan perintah berikut:

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

3. Menghubungkan ke TPU menggunakan SSH

   gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
     --project ${PROJECT_ID} \
     --zone ${ZONE}

4. Menetapkan beberapa variabel lingkungan

   export MODELS_REPO=/usr/share/tpu/models
   export PYTHONPATH="${MODELS_REPO}:${PYTHONPATH}"
   export MODEL_DIR=gcp-directory-to-store-model
   export DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet
   export NEXT_PLUGGABLE_DEVICE_USE_C_API=true
   export TF_PLUGGABLE_DEVICE_LIBRARY_PATH=/lib/libtpu.so

5. Ubah ke direktori repositori model dan persyaratan penginstalan.

   cd ${MODELS_REPO} && git checkout r2.15.0
   pip install -r official/requirements.txt

Melatih model

1. Jalankan skrip pelatihan.

   python3 official/vision/train.py \
     --tpu=local \
     --experiment=resnet_imagenet \
     --mode=train_and_eval \
     --config_file=official/vision/configs/experiments/image_classification/imagenet_resnet50_tpu.yaml \
     --model_dir=${MODEL_DIR} \
     --params_override="runtime.distribution_strategy=tpu,task.train_data.input_path=${DATA_DIR}/train*,task.validation_data.input_path=${DATA_DIR}/validation*,task.train_data.global_batch_size=2048,task.validation_data.global_batch_size=2048,trainer.train_steps=100"

Pembersihan

Hapus TPU dan resource dalam antrean.

Melatih ResNet pada v5p multi-host

Tutorial ini menjelaskan cara melatih ImageNet pada v5p-16 atau yang lebih besar menggunakan set data palsu. Jika Anda ingin menggunakan set data yang berbeda, lihat bagian Menyiapkan set data.

Catatan: Sebelum menjalankan tutorial ini, Anda perlu menyiapkan lingkungan Cloud TPU.

1. Buat variabel lingkungan:

   export PROJECT_ID=your_project_ID
   export TPU_NAME=your_tpu_name
   export ZONE=us-east1-c
   export ACCELERATOR_TYPE=v5p-16
   export RUNTIME_VERSION=tpu-vm-tf-2.16.1-pod-pjrt
   export QUEUED_RESOURCE_ID=your-queued-resource-id
   export QUOTA_TYPE=quota-type

   ACCELERATOR_TYPE dapat v5p-16 atau lebih besar.

2. Buat resource TPU:

   gcloud alpha 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} \
     --${QUOTA_TYPE}

   Catatan: Flag QUOTA_TYPE dapat berupa reserved atau best-effort. Jika tidak satu pun dari opsi tersebut yang ditentukan, defaultnya adalah on-demand. Lihat quotas untuk mengetahui informasi tentang berbagai jenis kuota yang didukung oleh Cloud TPU.

   Anda akan dapat terhubung ke VM TPU menggunakan SSH setelah resource yang diantrekan berada dalam status ACTIVE.

   Gunakan perintah describe untuk membuat kueri status resource dalam antrean:

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

3. Menghubungkan ke TPU (pekerja nol) menggunakan SSH

   gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
     --project ${PROJECT_ID} \
     --zone ${ZONE}

4. Menetapkan beberapa variabel lingkungan

   export TPU_NAME=your_tpu_name
   export MODELS_REPO=/usr/share/tpu/models
   export PYTHONPATH="${MODELS_REPO}:${PYTHONPATH}"
   export MODEL_DIR=gcp-directory-to-store-model
   export DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet
   export TPU_LOAD_LIBRARY=0

5. Ubah ke direktori repositori model dan persyaratan penginstalan.

   cd $MODELS_REPO && git checkout r2.15.0
   pip install -r official/requirements.txt

Melatih model

1. Jalankan skrip pelatihan.

   python3 official/vision/train.py \
     --tpu=${TPU_NAME} \
     --experiment=resnet_imagenet \
     --mode=train_and_eval \
     --config_file=official/vision/configs/experiments/image_classification/imagenet_resnet50_tpu.yaml \
     --model_dir=${MODEL_DIR} \
     --params_override="runtime.distribution_strategy=tpu,task.train_data.input_path=${DATA_DIR}/train*,task.validation_data.input_path=${DATA_DIR}/validation*,task.train_data.global_batch_size=2048,task.validation_data.global_batch_size=2048,trainer.train_steps=100"

Pembersihan

Hapus TPU dan resource dalam antrean.

PyTorch/XLA

Llama 2

Tutorial ini akan membahas cara melatih model Llama 2 7B pada v5p menggunakan fork repositori HuggingFace pada PyTorch/XLA dengan Paralelisasi Umum dan Skalabel untuk Grafik Komputasi ML (GSPMD).

Penyiapan

1. Buat variabel untuk project ID, jenis akselerator, zona, versi runtime, dan nama TPU.

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5p-8
   export ZONE=us-east5-a
   export RUNTIME_VERSION=v2-alpha-tpuv5
   export SERVICE_ACCOUNT=your_service_account
   export TPU_NAME=your_tpu_name
   export QUEUED_RESOURCE_ID=your_queued_resource_id
   export QUOTA_TYPE=quota_type
   export VALID_DURATION=1d

2. Membuat resource TPU

   gcloud alpha 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} \
   --valid-until-duration ${VALID_DURATION} \
   --service-account ${SERVICE_ACCOUNT} \
   --${QUOTA_TYPE}

   Catatan: Flag QUOTA_TYPE dapat berupa reserved atau best-effort. Jika tidak satu pun dari opsi tersebut yang ditentukan, defaultnya adalah on-demand. Lihat quotas untuk mengetahui informasi tentang berbagai jenis kuota yang didukung oleh Cloud TPU.

   Anda akan dapat terhubung ke VM TPU menggunakan SSH setelah QueuedResource berada dalam status ACTIVE:

   Gunakan perintah describe untuk membuat kueri status resource dalam antrean.

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

   Ketika resource yang diantrekan berada dalam status AKTIF, output akan serupa dengan berikut:

    state: ACTIVE
   

   Catatan: Jika Anda menjalankan lingkungan multislice, gunakan gcloud compute tpus queued-resources ssh --node=all ..., bukan gcloud compute tpus tpu-vm ssh ... untuk perintah Anda. Tindakan ini akan memastikan semua chip TPU Anda digunakan dengan benar.

3. Instal Pytorch/XLA dan dependensi yang diperlukan.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
   --project ${PROJECT_ID} \
   --zone  ${ZONE} \
   --worker=all \
   --command='
   sudo apt-get update
   sudo apt-get install libopenblas-dev -y
   pip3 install numpy
   pip3 install typing-extensions
   pip install torch~=2.2.0 torch_xla[tpu]~=2.2.0 -f https://storage.googleapis.com/libtpu-releases/index.html
   '

4. Download repositori HuggingFace dan persyaratan penginstalan.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME}
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
   git clone -b llama2-google-next-training https://github.com/pytorch-tpu/transformers.git
   cd transformers
   pip3 install git+file://$PWD
   pip3 install datasets accelerate evaluate scikit-learn'

5. Download konfigurasi model 7B.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command="curl https://huggingface.co/TheBloke/Llama-2-7B-fp16/raw/main/config.json --output ~/config.json"

6. Melatih model

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
   export PJRT_DEVICE=TPU
   export XLA_USE_BF16=1
   export XLA_IR_DEBUG=1
   export XLA_HLO_DEBUG=1
   
   export LIBTPU_INIT_ARGS="--xla_enable_async_collective_permute=true
   --xla_tpu_enable_async_collective_fusion_multiple_steps=true
   --xla_tpu_enable_async_collective_fusion=true
   --xla_tpu_overlap_compute_collective_tc=true
   --xla_enable_async_all_gather=true
   --xla_jf_spmd_threshold_for_windowed_einsum_mib=0"
   
   export PROFILE_EPOCH=0
   export PROFILE_STEP=3
   export PROFILE_DURATION_MS=20000
   export PROFILE_LOGDIR=/tmp/home/
   
   cd transformers
   python examples/pytorch/language-modeling/run_clm.py \
    --tokenizer_name hf-internal-testing/llama-tokenizer \
    --dataset_name wikitext \
    --dataset_config_name wikitext-2-raw-v1 \
    --per_device_train_batch_size 96 \
    --per_device_eval_batch_size 8 \
    --num_train_epochs 1 \
    --do_train \
    --output_dir /tmp/output \
    --overwrite_output_dir \
    --config_name ~/config.json \
    --save_strategy no \
    --logging_strategy no \
    --remove_unused_columns no \
    --optim adafactor \
    --torch_dtype bfloat16 \
    --dataloader_drop_last yes \
    --block_size 2048 \
    --spmd_2d_sharding 1 \
    --spmd_grad_chkpt
   '

Jika menjalankan di lingkungan multislice, Anda perlu menetapkan flag --spmd_dcn_parallelism ke jumlah irisan.

SPMD_USER_GUIDE menyediakan panduan pengguna yang lebih mendalam yang menjelaskan semua variabel lingkungan dan tombol alih skrip HF. Perlu diperhatikan, LIBTPU_INIT_ARGS akan disertakan ke dalam PyTorch/XLA dan aktif secara default dalam rilis mendatang.

Pembersihan

Hapus TPU dan resource dalam antrean.

Hasil benchmark

Throughput untuk ketiga ukuran model Llama 2 disertakan dalam tabel berikut.

	V5p-8	v5p-128	v5p-128
Ukuran model	7 M	13 M	70 M
Ukuran tumpukan global	96	1024	128
Bentuk mesh sharding	(4, 1)	(64, 1)	(16, 4)
Pemakaian model flop (MFU)	56,67%	55,80%	51,85%

Dukungan dan Masukan

Kami menerima semua masukan! Untuk memberikan masukan atau meminta dukungan, isi formulir Dukungan atau Masukan Cloud TPU.