Halaman ini diterjemahkan oleh Cloud Translation API.

Pelatihan Cloud TPU v5p

Cloud TPU v5p adalah Cloud TPU generasi kelima Google Cloud dan pengganti TPU v4. v5p dioptimalkan untuk pelatihan berskala besar dan menjadi platform terkemuka untuk pengembangan LLM dasar, model difusi, dan AI generatif. Pada tingkat tinggi, v5p memberikan performa hingga 2x lipat dari v4, sekaligus memuat TPU 2x lebih banyak ke dalam Pod (6 ribu slice terbesar versus 3 ribu di v4), sehingga meningkatkan performa hingga 4x lipat di tingkat Pod. Chip ini juga berjalan pada frekuensi clock yang lebih tinggi (1,75 Ghz vs. 1,05 Ghz), menambahkan SparseCore untuk penyematan skala besar, dan melipatgandakan kapasitas Memori Bandwidth Tinggi (HBM).

Konsep Cloud TPU v5p

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

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

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

Mengelola resource TPU

Semua perintah dalam dokumen ini mengasumsikan bahwa Anda membuat VM TPU v5p. Untuk mengetahui informasi selengkapnya tentang perintah untuk membuat VM TPU, lihat Mengelola TPU atau Panduan pengguna resource dalam antrean untuk mengelola resource dalam antrean. Agar perintah lebih mudah dijalankan, contoh kode dalam dokumen ini menggunakan variabel lingkungan berikut:

export PROJECT_ID=your-project
export ACCELERATOR_TYPE=v5p-8
export ZONE=us-east5-a
export RUNTIME_VERSION=v2-alpha-tpuv5
export TPU_NAME=your-tpu-name

Deskripsi variabel lingkungan

PROJECT_ID: Project Google Cloud tempat Anda membuat 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.
ZONE: Zona tempat Anda berencana membuat Cloud TPU.
RUNTIME_VERSION: Versi software TPU.
TPU_NAME: Nama TPU yang ditentukan pengguna yang Anda gunakan.

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 lebih besar dari 4 chip, Anda akan memiliki beberapa VM dalam satu slice. Dalam hal ini, Anda perlu menggunakan flag --worker=all untuk menjalankan penginstalan di semua VM TPU menggunakan satu perintah:

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'

Anda dapat menjalankan perintah berikut untuk memeriksa jumlah perangkat (output yang ditampilkan di sini dihasilkan dengan slice v5p-32). 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 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 slice yang diberikan. jax.local_device_count() menunjukkan jumlah chip yang dapat diakses oleh satu VM dalam 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 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 di 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
pip install numpy
pip install torch torch_xla[tpu]~=2.5.0 -f https://storage.googleapis.com/libtpu-releases/index.html
pip3 install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cpu
'

Gunakan skrip Python dengan PJRT untuk memvalidasi penginstalan Anda. Skrip ini menampilkan perangkat TPU yang tersedia (output yang ditampilkan di sini dihasilkan dengan slice 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 dalam dokumen ini untuk memulai pelatihan v5p menggunakan PyTorch.

Memantau dan membuat profil

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

Tutorial pelatihan

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

Diffusion 2.1
MaxText
ResNet di satu host v5p
ResNet di v5p multi-host

Tutorial JAX

Train Diffusion 2.1

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

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

Siapkan

Buat variabel lingkungan:

export GCS_BUCKET_NAME=your-bucket
export PROJECT_ID=your-project-ID
export ACCELERATOR_TYPE=v5p-32
export ZONE=europe-west4-b
export LOCATION=europe-west4
export RUNTIME_VERSION=v2-alpha-tpuv5
export SERVICE_ACCOUNT=your-service-account
export TPU_NAME=your-tpu-name
export QUEUED_RESOURCE_ID=your-qr-name
export QUOTA_TYPE=spot
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 mengetahui zona yang didukung.
LOKASI	Region Google Cloud tempat membuat bucket penyimpanan Cloud Storage.
RUNTIME_VERSION	Untuk v5p, gunakan v2-alpha-tpuv5 untuk RUNTIME_VERSION.
SERVICE_ACCOUNT	Ini adalah alamat akun layanan yang dapat Anda temukan di konsol Google Cloud -> IAM -> Service Accounts. Misalnya: tpu-service-account@myprojectID.iam.gserviceaccount.com
TPU_NAME	ID teks TPU yang ditetapkan pengguna yang dibuat saat permintaan resource yang diantrekan dialokasikan.
QUEUED_RESOURCE_ID	ID teks yang ditetapkan pengguna untuk permintaan resource yang diantrekan. Lihat dokumen Resource dalam Antrean untuk mengetahui informasi tentang resource dalam antrean.
QUOTA_TYPE	Dapat berupa `reserved` atau `spot`. Jika tidak ada yang ditentukan, QUOTA_TYPE akan ditetapkan secara default ke `on-demand`. Lihat kuota untuk mengetahui informasi tentang berbagai jenis kuota yang didukung oleh Cloud TPU.
VALID_UNTIL_DURATION	Durasi validitas permintaan. Lihat Resource dalam antrean untuk mengetahui informasi tentang berbagai durasi yang valid.

Siapkan bucket penyimpanan untuk output model Anda.

gcloud storage buckets create gs://$GCS_BUCKET_NAME \
 --project=$PROJECT_ID \
 --location=$LOCATION

Buat 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} \
--valid-until-duration ${VALID_UNTIL_DURATION} \
--service-account ${SERVICE_ACCOUNT} \
--${QUOTA_TYPE}
```
Catatan: Flag QUOTA_TYPE dapat berupa reserved atau spot. Jika tidak ada yang ditentukan, defaultnya adalah on-demand. Lihat kuota untuk mengetahui informasi tentang berbagai jenis kuota yang didukung oleh Cloud TPU.

Anda akan dapat menggunakan 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}
```
Jika resource yang diantrean berada dalam status ACTIVE, outputnya akan mirip dengan berikut ini:
```
state: ACTIVE
```

