Panduan performa Cloud TPU

Langkah pertama saat memecahkan masalah performa TPU adalah membuat profil model. Untuk 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.

1. Model terikat input

   TPU melakukan penghitungan dengan sangat cepat. Untuk memastikan TPU tidak nonaktif, penting untuk memastikan ada aliran data yang stabil yang dimuat ke TPU. Bagaimana hal ini dilakukan tergantung pada bagaimana Anda memuat dan melakukan pra-pemrosesan {i>dataset<i} Anda. Misalnya, Anda dapat membaca file data secara paralel menggunakan tf.data.TFRecordset() dan parameter num_parallel_reads.

2. Ukuran tumpukan terlalu kecil karena sharding (membagi batch di antara inti)

   Runtime TPU membagi batch di 8 core perangkat TPU (misalnya v2-8 atau v3-8). Jika Anda menentukan ukuran tumpukan global 128, setiap inti akan menerima ukuran tumpukan 16 (128 / 8).

   Untuk penggunaan memori yang optimal, gunakan ukuran tumpukan terbesar yang sesuai dengan memori TPU. Setiap inti TPU menggunakan register vektor berukuran 8 X 128 dua dimensi untuk memproses perkalian matriks. Secara umum, ukuran tumpukan harus dapat dibagi secara merata dengan 8 atau 128.

Pengoptimalan compiler XLA

XLA adalah compiler untuk machine learning yang dapat menghasilkan biner untuk TPU, CPU, GPU, dan platform lainnya. Meskipun merupakan bagian dari code base TensorFlow standar, XLA juga dapat digunakan pada model PyTorch dan JAX. Model untuk Cloud TPU diterjemahkan ke grafik XLA, yang kemudian dikompilasi XLA ke dalam TPU yang dapat dieksekusi. Untuk informasi selengkapnya tentang XLA, lihat XLA: Optimizing Compiler for Machine Learning.

Padding

Untuk menggunakan memori TPU secara efisien, susun data Anda agar dapat dikelompokkan menjadi potongan 128x8. Jika data untuk komputasi matriks tidak mengisi seluruh potongan 128 x 8, maka compiler XLA akan melakukan pad pada TensorFlow. Ada dua kelemahan dari padding:

1. Tensor dengan padding kurang memanfaatkan inti TPU.
2. Padding meningkatkan jumlah penyimpanan memori pada chip yang diperlukan untuk tensor dan dapat menyebabkan error kehabisan memori.

Meskipun padding otomatis dijalankan oleh compiler XLA saat diperlukan, Anda dapat menentukan jumlah padding yang dilakukan menggunakan alat penampil memori. Anda dapat menghindari padding dengan memilih dimensi TensorFlow yang cocok untuk TPU.

Dimensi Tensor

Compiler XLA membulatkan ukuran TensorFlow yang disimpan dalam memori TPU HBM untuk melakukan komputasi dengan lebih efisien. Padding ini terjadi secara transparan di tingkat hardware dan tidak memengaruhi hasil. Namun, dalam kasus tertentu, padding dapat mengakibatkan peningkatan penggunaan memori dan waktu eksekusi secara signifikan.

Runtime TPU menjabarkan TensorFlow di memori untuk memaksimalkan efisiensi komputasi dan meminimalkan padding. Untuk meminimalkan overhead memori dan memaksimalkan efisiensi komputasi, salah satu hal berikut harus terpenuhi:

1. Total ukuran tumpukan harus kelipatan 64 (8 per inti TPU), dan ukuran dimensi fitur harus kelipatan 128.

2. Total ukuran tumpukan harus kelipatan 1.024 (128 per inti TPU), dan ukuran dimensi fitur harus kelipatan 8.

Menggunakan ukuran tumpukan 1024 dan dimensi fitur yang merupakan kelipatan 128 akan menghasilkan efisiensi terbaik, meskipun hal ini mungkin tidak dapat dilakukan untuk semua model.

Catatan: Dimensi fitur mengacu pada ukuran tersembunyi dari lapisan yang terhubung sepenuhnya atau jumlah saluran output dalam konvolusi. Tidak semua lapisan dapat mematuhi aturan ini, terutama lapisan pertama dan terakhir dari jaringan. Ini tidak masalah. Sebagian besar model memerlukan sejumlah padding.

Penggabungan

Fusion adalah teknik umum yang digunakan compiler XLA untuk mengoptimalkan program. Operasi gabungan adalah kombinasi dari beberapa operasi konstituen yang akan dijalankan secara bersamaan.

Misalnya, pertimbangkan rangkaian operasi berikut:

    tmp = tf.add(x, y)
    result = tf.multiply(tmp, z)

Kode ini kurang lebih sama dengan kode pseudo berikut:

    for (i = 0; i < element_count; i++) {
      tmp[i] = x[i] + y[i];
    }

    for (i = 0; i < element_count; i++) {
      result = tmp[i] * z[i];
    }

Dengan fusi, akses array terjadi secara bersamaan:

    for (i = 0; i < element_count; i++) {
      result = (x[i] + y[i]) * z[i];
    }

Dalam contoh ini, jumlah round-trip memori dikurangi dan XLA tidak perlu mengalokasikan ruang untuk 'tmp'.

Fusion adalah pengoptimalan penting dan menguntungkan Cloud TPU dalam beberapa cara:

• Ini mengurangi transfer memori dengan meniadakan kebutuhan untuk menyimpan hasil menengah dalam memori utama, yang lambat.
• Hal ini memungkinkan pemanfaatan unit hardware yang lebih besar yang seharusnya tidak digunakan.
• Hal ini dapat mengurangi penggunaan memori model karena lebih sedikit buffering yang harus aktif pada saat yang sama.

Penyiaran

Penyiaran secara implisit terjadi jika dua TensorFlow dengan bentuk yang berbeda, tetapi kompatibel, digabungkan.

Misalnya, tf.add(vector, matrix) mengharuskan vektor untuk disiarkan sesuai bentuk matriks. Hasil operasi memiliki bentuk yang sama dengan matriks. Untuk detail selengkapnya, lihat panduan untuk menyiarkan array.

Meskipun siaran sering kali dapat digabungkan dengan konsumennya, memaksakan siaran dapat mengakibatkan performa yang buruk dan peningkatan penggunaan memori.

Pada contoh berikut, siaran implisit dalam penambahan vektor dan matriks tidak dapat digabungkan dengan argmax yang menghasilkan siaran terwujud:

`tf.argmax(tf.add(vector, zero_matrix), axis=0)`