Panduan performa Cloud TPU
Langkah pertama saat memecahkan masalah performa TPU adalah membuat profil model Anda. Untuk mengetahui informasi selengkapnya tentang merekam profil performa, lihat Membuat profil model di Cloud TPU.
Performa model TPU
Bagian ini menjelaskan masalah umum yang dapat mengurangi performa model dan cara mengatasinya.
Model terikat input
TPU melakukan perhitungan dengan sangat cepat. Untuk memastikan TPU tidak dalam kondisi idle, penting untuk memastikan ada aliran data yang stabil yang dimuat ke TPU. Cara melakukannya bergantung pada cara Anda memuat dan memproses awal set data.
Misalnya, Anda dapat membaca file data secara paralel menggunakan tf.data.TFRecordset()
dan parameter num_parallel_reads.
Ukuran batch kecil karena sharding
Runtime TPU membagi batch di semua 8 core perangkat TPU (misalnya v2-8 atau v3-8). Jika Anda menentukan ukuran batch global 128, setiap core akan menerima ukuran batch 16 (128 / 8).
Untuk penggunaan memori yang optimal, gunakan ukuran batch terbesar yang sesuai dengan memori TPU. Setiap core TPU menggunakan register vektor dua dimensi 8 X 128 untuk memproses perkalian matriks. Secara umum, ukuran batch Anda harus dapat dibagi rata dengan 8 atau 128.
Penyesuaian pengelolaan memori
Anda dapat menggunakan variabel lingkungan terkait memori untuk menyempurnakan perilaku runtime tingkat rendah.
TPU_PREMAPPED_BUFFER_SIZE
TPU_PREMAPPED_BUFFER_SIZE menetapkan ukuran buffer memori host (dalam byte)
yang dipetakan dan disematkan sebelumnya untuk digunakan oleh runtime TPU untuk transfer data
(misalnya, DMA). Nilai defaultnya adalah 4294967296 byte. Nilai harus kelipatan 2^12 (4 KB = 4 * 1024 Byte = 4096 = 2^12).
Contoh berikut adalah nilai TPU_PRE_MAPPED_BUFFER_SIZE yang valid.
17179869184 = 2^34 = 2^22 * 2^12 (2^22 4KB pages will be premapped).
40000000000 = 5^10 * 2^12 = (5^10 4KB pages will be premapped).
Meningkatkan ukuran ini berpotensi meningkatkan performa transfer data antara host dan perangkat TPU, terutama untuk beban kerja dengan tensor besar atau komunikasi host-perangkat yang sering. Namun, hal ini juga meningkatkan jumlah memori host yang disematkan, sehingga mengurangi memori yang tersedia untuk proses lain.
Memecahkan masalah memori
Jika region buffer yang dipetakan sebelumnya tidak cukup besar
untuk mengalokasikan memori selama runtime program, workload akan gagal
dan menampilkan error RESOURCE_EXHAUSTED yang mirip dengan:
"Mengalokasikan buffer dari region yang telah dipetakan sebelumnya gagal dengan: RESOURCE_EXHAUSTED:
Mencoba mengalokasikan allocation_size. Hal itu tidak mungkin. Ada available_size gratis."
Jika buffer terlalu besar, inisialisasi TPU dapat memerlukan waktu lebih lama (berpotensi lebih dari 15 detik), sehingga TPU tampak macet.
Untuk mendiagnosisnya, periksa log runtime TPU. Log ini menjelaskan operasi yang sedang dilakukan, termasuk pra-pemetaan buffer. Anda dapat menemukan log di /tmp/tpu_logs/tpu_driver.INFO atau mencetaknya langsung ke konsol dengan menyetel variabel lingkungan TPU_STDERR_LOG_LEVEL=0. Setelan ini akan menghasilkan output yang mirip dengan:
I0604 12:45:24.926233 62136 tpu_hal.cc:214] Starting premapped memory manager initialization...
I0604 12:45:29.411218 62136 system.cc:1059] tpu::System initialized, current host id: 0, logical device ids: 0
I0604 12:45:29.411244 61600 tfrt_tpu_system_state.cc:216] CreateTpuSystemState: TPU initialization is successful and it took 5.583190661s
I0604 12:45:29.411267 61600 tfrt_tpu_system_state.cc:220] CreateTpuSystemState: using TPU host premapped buffer of size: 4294967296
Output ini akan memberi tahu Anda berapa lama waktu yang dibutuhkan untuk menginisialisasi TPU dan ukuran buffer yang dipetakan sebelumnya.
Menetapkan ukuran buffer
Jika buffer yang dipetakan sebelumnya terlalu kecil atau terlalu besar, Anda dapat menetapkan ukuran buffer secara manual menggunakan variabel lingkungan berikut.
TPU_PREMAPPED_BUFFER_SIZE: Menetapkan ukuran total (dalam byte) dari wilayah buffer yang dipetakan sebelumnya.TPU_PREMAPPED_BUFFER_TRANSFER_THRESHOLD_BYTES: Menetapkan ukuran maksimum buffer tunggal yang dapat dialokasikan dari wilayah yang dipetakan sebelumnya.
Misalnya, Anda dapat:
export TPU_PREMAPPED_BUFFER_SIZE=4294967296
untuk menetapkan ukuran buffer dan:
export TPU_PREMAPPED_BUFFER_TRANSFER_THRESHOLD_BYTES
untuk mengaktifkannya.
Ekspor ini menyetel ukuran ke default.
Sesuaikan nilai TPU_PREMAPPED_BUFFER_SIZE jika Anda
mencurigai transfer data host-perangkat
menjadi hambatan. Pantau penggunaan memori host dan performa model untuk menemukan keseimbangan yang optimal. Nilai default biasanya sudah cukup untuk sebagian besar kasus penggunaan.
Konfigurasi tcmalloc
Library tcmalloc digunakan secara default di VM Cloud TPU untuk meningkatkan
performa model dengan alokasi memori yang cukup besar dan sering. Hal 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
Pengoptimalan performa jaringan
Bagian berikut menjelaskan cara mengoptimalkan performa jaringan dengan mengonfigurasi unit transmisi maksimum (MTU) dan menggunakan multi-NIC untuk lingkungan Multislice.
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 (kemungkinan maksimum). 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.
Menggunakan opsi multi-NIC untuk Multislice
Saat melatih model besar di lingkungan Multislice yang terdiri dari ribuan chip TPU, komunikasi antar-slice melalui jaringan pusat data (DCN) dapat menjadi hambatan. Untuk meningkatkan bandwidth jaringan untuk workload yang terikat jaringan, Anda dapat menggunakan multi-NIC untuk meningkatkan jumlah antarmuka jaringan di VM TPU. Saat Anda menggunakan multi-NIC, setiap VM TPU dialokasikan antarmuka jaringan tambahan, yang masing-masing terhubung ke jaringan VPC yang unik, sehingga meningkatkan throughput jaringan secara keseluruhan. NIC tambahan harus berada dalam rentang IP yang saling eksklusif.
Untuk mengetahui informasi selengkapnya tentang cara mengaktifkan multi-jaringan saat menggunakan Google Kubernetes Engine (GKE), lihat Meningkatkan performa jaringan tanpa hostNetwork di TPU Trillium atau Ironwood (TPU7x). Untuk contoh penggunaan multi-NIC dengan XPK, lihat Membuat cluster dengan dukungan multi-NIC menggunakan XPK.
Pengoptimalan compiler XLA
XLA adalah compiler untuk machine learning yang dapat menghasilkan biner untuk TPU, CPU, GPU, dan platform lainnya. Meskipun XLA adalah bagian dari codebase TensorFlow standar, XLA juga dapat digunakan pada model PyTorch dan JAX. Model untuk Cloud TPU diterjemahkan ke grafik XLA, yang kemudian dikompilasi XLA menjadi file yang dapat dieksekusi TPU. Untuk mengetahui informasi selengkapnya tentang XLA, lihat XLA: Pengoptimalan Compiler untuk Machine Learning.
Padding
Untuk menggunakan memori TPU secara efisien, susun data Anda sehingga dapat dibagi menjadi chunk 128 x 8. Jika data untuk komputasi matriks tidak mengisi seluruh potongan 128 x 8, compiler XLA akan mengisi tensor. Ada dua kelemahan padding:
- Tensor yang diberi padding kurang memanfaatkan inti TPU.
- Padding meningkatkan jumlah penyimpanan memori on-chip yang diperlukan untuk tensor dan dapat menyebabkan error kehabisan memori.
Meskipun padding dilakukan secara otomatis oleh compiler XLA jika diperlukan, Anda dapat menentukan jumlah padding yang dilakukan menggunakan alat penampil memori. Anda dapat menghindari padding dengan memilih dimensi tensor yang cocok untuk TPU.
Dimensi tensor
Untuk mencapai FLOP puncak, dimensi perkalian matriks harus lebih besar daripada ukuran MXU untuk versi TPU yang Anda gunakan. Ukuran MXU adalah 256 x 256 untuk v6e dan 128 x 128 untuk versi sebelum v6e. Untuk mengetahui informasi selengkapnya, lihat Arsitektur sistem Cloud TPU.
Ukuran batch
Compiler XLA membulatkan ukuran tensor yang disimpan dalam memori HBM TPU untuk melakukan komputasi secara lebih efisien. Pengisihan ini terjadi secara transparan di tingkat hardware dan tidak memengaruhi hasil. Namun, dalam kasus tertentu, padding dapat menyebabkan peningkatan penggunaan memori dan waktu eksekusi yang signifikan.
Runtime TPU menata tensor dalam memori untuk memaksimalkan efisiensi komputasi dan meminimalkan padding. Untuk meminimalkan overhead memori dan memaksimalkan efisiensi komputasi, salah satu kondisi berikut harus terpenuhi:
Ukuran batch total harus kelipatan 64 (8 per core TPU), dan ukuran dimensi fitur harus kelipatan 128.
Ukuran batch total harus kelipatan 1024 (128 per core TPU), dan ukuran dimensi fitur harus kelipatan 8.
Menggunakan ukuran batch 1024 dan dimensi fitur yang merupakan kelipatan 128 akan menghasilkan efisiensi terbaik, meskipun hal ini mungkin tidak dapat dilakukan untuk semua model.
Fusion
Penggabungan (Fusion) adalah teknik umum yang digunakan compiler XLA untuk mengoptimalkan program. Operasi gabungan adalah kombinasi dari beberapa operasi konstituen yang akan dijalankan bersama.
Misalnya, perhatikan serangkaian operasi berikut:
tmp = tf.add(x, y)
result = tf.multiply(tmp, z)
Kode ini kira-kira setara dengan pseudo-kode berikut:
for (i = 0; i < element_count; i++) {
tmp[i] = x[i] + y[i];
}
for (i = 0; i < element_count; i++) {
result[i] = tmp[i] * z[i];
}
Dengan penggabungan, akses array terjadi secara bersamaan:
for (i = 0; i < element_count; i++) {
result[i] = (x[i] + y[i]) * z[i];
}
Dalam contoh ini, jumlah perjalanan pulang pergi memori berkurang dan XLA tidak perlu mengalokasikan ruang untuk 'tmp'.
Penggabungan adalah pengoptimalan penting dan memberikan manfaat bagi Cloud TPU dalam beberapa cara:
- Hal ini mengurangi transfer memori dengan menghilangkan kebutuhan untuk menyimpan hasil sementara di memori utama, yang lambat.
- Hal ini memungkinkan pemanfaatan unit hardware yang lebih besar yang jika tidak, akan tidak digunakan.
- Hal ini dapat mengurangi pemanfaatan memori model karena lebih sedikit buffer yang perlu aktif secara bersamaan.
Penyiaran
Penyiaran terjadi secara implisit saat dua tensor dengan bentuk yang berbeda, tetapi kompatibel, digabungkan.
Misalnya, tf.add(vector, matrix) mengharuskan vektor disiarkan ke
bentuk matriks. Hasil operasi memiliki bentuk yang sama dengan matriks. Untuk mengetahui detail selengkapnya, lihat panduan untuk
penyiaran array.
Meskipun siaran sering kali dapat digabungkan dengan konsumennya, memaksa siaran dapat mengakibatkan performa yang buruk dan peningkatan penggunaan memori.
Dalam contoh berikut, siaran implisit dalam penambahan vektor dan matriks tidak dapat digabungkan dengan argmax yang menghasilkan siaran yang diwujudkan:
`tf.argmax(tf.add(vector, zero_matrix), axis=0)`
Rekomendasi performa untuk arsitektur dual-chiplet Ironwood
Model pemrograman Ironwood memungkinkan Anda mengakses dua perangkat TPU, bukan arsitektur inti logis tunggal (juga dikenal sebagai MegaCore) yang digunakan pada generasi sebelumnya (TPU v4 dan v5p). Perubahan ini meningkatkan efektivitas biaya dan efisiensi pembuatan chip. Meskipun hal ini merupakan perubahan arsitektur, desain baru memastikan Anda dapat menggunakan kembali model software yang ada dengan perubahan minimal.
Untuk mendapatkan performa terbaik dengan arsitektur dual-chiplet, sebaiknya gunakan pendekatan berikut:
Menggunakan paralelisme tensor di seluruh chiplet: Antarmuka D2D bandwidth tinggi dirancang untuk paralelisme tensor yang efisien. Sebaiknya pisahkan tensor di kedua perangkat on-chip.
Memanfaatkan kolektif hierarkis: Untuk memaksimalkan efisiensi komunikasi, manfaatkan hierarki jaringan dua tingkat: link D2D ultra-cepat antara chiplet on-chip dan link ICI cepat dalam slice. Saat menggunakan paralelisme otomatis dengan SPMD (single program, multiple data), compiler XLA akan menanganinya untuk Anda dengan otomatis membuat operasi kolektif hierarkis. Saat mempartisi model secara manual, Anda juga harus mendesain pola komunikasi di sekitar hierarki ini. Membuat prioritas komunikasi antara dua perangkat pada chip yang sama sebelum berkomunikasi dengan perangkat pada chip lain.
Tumpang-tindihkan komunikasi dengan komputasi: Untuk memaksimalkan pemanfaatan hardware, lakukan offload operasi komunikasi kolektif, seperti all-reduce, ke SparseCore. Operasi ini, yang tidak terikat ke unit perkalian matriks (MXU), dapat dieksekusi di SparseCores secara bersamaan sementara TensorCores melanjutkan komputasinya. Teknik ini dapat memulihkan beberapa manfaat performa yang melekat pada operasi gabungan dalam arsitektur MegaCore sebelumnya.
Offload ke SparseCore untuk penyematan: Dalam desain dual-chiplet, tabel penyematan dapat dipartisi di seluruh HBM kedua chiplet. Untuk menghindari penurunan performa akibat kurangnya berbagi memori ini, pindahkan operasi pengumpulan sematan ke SparseCore. Strategi ini memanfaatkan interkoneksi D2D berkecepatan tinggi untuk mentransfer vektor sematan secara efisien antar-chiplet. Untuk mengetahui informasi selengkapnya tentang SparseCore dan model embedding, lihat Pembahasan mendalam tentang SparseCore untuk Model Embedding Besar (LEM).
Untuk mengetahui informasi selengkapnya tentang arsitektur Ironwood di TPU7x, lihat TPU7x (Ironwood).