Melatih model

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
git reset --hard 57629bcf4fa32fe5a57096b60b09f41f2fa5c35d # This identifies the GitHub commit to use.
pip3 install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html # Install the latest version of JAX
pip3 install -r requirements.txt
pip3 install .
export LIBTPU_INIT_ARGS=""
python -m src.maxdiffusion.models.train src/maxdiffusion/configs/base_2_base.yml run_name=my_run base_output_directory=gs://$GCS_BUCKET_NAME enable_profiler=False"

Pembersihan

Menghapus TPU dan permintaan resource yang diantrekan di akhir sesi atau untuk menghapus permintaan resource yang diantrekan yang berada dalam status "GAGAL". Untuk menghapus resource yang diantrekan, hapus slice, 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 dalam antrean 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 Stable Diffusion berjalan di v5p-8, v5p-32, dan v5p-128. Tabel berikut menunjukkan throughput.

	v5p-8	v5p-32	v5p-128
Langkah Latihan	150	150	150
Ukuran batch global	32	64	64
Throughput (contoh/dtk)	12.10	18,08	19.10

MaxText

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

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

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

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 ditentukan pengguna untuk TPU Anda.
ACCELERATOR_TYPE	Lihat halaman versi TPU untuk versi TPU Anda.
ZONA	Lihat dokumen Region dan zona TPU untuk mengetahui zona yang didukung.
RUNTIME_VERSION	Untuk v5p, gunakan v2-alpha-tpuv5 untuk versi runtime.
RUN_NAME	Nama eksperimen yang dijalankan yang diberikan pengguna.

Penyiapan opsional yang direkomendasikan untuk Multislice:

export NETWORK_NAME=your_network_name
export FIREWALL_RULE_NAME=your_firewall_rule_name

Jika Anda menjalankan beban kerja Multislice dan menginginkan performa jaringan yang optimal, pertimbangkan untuk membuat jaringan khusus dengan Unit Transmisi Maksimum (MTU) 8896 byte dan mengonfigurasi aturan firewall yang sesuai. Meskipun opsional, langkah ini dapat meningkatkan performa secara signifikan, terutama saat menskalakan jumlah slice melalui jaringan pusat data (DCN). Perhatikan 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}

Meng-clone repositori MaxText

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

Melatih model

Bagian berikut menjelaskan dua opsi untuk melatih MaxText.

Opsi 1

Jika Anda ingin skrip mengelola seluruh alur kerja, mulai dari menyediakan Cloud TPU dan menginstal dependensi hingga menjalankan model dan menghapus 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"

Setelah memulai skrip, Anda akan melihat pesan yang mirip dengan berikut ini dalam 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.

Siapkan 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

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

Buat 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} \
--valid-until-duration ${VALID_DURATION} \
--service-account ${SERVICE_ACCOUNT} \
--${QUOTA_TYPE}
```
Catatan: Flag QUOTA_TYPE dapat berupa reserved atau spot. Jika tidak ada yang ditentukan, defaultnya adalah on-demand. Lihat kuota 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 yang diantrekan.
```
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 serupa dengan berikut:
```
 state: ACTIVE
```

