Pelatihan Pod TPU

Ringkasan

TPU didesain untuk diskalakan menjadi Pod TPU. Pod TPU adalah kumpulan perangkat TPU yang terhubung melalui antarmuka jaringan khusus berkecepatan tinggi. Dengan Pod TPU, Anda dapat mendistribusikan beban pemrosesan di beberapa TPU. Setiap board TPU terhubung ke mesin host berbasis CPU berperforma tinggi untuk hal-hal seperti memuat dan memproses data terlebih dahulu. Untuk memanfaatkan sepenuhnya jumlah TPU yang lebih besar, Anda harus menyesuaikan beberapa parameter tugas pelatihan.

Penyiapan untuk pelatihan dengan Pod TPU berbeda untuk setiap framework. Gunakan link berikut untuk melihat informasi mendetail tentang pelatihan Pod dengan setiap framework:

• TensorFlow

• PyTorch

• JAX

Bagian berikut menjelaskan beberapa masalah umum, perubahan yang perlu dilakukan pada model, dan praktik terbaik untuk mengurangi atau menghindari kegagalan Pod.

Menskalakan ukuran tumpukan dan melatih langkah

Untuk mencapai penskalaan linier pada jenis TPU yang lebih besar, pertahankan ukuran batch per core tetap sama.

Misalnya, jika Anda menggunakan ukuran tumpukan 1024 pada v2-8, gunakan ukuran tumpukan 4096 (4 * 1024) pada v2-32. Perangkat keras ini sepenuhnya menggunakan perangkat keras TPU. Anda dapat menggunakan ukuran tumpukan yang lebih kecil, tetapi pelatihan Anda tidak akan diskalakan secara linear jika Anda melakukannya.

Beberapa model menyertakan tanda train_steps yang satu langkah berkaitan dengan pemrosesan satu batch data. Saat Anda meningkatkan ukuran batch, perkecil jumlah langkah pelatihan sehingga jumlah total contoh pelatihan tetap sama.

Misalnya, jika Anda memiliki ukuran batch 1.000 untuk 100 langkah, 100.000 contoh akan diproses selama pelatihan. Jika sekarang Anda memiliki 4 pekerja dan ukuran batch efektif 4.000, Anda harus menyesuaikan jumlah langkah menjadi 25 untuk memproses 100.000 contoh yang sama. Jika model Anda menggunakan flag epochs, Anda tidak perlu menskalakan jumlah langkah.

Ukuran tumpukan yang lebih besar dapat mengubah perilaku konvergensi model, sehingga Anda juga dapat menyesuaikan beberapa hyperparameter, seperti kecepatan pemelajaran.

Menggunakan bucket Google Cloud Storage regional di region yang sama dengan Pod TPU

Secara umum, praktik terbaik untuk pelatihan TPU adalah selalu menggunakan resource di region yang sama. Region resource sangat penting saat menggunakan Pod TPU karena kecepatan transfer data akan lebih tinggi jika bucket Google Cloud Storage dan TPU Anda berada di region yang sama.

Pastikan Anda menggunakan bucket Google Cloud Storage regional di region yang sama dengan TPU untuk set data pelatihan dan checkpoint.

Praktik terbaik alur kerja untuk pengembangan di TPU Pod

Saat mengembangkan workload TPU baru, sebaiknya mulai pengembangan pada TPU terkecil dan secara bertahap melakukan iterasi ke ukuran TPU yang lebih besar. Mulailah dengan menggunakan versi TPU kecil (misalnya, v2-8 atau v3-8).

• Menguji fungsionalitas beban kerja Anda
• Menguji dan memvalidasi performa menggunakan alat performa

Setelah beban kerja Anda berfungsi dan mencapai target performa, tingkatkan skala ke jenis TPU yang lebih besar seperti v2-32 atau v3-32. Tingkatkan ukuran TPU secara bertahap dan berulang sambil memvalidasi skalabilitas (fungsi dan performa) hingga Anda mencapai ukuran TPU yang diinginkan.