Panduan performa Cloud TPU

Langkah pertama Anda saat memecahkan masalah performa TPU adalah membuat profil model. Untuk mengetahui informasi selengkapnya tentang cara mengambil 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 menganggur, Anda harus memastikan ada aliran data yang stabil yang dimuat ke TPU. Cara melakukannya bergantung pada cara Anda memuat dan memproses data set. Misalnya, Anda dapat membaca file data secara paralel menggunakan tf.data.TFRecordset() dan parameter num_parallel_reads.

  2. Ukuran batch terlalu kecil karena sharding (membagi batch di seluruh core)

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

    Untuk penggunaan memori yang optimal, gunakan ukuran batch terbesar yang sesuai dengan memori TPU. Setiap core TPU menggunakan register vektor 8 X 128 dua dimensi untuk memproses perkalian matriks. Secara umum, ukuran batch Anda harus dapat dibagi rata 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 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 oleh XLA ke TPU yang dapat dieksekusi. Untuk informasi selengkapnya tentang XLA, lihat XLA: Mengoptimalkan Compiler untuk Machine Learning.

Padding

Untuk menggunakan memori TPU secara efisien, strukturkan data Anda agar dapat disusun menjadi bagian 128 x 8. Jika data untuk komputasi matriks tidak mengisi seluruh bagian 128 x 8, compiler XLA akan menambahkan padding pada tensor. Ada dua kelemahan padding:

  1. Tensor dengan padding kurang memanfaatkan core TPU.
  2. 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 sangat cocok untuk TPU.

Dimensi tensor

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

Runtime TPU menata tensor dalam 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 batch harus kelipatan 64 (8 per core TPU), dan ukuran dimensi fitur harus kelipatan 128.

  2. Total ukuran batch harus kelipatan 1024 (128 per core TPU), dan ukuran dimensi fitur harus kelipatan 8.

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

Fusion

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

Misalnya, pertimbangkan rangkaian operasi berikut:

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

Kode ini kira-kira setara 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 penggabungan, akses array terjadi secara bersamaan:

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

Dalam contoh ini, jumlah perjalanan bolak-balik memori berkurang dan XLA tidak perlu mengalokasikan ruang apa pun untuk 'tmp'.

Fusi adalah pengoptimalan penting dan bermanfaat bagi Cloud TPU dengan beberapa cara:

  • Hal ini mengurangi transfer memori dengan menghilangkan kebutuhan untuk menyimpan hasil pertengahan di memori utama, yang lambat.
  • Hal ini memungkinkan penggunaan unit hardware yang lebih besar yang tidak akan digunakan.
  • Hal ini dapat mengurangi penggunaan memori model karena lebih sedikit buffering yang perlu aktif secara bersamaan.

Penyiaran

Siaran terjadi secara implisit saat dua tensor dengan bentuk yang berbeda, tetapi kompatibel, digabungkan.

Misalnya, tf.add(vector, matrix) memerlukan vektor untuk disiarkan ke bentuk matriks. Hasil operasi memiliki bentuk yang sama dengan matriks. Untuk mengetahui detail selengkapnya, lihat panduan untuk melakukan siaran array.

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

Dalam contoh berikut, siaran yang 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)`