Menghubungkan ke TPU menggunakan SSH

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

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'
```
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 yang diantrekan.

Hasil benchmark

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

Jumlah parameter	Jenis Akselerator	TFLOP/chip/dtk	Penggunaan flops model (MFU)
32B	v5p-128	3,28E+02	71,47%
64B	v5p-128	3,23E+02	70,31%
128B	v5p-256	3,15E+02	68,68%
128B	v5p-512	3,15E+02	68,53%
256B	v5p-1024	3,16E+02	68,82%
512B	v5p-1024	2,94E+02	63,99%
1024B	v5p-2048	2,49E+02	64,05%
1024B	v5p-4096	2,97E+02	64,80%
1160B	v5p-7680	2,95E+02	64,27%
1160B	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 batch 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
Penggunaan flops model (MFU)	66,98%	88,85%	67,09%	90,23%

Tutorial TensorFlow

Melatih ResNet di satu host v5p

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

Siapkan

Buat variabel lingkungan:

export PROJECT_ID=your-project-ID
export ACCELERATOR_TYPE=v5p-32
export ZONE=us-east1-c
export RUNTIME_VERSION=tpu-vm-tf-2.18.0-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.

Buat 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} \
  --${QUOTA_TYPE}
```
Catatan: Flag QUOTA_TYPE dapat berupa reserved atau spot. Jika tidak ada yang ditentukan, defaultnya adalah on-demand. Lihat kuota untuk mengetahui informasi tentang berbagai jenis kuota yang didukung oleh Cloud TPU.

Anda akan dapat terhubung ke VM TPU menggunakan SSH setelah resource dalam antrean 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}
```

Menghubungkan ke TPU menggunakan SSH

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

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

Ubah ke direktori repositori model dan instal persyaratan.

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

Melatih model

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 yang diantrekan.

Melatih ResNet di v5p multi-host

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

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.18.0-pod-pjrt
export QUEUED_RESOURCE_ID=your-queued-resource-id
export QUOTA_TYPE=quota-type

ACCELERATOR_TYPE dapat berupa v5p-16 atau lebih besar.

Buat 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} \
  --${QUOTA_TYPE}
```
Catatan: Flag QUOTA_TYPE dapat berupa reserved atau spot. Jika tidak ada yang ditentukan, defaultnya adalah on-demand. Lihat kuota untuk mengetahui informasi tentang berbagai jenis kuota yang didukung oleh Cloud TPU.

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

Gunakan perintah describe untuk membuat kueri status resource yang diantrekan:
```
gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
  --project ${PROJECT_ID} \
  --zone ${ZONE}
```

Menghubungkan ke TPU (pekerja nol) menggunakan SSH

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

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

Ubah ke direktori repositori model dan instal persyaratan.

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

Melatih model

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 yang diantrekan.

PyTorch/XLA

Llama 2

Tutorial ini akan membahas cara melatih model Llama 2 7B di v5p menggunakan fork repositori HuggingFace di PyTorch/XLA dengan Paralelisasi General dan Scalable untuk ML Computation Graphs (GSPMD).

Penyiapan

Buat variabel lingkungan.

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

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} \
--valid-until-duration ${VALID_DURATION} \
--service-account ${SERVICE_ACCOUNT} \
--${QUOTA_TYPE}
```
Catatan: Flag QUOTA_TYPE dapat berupa reserved atau spot. Jika tidak ada yang ditentukan, defaultnya adalah on-demand. Lihat kuota 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 yang diantrekan.
```
gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
--project ${PROJECT_ID} \
--zone ${ZONE}
```
Jika resource yang diantrekan berada dalam status AKTIF, outputnya akan mirip dengan berikut ini:
```
 state: ACTIVE
```
Catatan: Jika Anda menjalankan di lingkungan multislice, gunakan gcloud compute tpus queued-resources ssh --node=all ..., bukan gcloud compute tpus tpu-vm ssh ... untuk perintah Anda. Hal ini memastikan semua chip TPU Anda digunakan dengan benar.

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 torch_xla[tpu]~=2.5.0 -f https://storage.googleapis.com/libtpu-releases/index.html
'

Download repositori HuggingFace dan instal persyaratan.

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'

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"

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 Anda menjalankan di lingkungan multislice, Anda perlu menetapkan flag --spmd_dcn_parallelism ke jumlah slice.

SPMD_USER_GUIDE menyediakan panduan pengguna yang lebih mendalam yang menjelaskan semua variabel lingkungan dan tombol yang berbeda dari skrip HF. Perlu diperhatikan bahwa LIBTPU_INIT_ARGS akan digabungkan ke dalam PyTorch/XLA dan diaktifkan secara default dalam rilis mendatang.

Pembersihan

Hapus TPU dan resource yang diantrekan.

Hasil benchmark

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

	v5p-8	v5p-128	v5p-128
Ukuran model	7B	13B	70 M
Ukuran batch global	96	1024	128
Bentuk mesh sharding	(4, 1)	(64, 1)	(16, 4)
Penggunaan flops model (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.